openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/datalake/connection.py

@@ -12,15 +12,10 @@
 """
 Source connection handler
 """
-import os
-from copy import deepcopy
 from dataclasses import dataclass
-from functools import partial, singledispatch
+from functools import singledispatch
 from typing import Optional
 
-from google.cloud import storage
-
-from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -36,13 +31,13 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
-from metadata.generated.schema.security.credentials.gcpValues import (
-    MultipleProjectId,
-    SingleProjectId,
-)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.credentials import GOOGLE_CREDENTIALS, set_google_credentials
+from metadata.ingestion.source.database.datalake.clients.azure_blob import (
+    DatalakeAzureBlobClient,
+)
+from metadata.ingestion.source.database.datalake.clients.gcs import DatalakeGcsClient
+from metadata.ingestion.source.database.datalake.clients.s3 import DatalakeS3Client
 
 
 # Only import specific datalake dependencies if necessary
@@ -66,44 +61,17 @@ def get_datalake_client(config):
 
 @get_datalake_client.register
 def _(config: S3Config):
-    from metadata.clients.aws_client import AWSClient
-
-    s3_client = AWSClient(config.securityConfig).get_client(service_name="s3")
-    return s3_client
+    return DatalakeS3Client.from_config(config)
 
 
 @get_datalake_client.register
 def _(config: GCSConfig):
-    gcs_config = deepcopy(config)
-    if hasattr(config.securityConfig, "gcpConfig") and isinstance(
-        config.securityConfig.gcpConfig.projectId, MultipleProjectId
-    ):
-        gcs_config: GCSConfig = deepcopy(config)
-        gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-            gcs_config.securityConfig.gcpConfig.projectId.__root__[0]
-        )
-    set_google_credentials(gcp_credentials=gcs_config.securityConfig)
-    gcs_client = storage.Client()
-    return gcs_client
+    return DatalakeGcsClient.from_config(config)
 
 
 @get_datalake_client.register
 def _(config: AzureConfig):
-    try:
-        return AzureClient(config.securityConfig).create_blob_client()
-    except Exception as exc:
-        raise RuntimeError(
-            f"Unknown error connecting with {config.securityConfig}: {exc}."
-        )
-
-
-def set_gcs_datalake_client(config: GCSConfig, project_id: str):
-    gcs_config = deepcopy(config)
-    if hasattr(gcs_config.securityConfig, "gcpConfig"):
-        gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-            project_id
-        )
-    return get_datalake_client(gcs_config)
+    return DatalakeAzureBlobClient.from_config(config)
 
 
 def get_connection(connection: DatalakeConnection) -> DatalakeClient:
