openmetadata-ingestion 1.2.5.2__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

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

@@ -13,6 +13,7 @@
 DataLake connector to fetch metadata from a files stored s3, gcs and Hdfs
 """
 import json
+import os
 import traceback
 from typing import Any, Iterable, Tuple, Union
 
@@ -41,6 +42,9 @@ 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.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
@@ -50,12 +54,18 @@ from metadata.generated.schema.metadataIngestion.storage.containerMetadataConfig
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+)
+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,
@@ -65,12 +75,13 @@ 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,
-    get_columns,
     get_file_format_type,
 )
-from metadata.utils.filters import filter_by_schema, filter_by_table
+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
 
@@ -93,8 +104,10 @@ class DatalakeSource(DatabaseServiceSource):
         )
         self.metadata = metadata
         self.service_connection = self.config.serviceConnection.__root__.config
+        self.temp_credentials_file_path = []
         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()
@@ -122,8 +135,47 @@ class DatalakeSource(DatabaseServiceSource):
         Sources with multiple databases should overwrite this and
         apply the necessary filters.
         """
-        database_name = self.service_connection.databaseName or DEFAULT_DATABASE
-        yield database_name
+        if isinstance(self.config_source, GCSConfig):
+            project_id_list = (
+                self.service_connection.configSource.securityConfig.gcpConfig.projectId.__root__
+            )
+            if not isinstance(
+                project_id_list,
+                list,
+            ):
+                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.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
 
     def yield_database(
         self, database_name: str
@@ -132,6 +184,8 @@ class DatalakeSource(DatabaseServiceSource):
         From topology.
         Prepare a database request and pass it to the sink
         """
+        if isinstance(self.config_source, GCSConfig):
+            database_name = self.client.project
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
@@ -140,24 +194,42 @@ class DatalakeSource(DatabaseServiceSource):
         )
 
     def fetch_gcs_bucket_names(self):
-        for bucket in self.client.list_buckets():
-            schema_fqn = fqn.build(
-                self.metadata,
-                entity_type=DatabaseSchema,
-                service_name=self.context.database_service,
-                database_name=self.context.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
+        """
+        Fetch Google cloud storage buckets
+        """
+        try:
+            # List all the buckets in the project
+            for bucket in self.client.list_buckets():
+                # Build a fully qualified name (FQN) for each bucket
+                schema_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=DatabaseSchema,
+                    service_name=self.context.database_service,
+                    database_name=self.context.database,
+                    schema_name=bucket.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,
+                ):
+                    # If it does not match, the bucket is filtered out
+                    self.status.filter(schema_fqn, "Bucket Filtered Out")
+                    continue
 
-            yield bucket.name
+                # If it does match, the bucket name is yielded
+                yield bucket.name
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name="Bucket",
+                    error=f"Unexpected exception to yield bucket: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
 
     def fetch_s3_bucket_names(self):
         for bucket in self.client.list_buckets()["Buckets"]:
@@ -344,9 +416,14 @@ class DatalakeSource(DatabaseServiceSource):
                     file_extension=table_extension,
                 ),
             )
-
-            # If no data_frame (due to unsupported type), ignore
-            columns = get_columns(data_frame[0]) if data_frame else None
+            if data_frame:
+                column_parser = DataFrameColumnParser.create(
+                    data_frame[0], table_extension
+                )
+                columns = column_parser.get_columns()
+            else:
+                # If no data_frame (due to unsupported type), ignore
+                columns = None
             if columns:
                 table_request = CreateTableRequest(
                     name=table_name,
@@ -369,7 +446,7 @@ class DatalakeSource(DatabaseServiceSource):
                 left=StackTraceError(
                     name="Table",
                     error=f"Unexpected exception to yield table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -431,3 +508,12 @@ class DatalakeSource(DatabaseServiceSource):
     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)

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

@@ -172,6 +172,11 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
             params_data["job_definition_id"] = job_id
 
         response = client.get(f"/accounts/{account_id}/runs", data=params_data)
+        if not response and not response.get("data"):
+            raise DBTConfigException(
+                "Unable to get the dbt job runs information.\n"
+                "Please check if the auth token is correct and has the necessary scopes to fetch dbt runs"
+            )
         runs_data = response.get("data")
         if runs_data:
             run_id = runs_data[0]["id"]

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

@@ -31,8 +31,8 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.ingestion.source.database.dbt.dbt_config import get_dbt_details
@@ -56,11 +56,6 @@ class DbtServiceTopology(ServiceTopology):
 
     root = TopologyNode(
         producer="get_dbt_files",
-        stages=[],
-        children=["process_dbt_files"],
-    )
-    process_dbt_files = TopologyNode(
-        producer="process_dbt_files",
         stages=[
             NodeStage(
                 type_=DbtFiles,
@@ -82,12 +77,13 @@ class DbtServiceTopology(ServiceTopology):
                 context="tags",
                 processor="yield_dbt_tags",
                 nullable=True,
-                cache_all=True,
+                store_all_in_context=True,
             ),
             NodeStage(
                 type_=DataModelLink,
                 processor="yield_data_models",
                 nullable=True,
+                consumer=["validate_dbt_files"],
             ),
         ],
     )
@@ -97,6 +93,7 @@ class DbtServiceTopology(ServiceTopology):
             NodeStage(
                 type_=AddLineageRequest,
                 processor="create_dbt_lineage",
+                consumer=["yield_data_models"],
             ),
             NodeStage(
                 type_=AddLineageRequest,
@@ -115,6 +112,7 @@ class DbtServiceTopology(ServiceTopology):
             NodeStage(
                 type_=CreateTestDefinitionRequest,
                 processor="create_dbt_tests_definition",
+                consumer=["yield_data_models"],
             ),
             NodeStage(
                 type_=CreateTestCaseRequest,
@@ -135,9 +133,13 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
     """
 
     topology = DbtServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
     source_config: DbtPipeline
 
+    @property
+    def name(self) -> str:
+        return "dbt"
+
     def remove_manifest_non_required_keys(self, manifest_dict: dict):
         """
         Method to remove the non required keys from manifest file
@@ -157,12 +159,6 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             }
         )
 
-    def process_dbt_files(self) -> Iterable[DbtFiles]:
-        """
-        Method return the dbt file from topology
-        """
-        yield self.context.dbt_file
-
     def get_dbt_files(self) -> Iterable[DbtFiles]:
         dbt_files = get_dbt_details(self.source_config.dbtConfigSource)
         for dbt_file in dbt_files:

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

@@ -14,6 +14,7 @@ DBT utils methods.
 import traceback
 from typing import Optional, Union
 
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -81,7 +82,8 @@ def generate_entity_link(dbt_test):
     manifest_node = dbt_test.get(DbtCommonEnum.MANIFEST_NODE.value)
     entity_link_list = [
         entity_link.get_entity_link(
-            table_fqn=table_fqn,
+            Table,
+            fqn=table_fqn,
             column_name=manifest_node.column_name
             if hasattr(manifest_node, "column_name")
             else None,

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

@@ -28,14 +28,15 @@ from metadata.generated.schema.entity.data.table import (
     Table,
 )
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.tests.basic import (
-    TestCaseFailureStatus,
-    TestCaseFailureStatusType,
     TestCaseResult,
     TestCaseStatus,
     TestResultValue,
@@ -46,11 +47,11 @@ from metadata.generated.schema.tests.testDefinition import (
     TestDefinition,
     TestPlatform,
 )
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName, Timestamp
+from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -270,7 +271,7 @@ class DbtSource(DbtServiceSource):
                         left=StackTraceError(
                             name=key,
                             error=f"Unable to process DBT tags for node: f{key} - {exc}",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
             try:
@@ -298,7 +299,7 @@ class DbtSource(DbtServiceSource):
                     left=StackTraceError(
                         name="Tags and Classification",
                         error=f"Unexpected exception creating DBT tags: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -469,7 +470,7 @@ class DbtSource(DbtServiceSource):
                         left=StackTraceError(
                             name=key,
                             error=f"Unexpected exception parsing DBT node due to {exc}",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
 
@@ -677,7 +678,7 @@ class DbtSource(DbtServiceSource):
                         f"Failed to parse the query {data_model_link.datamodel.sql.__root__}"
                         f" to capture lineage: {exc}"
                     ),
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -691,7 +692,6 @@ class DbtSource(DbtServiceSource):
         )
         if table_entity:
             try:
-
                 service_name, database_name, schema_name, table_name = fqn.split(
                     table_entity.fullyQualifiedName.__root__
                 )
@@ -783,7 +783,7 @@ class DbtSource(DbtServiceSource):
                 left=StackTraceError(
                     name="Test Definition",
                     error=f"Failed to parse the node to capture tests {err}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -821,11 +821,10 @@ class DbtSource(DbtServiceSource):
                 left=StackTraceError(
                     name="Test Cases",
                     error=f"Failed to parse the node to capture tests {err}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
-    # pylint: disable=too-many-locals
     def add_dbt_test_result(self, dbt_test: dict):
         """
         After test cases has been processed, add the tests results info
@@ -840,7 +839,6 @@ class DbtSource(DbtServiceSource):
                 dbt_test_result = dbt_test.get(DbtCommonEnum.RESULTS.value)
                 test_case_status = TestCaseStatus.Aborted
                 test_result_value = 0
-                test_case_failure_status = TestCaseFailureStatus()  # type: ignore
                 if dbt_test_result.status.value in [
                     item.value for item in DbtTestSuccessEnum
                 ]:
@@ -851,17 +849,6 @@ class DbtSource(DbtServiceSource):
                 ]:
                     test_case_status = TestCaseStatus.Failed
                     test_result_value = 0
