Rest Users Modification

This commit is contained in:
Ayush Shah 2021-08-14 10:10:32 +05:30
parent 165925f6c6
commit d70d44053e
2 changed files with 11 additions and 9 deletions

View File

@ -14,13 +14,13 @@
# limitations under the License.
import logging
import requests
from metadata.config.common import ConfigModel
from metadata.ingestion.api.common import WorkflowContext, Record
from metadata.ingestion.api.sink import Sink, SinkStatus
from metadata.ingestion.models.user import MetadataUser
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig
from metadata.ingestion.ometa.client import REST
logger = logging.getLogger(__name__)
@ -40,6 +40,7 @@ class LdapRestUsersSink(Sink):
self.status = SinkStatus()
self.api_users = self.metadata_config.api_endpoint + "/v1/users"
self.headers = {'Content-type': 'application/json'}
self.rest = REST(metadata_config)
@classmethod
def create(cls, config_dict: dict, metadata_config_dict: dict, ctx: WorkflowContext):
@ -54,7 +55,7 @@ class LdapRestUsersSink(Sink):
metadata_user = MetadataUser(name=record.github_username[0],
display_name=record.name[0],
email=record.email[0])
r = requests.post(self.api_users, data=metadata_user.to_json(), headers=self.headers)
r = self.rest.post(self.api_users, data=metadata_user.to_json(), headers=self.headers)
if r.status_code == 200 or r.status_code == 201:
self.status.records_written(record.name[0])
else:

View File

@ -15,13 +15,12 @@
import logging
import requests
from metadata.config.common import ConfigModel
from metadata.ingestion.api.common import WorkflowContext, Record
from metadata.ingestion.api.sink import Sink, SinkStatus
from metadata.ingestion.models.user import MetadataTeam, MetadataUser
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig
from metadata.ingestion.ometa.client import REST
logger = logging.getLogger(__name__)
@ -47,6 +46,7 @@ class MetadataRestUsersSink(Sink):
self.org_entities = {}
self.role_entities = {}
self.team_entities = {}
self.rest = REST(metadata_config)
self._bootstrap_entities()
@classmethod
@ -60,7 +60,8 @@ class MetadataRestUsersSink(Sink):
def _bootstrap_entities(self):
# Fetch teams per org
r = requests.get(self.api_team_get, headers=self.headers)
r = self.rest.get(self.api_team_get, headers=self.headers)
if r.status_code == 200:
team_response = r.json()
for team in team_response['data']:
@ -69,9 +70,9 @@ class MetadataRestUsersSink(Sink):
def _create_team(self, record: MetadataUser) -> None:
team_name = record.team_name
metadata_team = MetadataTeam(team_name, 'Team Name')
r = requests.post(self.api_team_post,
data=metadata_team.to_json(),
headers=self.headers)
r = self.rest.post(self.api_team_post,
data=metadata_team.to_json(),
headers=self.headers)
if r.status_code == 200 or r.status_code == 201:
instance_id = r.json()['id']
self.team_entities[team_name] = instance_id
@ -85,7 +86,7 @@ class MetadataRestUsersSink(Sink):
display_name=record.name,
email=record.email,
teams=teams)
r = requests.post(self.api_users, data=metadata_user.to_json(), headers=self.headers)
r = self.rest.post(self.api_users, data=metadata_user.to_json(), headers=self.headers)
if r.status_code == 200 or r.status_code == 201:
self.status.records_written(record.github_username)
logger.info("Sink: {}".format(record.github_username))