@@ -128,36 +96,10 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    config = connection.config.configSource
-    func = None
-    if isinstance(config, GCSConfig):
-        if connection.config.bucketName:
-            func = partial(connection.client.get_bucket, connection.config.bucketName)
-        else:
-            func = connection.client.list_buckets
-        os.environ.pop("GOOGLE_CLOUD_PROJECT", "")
-        if GOOGLE_CREDENTIALS in os.environ:
-            os.remove(os.environ[GOOGLE_CREDENTIALS])
-            del os.environ[GOOGLE_CREDENTIALS]
-
-    if isinstance(config, S3Config):
-        if connection.config.bucketName:
-            func = partial(
-                connection.client.list_objects, Bucket=connection.config.bucketName
-            )
-        else:
-            func = connection.client.list_buckets
-
-    if isinstance(config, AzureConfig):
-
-        def list_connection(connection):
-            conn = connection.client.list_containers(name_starts_with="")
-            list(conn)
-
-        func = partial(list_connection, connection)
-
     test_fn = {
-        "ListBuckets": func,
+        "ListBuckets": connection.client.get_test_list_buckets_fn(
+            connection.config.bucketName
+        ),
     }
 
     test_connection_steps(

metadata/ingestion/source/database/datalake/metadata.py

@@ -13,7 +13,6 @@
 DataLake connector to fetch metadata from a files stored s3, gcs and Hdfs
 """
 import json
-import os
 import traceback
 from typing import Any, Iterable, Optional, Tuple, Union
 
@@ -30,15 +29,9 @@ from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import Table, TableType
-from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
-    AzureConfig,
-)
 from metadata.generated.schema.entity.services.connections.database.datalake.gcsConfig import (
     GCSConfig,
 )
-from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
-    S3Config,
-)
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
@@ -54,28 +47,22 @@ from metadata.generated.schema.metadataIngestion.storage.containerMetadataConfig
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.security.credentials.gcpValues import (
-    GcpCredentialsValues,
-)
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.database_service import DatabaseServiceSource
-from metadata.ingestion.source.database.datalake.connection import (
-    set_gcs_datalake_client,
-)
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.ingestion.source.storage.storage_service import (
     OPENMETADATA_TEMPLATE_FILE_NAME,
 )
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
+from metadata.readers.dataframe.reader_factory import SupportedTypes
 from metadata.readers.file.base import ReadException
 from metadata.readers.file.config_source_factory import get_reader
 from metadata.utils import fqn
-from metadata.utils.constants import DEFAULT_DATABASE
-from metadata.utils.credentials import GOOGLE_CREDENTIALS
 from metadata.utils.datalake.datalake_utils import (
     DataFrameColumnParser,
     fetch_dataframe,
@@ -83,7 +70,6 @@ from metadata.utils.datalake.datalake_utils import (
 )
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.s3_utils import list_s3_objects
 
 logger = ingestion_logger()
 
@@ -103,25 +89,24 @@ class DatalakeSource(DatabaseServiceSource):
             self.config.sourceConfig.config
         )
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
-        self.temp_credentials_file_path = []
+        self.service_connection = self.config.serviceConnection.root.config
         self.connection = get_connection(self.service_connection)
-        if GOOGLE_CREDENTIALS in os.environ:
-            self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
         self.client = self.connection.client
         self.table_constraints = None
         self.database_source_state = set()
         self.config_source = self.service_connection.configSource
         self.connection_obj = self.connection
         self.test_connection()
-        self.reader = get_reader(config_source=self.config_source, client=self.client)
+        self.reader = get_reader(
+            config_source=self.config_source, client=self.client._client
+        )
 
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatalakeConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: DatalakeConnection = config.serviceConnection.root.config
         if not isinstance(connection, DatalakeConnection):
             raise InvalidSourceException(
                 f"Expected DatalakeConnection, but got {connection}"
@@ -137,47 +122,31 @@ class DatalakeSource(DatabaseServiceSource):
         Sources with multiple databases should overwrite this and
         apply the necessary filters.
         """
-        if isinstance(self.config_source, GCSConfig):
-            project_id_list = (
-                self.service_connection.configSource.securityConfig.gcpConfig.projectId.__root__
+        for database_name in self.client.get_database_names(self.service_connection):
+            database_fqn = fqn.build(
+                self.metadata,
+                entity_type=Database,
+                service_name=self.context.get().database_service,
+                database_name=database_name,
             )
-            if not isinstance(
-                project_id_list,
-                list,
+            if filter_by_database(
+                self.source_config.databaseFilterPattern,
+                database_fqn
+                if self.source_config.useFqnForFiltering
+                else database_name,
             ):
-                project_id_list = [project_id_list]
-            for project_id in project_id_list:
-                database_fqn = fqn.build(
-                    self.metadata,
-                    entity_type=Database,
-                    service_name=self.context.get().database_service,
-                    database_name=project_id,
-                )
-                if filter_by_database(
-                    self.source_config.databaseFilterPattern,
-                    database_fqn
-                    if self.source_config.useFqnForFiltering
-                    else project_id,
-                ):
-                    self.status.filter(database_fqn, "Database Filtered out")
-                else:
-                    try:
-                        self.client = set_gcs_datalake_client(
-                            config=self.config_source, project_id=project_id
-                        )
-                        if GOOGLE_CREDENTIALS in os.environ:
-                            self.temp_credentials_file_path.append(
-                                os.environ[GOOGLE_CREDENTIALS]
-                            )
-                        yield project_id
-                    except Exception as exc:
-                        logger.debug(traceback.format_exc())
-                        logger.error(
-                            f"Error trying to connect to database {project_id}: {exc}"
-                        )
-        else:
-            database_name = self.service_connection.databaseName or DEFAULT_DATABASE
-            yield database_name
+                self.status.filter(database_fqn, "Database Filtered out")
+            else:
+                try:
+                    self.client.update_client_database(
+                        self.config_source, database_name
+                    )
+                    yield database_name
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.error(
+                        f"Error trying to connect to database {database_name}: {exc}"
+                    )
 
     def yield_database(
         self, database_name: str
@@ -190,40 +159,37 @@ class DatalakeSource(DatabaseServiceSource):
             database_name = self.client.project
         yield Either(
             right=CreateDatabaseRequest(
-                name=database_name,
-                service=self.context.get().database_service,
+                name=EntityName(database_name),
+                service=FullyQualifiedEntityName(self.context.get().database_service),
             )
         )
 
-    def fetch_gcs_bucket_names(self):
+    def get_database_schema_names(self) -> Iterable[str]:
         """
-        Fetch Google cloud storage buckets
+        return schema names
         """
         try:
-            # List all the buckets in the project
-            for bucket in self.client.list_buckets():
-                # Build a fully qualified name (FQN) for each bucket
+            for schema_name in self.client.get_database_schema_names(
+                self.service_connection.bucketName
+            ):
                 schema_fqn = fqn.build(
                     self.metadata,
                     entity_type=DatabaseSchema,
                     service_name=self.context.get().database_service,
                     database_name=self.context.get().database,
-                    schema_name=bucket.name,
+                    schema_name=schema_name,
                 )
 
-                # Check if the bucket matches a certain filter pattern
                 if filter_by_schema(
                     self.config.sourceConfig.config.schemaFilterPattern,
                     schema_fqn
                     if self.config.sourceConfig.config.useFqnForFiltering
-                    else bucket.name,
+                    else schema_name,
                 ):
-                    # If it does not match, the bucket is filtered out
                     self.status.filter(schema_fqn, "Bucket Filtered Out")
                     continue
 
-                # If it does match, the bucket name is yielded
-                yield bucket.name
+                yield schema_name
         except Exception as exc:
             yield Either(
                 left=StackTraceError(
@@ -233,74 +199,6 @@ class DatalakeSource(DatabaseServiceSource):
                 )
             )
 
-    def fetch_s3_bucket_names(self):
-        for bucket in self.client.list_buckets()["Buckets"]:
-            schema_fqn = fqn.build(
-                self.metadata,
-                entity_type=DatabaseSchema,
-                service_name=self.context.get().database_service,
-                database_name=self.context.get().database,
-                schema_name=bucket["Name"],
-            )
-            if filter_by_schema(
-                self.config.sourceConfig.config.schemaFilterPattern,
-                schema_fqn
-                if self.config.sourceConfig.config.useFqnForFiltering
-                else bucket["Name"],
-            ):
-                self.status.filter(schema_fqn, "Bucket Filtered Out")
-                continue
-            yield bucket["Name"]
-
-    def get_database_schema_names(self) -> Iterable[str]:
-        """
-        return schema names
-        """
-        bucket_name = self.service_connection.bucketName
-        if isinstance(self.config_source, GCSConfig):
-            if bucket_name:
-                yield bucket_name
-            else:
-                yield from self.fetch_gcs_bucket_names()
-
-        if isinstance(self.config_source, S3Config):
-            if bucket_name:
-                yield bucket_name
-            else:
-                yield from self.fetch_s3_bucket_names()
-
-        if isinstance(self.config_source, AzureConfig):
-            yield from self.get_container_names()
-
-    def get_container_names(self) -> Iterable[str]:
-        """
-        To get schema names
-        """
-        prefix = (
-            self.service_connection.bucketName
-            if self.service_connection.bucketName
-            else ""
-        )
-        schema_names = self.client.list_containers(name_starts_with=prefix)
-        for schema in schema_names:
-            schema_fqn = fqn.build(
-                self.metadata,
-                entity_type=DatabaseSchema,
-                service_name=self.context.get().database_service,
-                database_name=self.context.get().database,
-                schema_name=schema["name"],
-            )
-            if filter_by_schema(
-                self.config.sourceConfig.config.schemaFilterPattern,
-                schema_fqn
-                if self.config.sourceConfig.config.useFqnForFiltering
-                else schema["name"],
-            ):
-                self.status.filter(schema_fqn, "Container Filtered Out")
-                continue
-
-            yield schema["name"]
-
     def yield_database_schema(
         self, schema_name: str
     ) -> Iterable[Either[CreateDatabaseSchemaRequest]]:
@@ -310,19 +208,21 @@ class DatalakeSource(DatabaseServiceSource):
         """
         yield Either(
             right=CreateDatabaseSchemaRequest(
-                name=schema_name,
-                database=fqn.build(
-                    metadata=self.metadata,
-                    entity_type=Database,
-                    service_name=self.context.get().database_service,
-                    database_name=self.context.get().database,
+                name=EntityName(schema_name),
+                database=FullyQualifiedEntityName(
+                    fqn.build(
+                        metadata=self.metadata,
+                        entity_type=Database,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                    )
                 ),
             )
         )
 
     def get_tables_name_and_type(  # pylint: disable=too-many-branches
         self,
-    ) -> Iterable[Tuple[str, TableType]]:
+    ) -> Iterable[Tuple[str, TableType, SupportedTypes]]:
         """
         Handle table and views.
 
@@ -344,62 +244,26 @@ class DatalakeSource(DatabaseServiceSource):
         except ReadException:
             metadata_entry = None
         if self.source_config.includeTables:
-            if isinstance(self.config_source, GCSConfig):
-                bucket = self.client.get_bucket(bucket_name)
-                for key in bucket.list_blobs(prefix=prefix):
-                    table_name = self.standardize_table_name(bucket_name, key.name)
-                    # adding this condition as the gcp blobs also contains directory, which we can filter out
-                    if self.filter_dl_table(table_name):
-                        continue
-                    file_extension = get_file_format_type(
-                        key_name=key.name, metadata_entry=metadata_entry
-                    )
-                    if table_name.endswith("/") or not file_extension:
-                        logger.debug(
-                            f"Object filtered due to unsupported file type: {key.name}"
-                        )
-                        continue
-
-                    yield table_name, TableType.Regular, file_extension
-            if isinstance(self.config_source, S3Config):
-                kwargs = {"Bucket": bucket_name}
-                if prefix:
-                    kwargs["Prefix"] = prefix if prefix.endswith("/") else f"{prefix}/"
-                for key in list_s3_objects(self.client, **kwargs):
-                    table_name = self.standardize_table_name(bucket_name, key["Key"])
-                    if self.filter_dl_table(table_name):
-                        continue
-                    file_extension = get_file_format_type(
-                        key_name=key["Key"], metadata_entry=metadata_entry
-                    )
-                    if not file_extension:
-                        logger.debug(
-                            f"Object filtered due to unsupported file type: {key['Key']}"
-                        )
-                        continue
+            for key_name in self.client.get_table_names(bucket_name, prefix):
+                table_name = self.standardize_table_name(bucket_name, key_name)
+
+                if self.filter_dl_table(table_name):
+                    continue
 
-                    yield table_name, TableType.Regular, file_extension
-            if isinstance(self.config_source, AzureConfig):
-                container_client = self.client.get_container_client(bucket_name)
+                file_extension = get_file_format_type(
+                    key_name=key_name, metadata_entry=metadata_entry
+                )
 
-                for file in container_client.list_blobs(
-                    name_starts_with=prefix or None
-                ):
-                    table_name = self.standardize_table_name(bucket_name, file.name)
-                    if self.filter_dl_table(table_name):
-                        continue
-                    file_extension = get_file_format_type(
-                        key_name=file.name, metadata_entry=metadata_entry
+                if table_name.endswith("/") or not file_extension:
+                    logger.debug(
+                        f"Object filtered due to unsupported file type: {key_name}"
                     )
-                    if not file_extension:
-                        logger.debug(
-                            f"Object filtered due to unsupported file type: {file.name}"
-                        )
-                        continue
-                    yield table_name, TableType.Regular, file_extension
+                    continue
+
+                yield table_name, TableType.Regular, file_extension
 
     def yield_table(
-        self, table_name_and_type: Tuple[str, str]
+        self, table_name_and_type: Tuple[str, TableType, SupportedTypes]
     ) -> Iterable[Either[CreateTableRequest]]:
         """
         From topology.
@@ -411,7 +275,7 @@ class DatalakeSource(DatabaseServiceSource):
             table_constraints = None
             data_frame, raw_data = fetch_dataframe(
                 config_source=self.config_source,
-                client=self.client,
+                client=self.client._client,
                 file_fqn=DatalakeTableSchemaWrapper(
                     key=table_name,
                     bucket_name=schema_name,
@@ -433,12 +297,14 @@ class DatalakeSource(DatabaseServiceSource):
                     tableType=table_type,
                     columns=columns,
                     tableConstraints=table_constraints if table_constraints else None,
-                    databaseSchema=fqn.build(
-                        metadata=self.metadata,
-                        entity_type=DatabaseSchema,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                        schema_name=schema_name,
+                    databaseSchema=FullyQualifiedEntityName(
+                        fqn.build(
+                            metadata=self.metadata,
+                            entity_type=DatabaseSchema,
+                            service_name=self.context.get().database_service,
+                            database_name=self.context.get().database,
+                            schema_name=schema_name,
+                        )
                     ),
                     fileFormat=table_extension.value if table_extension else None,
                 )
@@ -509,14 +375,4 @@ class DatalakeSource(DatabaseServiceSource):
         return False
 
     def close(self):
-        if isinstance(self.config_source, AzureConfig):
-            self.client.close()
-        if isinstance(self.config_source, GCSConfig):
-            os.environ.pop("GOOGLE_CLOUD_PROJECT", "")
-            if isinstance(self.service_connection, GcpCredentialsValues) and (
-                GOOGLE_CREDENTIALS in os.environ
-            ):
-                del os.environ[GOOGLE_CREDENTIALS]
-                for temp_file_path in self.temp_credentials_file_path:
-                    if os.path.exists(temp_file_path):
-                        os.remove(temp_file_path)
+        self.client.close(self.service_connection)

metadata/ingestion/source/database/db2/metadata.py

@@ -10,12 +10,12 @@
 #  limitations under the License.
 """Db2 source module"""
 import traceback
-from typing import Optional
+from typing import Iterable, Optional
 
-from ibm_db_sa.base import DB2Dialect
-from sqlalchemy.engine import reflection
+from ibm_db_sa.base import ischema_names
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.row import LegacyRow
+from sqlalchemy.sql.sqltypes import BOOLEAN
 
 from metadata.generated.schema.entity.services.connections.database.db2Connection import (
     Db2Connection,
@@ -31,14 +31,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
-@reflection.cache
-def get_pk_constraint(
-    self, bind, table_name, schema=None, **kw
-):  # pylint: disable=unused-argument
-    return {"constrained_columns": [], "name": "undefined"}
-
-
-DB2Dialect.get_pk_constraint = get_pk_constraint
+ischema_names.update({"BOOLEAN": BOOLEAN})
 
 
 class Db2Source(CommonDbSourceService):
@@ -51,14 +44,21 @@ class Db2Source(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: Db2Connection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: Db2Connection = config.serviceConnection.root.config
         if not isinstance(connection, Db2Connection):
             raise InvalidSourceException(
                 f"Expected Db2Connection, but got {connection}"
             )
         return cls(config, metadata)
 
+    def get_raw_database_schema_names(self) -> Iterable[str]:
+        if self.service_connection.__dict__.get("databaseSchema"):
+            yield self.service_connection.databaseSchema
+        else:
+            for schema_name in self.inspector.get_schema_names():
+                yield schema_name.rstrip()
+
     @staticmethod
     def get_table_description(
         schema_name: str, table_name: str, inspector: Inspector

metadata/ingestion/source/database/dbt/dbt_config.py

@@ -151,7 +151,7 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
         expiry = 0
         auth_token = config.dbtCloudAuthToken.get_secret_value(), expiry
         client_config = ClientConfig(
-            base_url=config.dbtCloudUrl,
+            base_url=str(config.dbtCloudUrl),
             api_version="api/v2",
             auth_token=lambda: auth_token,
             auth_header="Authorization",
@@ -182,7 +182,8 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
             last_run = runs_data[0]
             run_id = last_run["id"]
             logger.info(
-                f"Retrieved last successful run [{run_id}] finished {last_run['finished_at_humanized']} (duration: {last_run['duration_humanized']})"
+                f"Retrieved last successful run [{str(run_id)}]: "
+                f"Finished {str(last_run['finished_at_humanized'])} (duration: {str(last_run['duration_humanized'])})"
             )
             try:
                 logger.debug("Requesting [dbt_catalog]")