Sample User pipeline fixed and modified imports (#322)

This commit is contained in:
Ayush Shah 2021-08-27 14:18:33 +05:30 committed by GitHub
parent c36727aaf9
commit 5703a69e8c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 17 additions and 16 deletions

View File

@ -19,7 +19,7 @@ import pathlib
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.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)

View File

@ -19,7 +19,7 @@ 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.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.ometa.client import REST from metadata.ingestion.ometa.client import REST
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)

View File

@ -19,8 +19,8 @@ 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.openmetadata_rest import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig, OpenMetadataAPIClient
from metadata.ingestion.ometa.client import REST, APIError from metadata.ingestion.ometa.client import APIError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -45,7 +45,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.client = OpenMetadataAPIClient(self.metadata_config).client
self._bootstrap_entities() self._bootstrap_entities()
@classmethod @classmethod
@ -60,7 +60,7 @@ class MetadataRestUsersSink(Sink):
def _bootstrap_entities(self): def _bootstrap_entities(self):
# Fetch teams per org # Fetch teams per org
team_response = self.rest.get(self.api_team_get) team_response = self.client.get(self.api_team_get)
for team in team_response['data']: for team in team_response['data']:
self.team_entities[team['displayName']] = team['id'] self.team_entities[team['displayName']] = team['id']
@ -68,7 +68,7 @@ class MetadataRestUsersSink(Sink):
team_name = record.team_name team_name = record.team_name
metadata_team = MetadataTeam(team_name, 'Team Name') metadata_team = MetadataTeam(team_name, 'Team Name')
try: try:
r = self.rest.post(self.api_team_post, r = self.client.post(self.api_team_post,
data=metadata_team.to_json() data=metadata_team.to_json()
) )
instance_id = r['id'] instance_id = r['id']
@ -86,7 +86,7 @@ class MetadataRestUsersSink(Sink):
email=record.email, email=record.email,
teams=teams) teams=teams)
try: try:
self.rest.post(self.api_users, data=metadata_user.to_json()) self.client.post(self.api_users, data=metadata_user.to_json())
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))
except APIError: except APIError:

View File

@ -19,7 +19,7 @@ from metadata.ingestion.models.table_queries import TableQuery
from google.cloud import logging from google.cloud import logging
import collections import collections
from datetime import datetime from datetime import datetime
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.api.source import Source, SourceStatus from metadata.ingestion.api.source import Source, SourceStatus
from typing import Dict, Any, Iterable from typing import Dict, Any, Iterable
from metadata.ingestion.source.sql_alchemy_helper import SQLSourceStatus from metadata.ingestion.source.sql_alchemy_helper import SQLSourceStatus

View File

@ -21,7 +21,7 @@ from metadata.config.common import ConfigModel
from metadata.ingestion.api.common import WorkflowContext from metadata.ingestion.api.common import WorkflowContext
from metadata.ingestion.api.source import SourceStatus, Source from metadata.ingestion.api.source import SourceStatus, Source
from metadata.ingestion.models.user import MetadataUser, User from metadata.ingestion.models.user import MetadataUser, User
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)

View File

@ -16,7 +16,7 @@
import logging import logging
from typing import Optional from typing import Optional
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.source.sql_source import SQLSource, SQLConnectionConfig from metadata.ingestion.source.sql_source import SQLSource, SQLConnectionConfig
from metadata.ingestion.api.source import SourceStatus from metadata.ingestion.api.source import SourceStatus

View File

@ -21,7 +21,7 @@ from typing import Iterable, List
from dataclasses import dataclass, field from dataclasses import dataclass, field
from metadata.config.common import ConfigModel from metadata.config.common import ConfigModel
from metadata.ingestion.api.source import Source, SourceStatus from metadata.ingestion.api.source import Source, SourceStatus
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig, OpenMetadataAPIClient
from metadata.ingestion.models.table_metadata import DatabaseMetadata from metadata.ingestion.models.table_metadata import DatabaseMetadata
from metadata.ingestion.models.user import User from metadata.ingestion.models.user import User
@ -86,6 +86,7 @@ class SampleUsersSource(Source):
def __init__(self, config: SampleUserSourceConfig, metadata_config: MetadataServerConfig, ctx): def __init__(self, config: SampleUserSourceConfig, metadata_config: MetadataServerConfig, ctx):
super().__init__(ctx) super().__init__(ctx)
self.client = OpenMetadataAPIClient(metadata_config)
self.status = SampleUserSourceStatus() self.status = SampleUserSourceStatus()
metadata_gen = SampleUserMetadataGenerator(config.no_of_users) metadata_gen = SampleUserMetadataGenerator(config.no_of_users)
self.sample_columns = metadata_gen.generate_sample_user() self.sample_columns = metadata_gen.generate_sample_user()

View File

@ -15,7 +15,7 @@
# This import verifies that the dependencies are available. # This import verifies that the dependencies are available.
from metadata.ingestion.models.table_queries import TableQuery from metadata.ingestion.models.table_queries import TableQuery
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.source.sql_alchemy_helper import SQLAlchemyHelper, SQLSourceStatus from metadata.ingestion.source.sql_alchemy_helper import SQLAlchemyHelper, SQLSourceStatus
from metadata.ingestion.api.source import Source, SourceStatus from metadata.ingestion.api.source import Source, SourceStatus
from typing import Iterator, Union, Dict, Any, Iterable from typing import Iterator, Union, Dict, Any, Iterable

View File

@ -18,7 +18,7 @@ from typing import Any, Iterable
from metadata.ingestion.api.common import WorkflowContext from metadata.ingestion.api.common import WorkflowContext
from sqlalchemy import create_engine from sqlalchemy import create_engine
from .sql_source import SQLConnectionConfig, SQLSourceStatus from .sql_source import SQLConnectionConfig, SQLSourceStatus
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
class SQLAlchemyHelper: class SQLAlchemyHelper: