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

View File

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