diff --git a/ingestion/src/metadata/ingestion/sink/file.py b/ingestion/src/metadata/ingestion/sink/file.py index 5129da3ae57..9aab79da6bd 100644 --- a/ingestion/src/metadata/ingestion/sink/file.py +++ b/ingestion/src/metadata/ingestion/sink/file.py @@ -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__) diff --git a/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py b/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py index 1cd422b23fd..e6256005ee1 100644 --- a/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py +++ b/ingestion/src/metadata/ingestion/sink/ldap_rest_users.py @@ -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__) diff --git a/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py b/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py index 174fd8416f5..f39afb76791 100644 --- a/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py +++ b/ingestion/src/metadata/ingestion/sink/metadata_rest_users.py @@ -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: diff --git a/ingestion/src/metadata/ingestion/source/bigquery_usage.py b/ingestion/src/metadata/ingestion/source/bigquery_usage.py index d1995a15e55..08eb0e7e51f 100644 --- a/ingestion/src/metadata/ingestion/source/bigquery_usage.py +++ b/ingestion/src/metadata/ingestion/source/bigquery_usage.py @@ -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 diff --git a/ingestion/src/metadata/ingestion/source/ldap_users.py b/ingestion/src/metadata/ingestion/source/ldap_users.py index 52900711478..206dc7bf426 100644 --- a/ingestion/src/metadata/ingestion/source/ldap_users.py +++ b/ingestion/src/metadata/ingestion/source/ldap_users.py @@ -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__) diff --git a/ingestion/src/metadata/ingestion/source/redshift.py b/ingestion/src/metadata/ingestion/source/redshift.py index da576a6ef06..e06be6ce623 100644 --- a/ingestion/src/metadata/ingestion/source/redshift.py +++ b/ingestion/src/metadata/ingestion/source/redshift.py @@ -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 diff --git a/ingestion/src/metadata/ingestion/source/sample_users.py b/ingestion/src/metadata/ingestion/source/sample_users.py index fa7a82ccfdc..557ab9c5f3a 100644 --- a/ingestion/src/metadata/ingestion/source/sample_users.py +++ b/ingestion/src/metadata/ingestion/source/sample_users.py @@ -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() diff --git a/ingestion/src/metadata/ingestion/source/snowflake_usage.py b/ingestion/src/metadata/ingestion/source/snowflake_usage.py index 53d413297c8..7ec61a94bd9 100644 --- a/ingestion/src/metadata/ingestion/source/snowflake_usage.py +++ b/ingestion/src/metadata/ingestion/source/snowflake_usage.py @@ -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 diff --git a/ingestion/src/metadata/ingestion/source/sql_alchemy_helper.py b/ingestion/src/metadata/ingestion/source/sql_alchemy_helper.py index 38c7b469b5a..8da5bb9d90d 100644 --- a/ingestion/src/metadata/ingestion/source/sql_alchemy_helper.py +++ b/ingestion/src/metadata/ingestion/source/sql_alchemy_helper.py @@ -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: