From d70d44053e12b2ecacaa68ff56e8b73d85bcaf34 Mon Sep 17 00:00:00 2001 From: Ayush Shah Date: Sat, 14 Aug 2021 10:10:32 +0530 Subject: [PATCH] Rest Users Modification --- .../metadata/ingestion/sink/ldap_rest_users.py | 5 +++-- .../ingestion/sink/metadata_rest_users.py | 15 ++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py b/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py index 5975140f026..6056c245fa2 100644 --- a/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py +++ b/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py @@ -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: diff --git a/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py b/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py index 43ba7e3111c..071d1f40dbc 100644 --- a/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py +++ b/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py @@ -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))