-                    test_case_failure_status = TestCaseFailureStatus(
-                        testCaseFailureStatusType=TestCaseFailureStatusType.New,
-                        testCaseFailureReason=None,
-                        testCaseFailureComment=None,
-                        updatedAt=Timestamp(
-                            __root__=convert_timestamp_to_milliseconds(
-                                datetime.utcnow().timestamp()
-                            )
-                        ),
-                        updatedBy=None,
-                    )
 
                 # Process the Test Timings
                 dbt_test_timings = dbt_test_result.timing
@@ -893,7 +880,6 @@ class DbtSource(DbtServiceSource):
                             value=str(test_result_value),
                         )
                     ],
-                    testCaseFailureStatus=test_case_failure_status,
                     sampleData=None,
                     result=None,
                 )

metadata/ingestion/source/database/deltalake/metadata.py

@@ -34,13 +34,16 @@ from metadata.generated.schema.entity.data.table import Column, Table, TableType
 from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
     DeltaLakeConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+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
@@ -201,7 +204,7 @@ class DeltalakeSource(DatabaseServiceSource):
         :return: tables or views, depending on config
         """
         schema_name = self.context.database_schema
-        for table in self.spark.catalog.listTables(schema_name):
+        for table in self.spark.catalog.listTables(dbName=schema_name):
             try:
                 table_name = table.name
                 table_fqn = fqn.build(
@@ -289,7 +292,7 @@ class DeltalakeSource(DatabaseServiceSource):
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception to yield table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

metadata/ingestion/source/database/domodatabase/metadata.py

@@ -32,6 +32,9 @@ from metadata.generated.schema.entity.data.table import Column, Table, TableType
 from metadata.generated.schema.entity.services.connections.database.domoDatabaseConnection import (
     DomoDatabaseConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
@@ -39,7 +42,7 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.api.models import Either, StackTraceError
+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
@@ -155,7 +158,7 @@ class DomodatabaseSource(DatabaseServiceSource):
                 StackTraceError(
                     name=schema_name,
                     error=f"Fetching tables names failed for schema {schema_name} due to - {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -163,9 +166,7 @@ class DomodatabaseSource(DatabaseServiceSource):
         try:
             owner_details = User(**self.domo_client.users_get(owner.id))
             if owner_details.email:
-                user = self.metadata.get_user_by_email(owner_details.email)
-                if user:
-                    return EntityReference(id=user.id.__root__, type="user")
+                return self.metadata.get_reference_by_email(owner_details.email)
         except Exception as exc:
             logger.warning(f"Error while getting details of user {owner.name} - {exc}")
         return None
@@ -208,7 +209,7 @@ class DomodatabaseSource(DatabaseServiceSource):
                 left=StackTraceError(
                     name=table_id,
                     error=f"Unexpected exception for table [{table_id}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -0,0 +1,72 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Source connection handler
+"""
+from typing import Optional
+
+from sqlalchemy.engine import Engine
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
+    DorisConnection,
+)
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    get_connection_url_common,
+    init_empty_connection_options,
+)
+from metadata.ingestion.connections.test_connections import (
+    test_connection_db_schema_sources,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+
+def get_connection(connection: DorisConnection) -> Engine:
+    """
+    Create connection
+    """
+    if connection.sslCA or connection.sslCert or connection.sslKey:
+        if not connection.connectionOptions:
+            connection.connectionOptions = init_empty_connection_options()
+        if connection.sslCA:
+            connection.connectionOptions.__root__["ssl_ca"] = connection.sslCA
+        if connection.sslCert:
+            connection.connectionOptions.__root__["ssl_cert"] = connection.sslCert
+        if connection.sslKey:
+            connection.connectionOptions.__root__["ssl_key"] = connection.sslKey
+    return create_generic_db_connection(
+        connection=connection,
+        get_connection_url_fn=get_connection_url_common,
+        get_connection_args_fn=get_connection_args_common,
+    )
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    engine: Engine,
+    service_connection: DorisConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+    test_connection_db_schema_sources(
+        metadata=metadata,
+        engine=engine,
+        service_connection=service_connection,
+        automation_workflow=automation_workflow,
+    )