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.ingestion.api.common import WorkflowContext, Record
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__)

View File

@ -19,7 +19,7 @@ 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.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.ometa.client import REST
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.sink import Sink, SinkStatus
from metadata.ingestion.models.user import MetadataTeam, MetadataUser
from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
from metadata.ingestion.ometa.client import REST, APIError
from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig, OpenMetadataAPIClient
from metadata.ingestion.ometa.client import APIError
logger = logging.getLogger(__name__)
@ -45,7 +45,7 @@ class MetadataRestUsersSink(Sink):
self.org_entities = {}
self.role_entities = {}
self.team_entities = {}
self.rest = REST(metadata_config)
self.client = OpenMetadataAPIClient(self.metadata_config).client
self._bootstrap_entities()
@classmethod
@ -60,7 +60,7 @@ class MetadataRestUsersSink(Sink):
def _bootstrap_entities(self):
# 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']:
self.team_entities[team['displayName']] = team['id']
@ -68,9 +68,9 @@ class MetadataRestUsersSink(Sink):
team_name = record.team_name
metadata_team = MetadataTeam(team_name, 'Team Name')
try:
r = self.rest.post(self.api_team_post,
data=metadata_team.to_json()
)
r = self.client.post(self.api_team_post,
data=metadata_team.to_json()
)
instance_id = r['id']
self.team_entities[team_name] = instance_id
except APIError:
@ -86,7 +86,7 @@ class MetadataRestUsersSink(Sink):
email=record.email,
teams=teams)
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)
logger.info("Sink: {}".format(record.github_username))
except APIError:

View File

@ -19,7 +19,7 @@ from metadata.ingestion.models.table_queries import TableQuery
from google.cloud import logging
import collections
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 typing import Dict, Any, Iterable
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.source import SourceStatus, Source
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__)

View File

@ -16,7 +16,7 @@
import logging
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.api.source import SourceStatus

View File

@ -21,7 +21,7 @@ from typing import Iterable, List
from dataclasses import dataclass, field
from metadata.config.common import ConfigModel
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.user import User
@ -86,6 +86,7 @@ class SampleUsersSource(Source):
def __init__(self, config: SampleUserSourceConfig, metadata_config: MetadataServerConfig, ctx):
super().__init__(ctx)
self.client = OpenMetadataAPIClient(metadata_config)
self.status = SampleUserSourceStatus()
metadata_gen = SampleUserMetadataGenerator(config.no_of_users)
self.sample_columns = metadata_gen.generate_sample_user()

View File

@ -15,7 +15,7 @@
# This import verifies that the dependencies are available.
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.api.source import Source, SourceStatus
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 sqlalchemy import create_engine
from .sql_source import SQLConnectionConfig, SQLSourceStatus
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig
from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
class SQLAlchemyHelper: