openmetadata-ingestion 1.5.0.0rc1__py3-none-any.whl → 1.5.1.0__py3-none-any.whl

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -49,7 +49,12 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -61,6 +66,7 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 )
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
+from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.connection import get_connection
 from metadata.ingestion.source.database.unitycatalog.models import (
     ColumnJson,
@@ -99,6 +105,7 @@ class UnitycatalogSource(
         )
         self.external_location_map = {}
         self.client = get_connection(self.service_connection)
+        self.api_client = UnityCatalogClient(self.service_connection)
         self.connection_obj = self.client
         self.table_constraints = []
         self.context.storage_location = None
@@ -147,9 +154,11 @@ class UnitycatalogSource(
                     )
                     if filter_by_database(
                         self.config.sourceConfig.config.databaseFilterPattern,
-                        database_fqn
-                        if self.config.sourceConfig.config.useFqnForFiltering
-                        else catalog_name,
+                        (
+                            database_fqn
+                            if self.config.sourceConfig.config.useFqnForFiltering
+                            else catalog_name
+                        ),
                     ):
                         self.status.filter(
                             database_fqn,
@@ -196,9 +205,11 @@ class UnitycatalogSource(
                 )
                 if filter_by_schema(
                     self.config.sourceConfig.config.schemaFilterPattern,
-                    schema_fqn
-                    if self.config.sourceConfig.config.useFqnForFiltering
-                    else schema.name,
+                    (
+                        schema_fqn
+                        if self.config.sourceConfig.config.useFqnForFiltering
+                        else schema.name
+                    ),
                 ):
                     self.status.filter(schema_fqn, "Schema Filtered Out")
                     continue
@@ -260,9 +271,11 @@ class UnitycatalogSource(
                 )
                 if filter_by_table(
                     self.config.sourceConfig.config.tableFilterPattern,
-                    table_fqn
-                    if self.config.sourceConfig.config.useFqnForFiltering
-                    else table_name,
+                    (
+                        table_fqn
+                        if self.config.sourceConfig.config.useFqnForFiltering
+                        else table_name
+                    ),
                 ):
                     self.status.filter(
                         table_fqn,
@@ -330,6 +343,7 @@ class UnitycatalogSource(
                         schema_name=schema_name,
                     )
                 ),
+                owners=self.get_owner_ref(table_name),
             )
             yield Either(right=table_request)
 
@@ -447,12 +461,12 @@ class UnitycatalogSource(
         """
         try:
             if column.children is None:
-                if column_json.metadata:
-                    column.description = column_json.metadata.comment
+                if column_json.metadata and column_json.metadata.comment:
+                    column.description = Markdown(column_json.metadata.comment)
             else:
                 for i, child in enumerate(column.children):
-                    if column_json.metadata:
-                        column.description = column_json.metadata.comment
+                    if column_json.metadata and column_json.metadata.comment:
+                        column.description = Markdown(column_json.metadata.comment)
                     if (
                         column_json.type
                         and isinstance(column_json.type, Type)
@@ -483,17 +497,23 @@ class UnitycatalogSource(
         """
 
         for column in column_data:
-            if column.type_text.lower().startswith("union"):
-                column.type_text = column.Type.replace(" ", "")
-            if column.type_text.lower() == "struct":
-                column.type_text = "struct<>"
+            parsed_string = {}
+            if column.type_text:
+                if column.type_text.lower().startswith("union"):
+                    column.type_text = column.type_text.replace(" ", "")
+                if (
+                    column.type_text.lower() == "struct"
+                    or column.type_text.lower() == "array"
+                ):
+                    column.type_text = column.type_text.lower() + "<>"
 
-            parsed_string = ColumnTypeParser._parse_datatype_string(  # pylint: disable=protected-access
-                column.type_text.lower()
-            )
+                parsed_string = ColumnTypeParser._parse_datatype_string(  # pylint: disable=protected-access
+                    column.type_text.lower()
+                )
             parsed_string["name"] = column.name[:256]
             parsed_string["dataLength"] = parsed_string.get("dataLength", 1)
-            parsed_string["description"] = column.comment
+            if column.comment:
+                parsed_string["description"] = Markdown(column.comment)
             parsed_column = Column(**parsed_string)
             self.add_complex_datatype_descriptions(
                 column=parsed_column,
@@ -537,3 +557,19 @@ class UnitycatalogSource(
 
     def close(self):
         """Nothing to close"""
+
+    def get_owner_ref(self, table_name: str) -> Optional[EntityReferenceList]:
+        """
+        Method to process the table owners
+        """
+        try:
+            full_table_name = f"{self.context.get().database}.{self.context.get().database_schema}.{table_name}"
+            owner = self.api_client.get_owner_info(full_table_name)
+            if not owner:
+                return
+            owner_ref = self.metadata.get_reference_by_email(email=owner)
+            return owner_ref
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error processing owner for table {table_name}: {exc}")
+        return

metadata/ingestion/source/pipeline/dagster/metadata.py

@@ -133,9 +133,11 @@ class DagsterSource(PipelineServiceSource):
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.id.replace(":", "")),
                 displayName=pipeline_details.name,
-                description=Markdown(pipeline_details.description)
-                if pipeline_details.description
-                else None,
+                description=(
+                    Markdown(pipeline_details.description)
+                    if pipeline_details.description
+                    else None
+                ),
                 tasks=self._get_task_list(pipeline_name=pipeline_details.name),
                 service=FullyQualifiedEntityName(self.context.get().pipeline_service),
                 tags=get_tag_labels(
@@ -178,24 +180,25 @@ class DagsterSource(PipelineServiceSource):
                 executionStatus=STATUS_MAP.get(
                     run.status.lower(), StatusType.Pending.value
                 ),
-                startTime=round(convert_timestamp_to_milliseconds(run.startTime))
-                if run.startTime
-                else None,
-                endTime=round(convert_timestamp_to_milliseconds(run.endTime))
-                if run.endTime
-                else None,
+                startTime=(
+                    round(convert_timestamp_to_milliseconds(run.startTime))
+                    if run.startTime
+                    else None
+                ),
+                endTime=(
+                    round(convert_timestamp_to_milliseconds(run.endTime))
+                    if run.endTime
+                    else None
+                ),
             )
-
             pipeline_status = PipelineStatus(
                 taskStatus=[task_status],
                 executionStatus=STATUS_MAP.get(
                     run.status.lower(), StatusType.Pending.value
                 ),
                 timestamp=Timestamp(
-                    round(convert_timestamp_to_milliseconds(run.endTime))
-                )
-                if run.endTime
-                else None,
+                    round(convert_timestamp_to_milliseconds(timestamp=run.startTime))
+                ),
             )
             pipeline_fqn = fqn.build(
                 metadata=self.metadata,

metadata/ingestion/source/pipeline/dbtcloud/client.py

@@ -26,7 +26,10 @@ from metadata.ingestion.source.pipeline.dbtcloud.models import (
     DBTRun,
     DBTRunList,
 )
-from metadata.ingestion.source.pipeline.dbtcloud.queries import DBT_QUERY
+from metadata.ingestion.source.pipeline.dbtcloud.queries import (
+    DBT_GET_MODEL_DEPENDS_ON,
+    DBT_GET_MODELS_SEEDS,
+)
 from metadata.utils.constants import AUTHORIZATION_HEADER
 from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ometa_logger
@@ -45,7 +48,7 @@ class DBTCloudClient:
             base_url=clean_uri(self.config.host),
             api_version="api/v2",
             auth_header=AUTHORIZATION_HEADER,
-            auth_token=lambda: (self.config.token, 0),
+            auth_token=lambda: (self.config.token.get_secret_value(), 0),
             allow_redirects=True,
         )
 
@@ -53,7 +56,7 @@ class DBTCloudClient:
             base_url=clean_uri(self.config.discoveryAPI),
             api_version="",
             auth_header=AUTHORIZATION_HEADER,
-            auth_token=lambda: (self.config.token, 0),
+            auth_token=lambda: (self.config.token.get_secret_value(), 0),
             allow_redirects=True,
         )
 
@@ -102,7 +105,7 @@ class DBTCloudClient:
         """
         try:
             query_params = {
-                "query": DBT_QUERY,
+                "query": DBT_GET_MODEL_DEPENDS_ON,
                 "variables": {"jobId": job_id, "runId": run_id},
             }
 
@@ -116,3 +119,25 @@ class DBTCloudClient:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unable to get model info :{exc}")
         return None
+
+    def get_models_and_seeds_details(self, job_id: int, run_id: int):
+        """
+        get model details for a job in dbt cloud for lineage
+        """
+        try:
+            query_params = {
+                "query": DBT_GET_MODELS_SEEDS,
+                "variables": {"jobId": job_id, "runId": run_id},
+            }
+
+            result = self.graphql_client.post("", json=query_params)
+
+            if result.get("data") and result["data"].get("job"):
+                result = DBTModelList(**result["data"]["job"])
+                parents_list = result.models + result.seeds
+                return parents_list
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Unable to get model info :{exc}")
+        return None

metadata/ingestion/source/pipeline/dbtcloud/metadata.py

@@ -111,6 +111,7 @@ class DbtcloudSource(PipelineServiceSource):
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Failed to get tasks list due to : {exc}")
+        return None
 
     def yield_pipeline(
         self, pipeline_details: DBTJob
@@ -119,16 +120,22 @@ class DbtcloudSource(PipelineServiceSource):
         Method to Get Pipeline Entity
         """
         try:
-            connection_url = f"{clean_uri(self.service_connection.host)}/deploy/{self.service_connection.accountId}/projects/{pipeline_details.project_id}/jobs/{pipeline_details.id}"
+            connection_url = (
+                f"{clean_uri(self.service_connection.host)}/deploy/"
+                f"{self.service_connection.accountId}/projects/"
+                f"{pipeline_details.project_id}/jobs/{pipeline_details.id}"
+            )
 
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.name),
                 description=Markdown(pipeline_details.description),
                 sourceUrl=SourceUrl(connection_url),
                 tasks=self._get_task_list(job_id=int(pipeline_details.id)),
-                scheduleInterval=str(pipeline_details.schedule.cron)
-                if pipeline_details.schedule
-                else None,
+                scheduleInterval=(
+                    str(pipeline_details.schedule.cron)
+                    if pipeline_details.schedule
+                    else None
+                ),
                 service=FullyQualifiedEntityName(self.context.get().pipeline_service),
             )
             yield Either(right=pipeline_request)
@@ -148,77 +155,88 @@ class DbtcloudSource(PipelineServiceSource):
         """
         Get lineage between pipeline and data sources
         """
-        try:
-            pipeline_fqn = fqn.build(
-                metadata=self.metadata,
-                entity_type=Pipeline,
-                service_name=self.context.get().pipeline_service,
-                pipeline_name=self.context.get().pipeline,
-            )
-
-            pipeline_entity = self.metadata.get_by_name(
-                entity=Pipeline, fqn=pipeline_fqn
-            )
+        try:  # pylint: disable=too-many-nested-blocks
+            if self.source_config.lineageInformation:
+                pipeline_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Pipeline,
+                    service_name=self.context.get().pipeline_service,
+                    pipeline_name=self.context.get().pipeline,
+                )
 
-            lineage_details = LineageDetails(
-                pipeline=EntityReference(id=pipeline_entity.id.root, type="pipeline"),
-                source=LineageSource.PipelineLineage,
-            )
+                pipeline_entity = self.metadata.get_by_name(
+                    entity=Pipeline, fqn=pipeline_fqn
+                )
 
-            dbt_models = self.client.get_model_details(
-                job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
-            )
+                lineage_details = LineageDetails(
+                    pipeline=EntityReference(
+                        id=pipeline_entity.id.root, type="pipeline"
+                    ),
+                    source=LineageSource.PipelineLineage,
+                )
 
-            for model in dbt_models or []:
-                for dbservicename in (
-                    self.source_config.lineageInformation.dbServiceNames or []
-                ):
-                    to_entity = self.metadata.get_by_name(
-                        entity=Table,
-                        fqn=fqn.build(
-                            metadata=self.metadata,
-                            entity_type=Table,
-                            table_name=model.alias,
-                            database_name=model.database,
-                            schema_name=model.dbtschema,
-                            service_name=dbservicename,
-                        ),
-                    )
+                dbt_models = self.client.get_model_details(
+                    job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
+                )
 
-                    if to_entity is None:
-                        continue
+                dbt_parents = self.client.get_models_and_seeds_details(
+                    job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
+                )
 
-                    for dest in model.parentsSources or []:
-                        from_entity = self.metadata.get_by_name(
+                for model in dbt_models or []:
+                    for dbservicename in (
+                        self.source_config.lineageInformation.dbServiceNames or []
+                    ):
+                        to_entity = self.metadata.get_by_name(
                             entity=Table,
                             fqn=fqn.build(
                                 metadata=self.metadata,
                                 entity_type=Table,
-                                table_name=dest.name,
-                                database_name=dest.database,
-                                schema_name=dest.dbtschema,
+                                table_name=model.name,
+                                database_name=model.database,
+                                schema_name=model.dbtschema,
                                 service_name=dbservicename,
                             ),
                         )
 
-                        if from_entity is None:
+                        if to_entity is None:
                             continue
 
-                        yield Either(
-                            right=AddLineageRequest(
-                                edge=EntitiesEdge(
-                                    fromEntity=EntityReference(
-                                        id=from_entity.id,
-                                        type="table",
-                                    ),
-                                    toEntity=EntityReference(
-                                        id=to_entity.id,
-                                        type="table",
+                        for unique_id in model.dependsOn or []:
+                            parents = [
+                                d for d in dbt_parents if d.uniqueId == unique_id
+                            ]
+                            if parents:
+                                from_entity = self.metadata.get_by_name(
+                                    entity=Table,
+                                    fqn=fqn.build(
+                                        metadata=self.metadata,
+                                        entity_type=Table,
+                                        table_name=parents[0].name,
+                                        database_name=parents[0].database,
+                                        schema_name=parents[0].dbtschema,
+                                        service_name=dbservicename,
                                     ),
-                                    lineageDetails=lineage_details,
                                 )
-                            )
-                        )
+
+                                if from_entity is None:
+                                    continue
+
+                                yield Either(
+                                    right=AddLineageRequest(
+                                        edge=EntitiesEdge(
+                                            fromEntity=EntityReference(
+                                                id=from_entity.id,
+                                                type="table",
+                                            ),
+                                            toEntity=EntityReference(
+                                                id=to_entity.id,
+                                                type="table",
+                                            ),
+                                            lineageDetails=lineage_details,
+                                        )
+                                    )
+                                )
 
         except Exception as exc:
             yield Either(
@@ -238,7 +256,6 @@ class DbtcloudSource(PipelineServiceSource):
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Failed to get pipeline list due to : {exc}")
-        return None
 
     def get_pipeline_name(self, pipeline_details: DBTJob) -> str:
         """
@@ -263,21 +280,27 @@ class DbtcloudSource(PipelineServiceSource):
                 TaskStatus(
                     name=str(task.id),
                     executionStatus=STATUS_MAP.get(task.state, StatusType.Pending),
-                    startTime=Timestamp(
-                        datetime_to_ts(
-                            datetime.strptime(task.started_at, "%Y-%m-%d %H:%M:%S.%f%z")
-                            if task.started_at
-                            else datetime.now()
+                    startTime=(
+                        Timestamp(
+                            datetime_to_ts(
+                                datetime.strptime(
+                                    task.started_at, "%Y-%m-%d %H:%M:%S.%f%z"
+                                )
+                            )
                         )
+                        if task.started_at
+                        else None
                     ),
-                    endTime=Timestamp(
-                        datetime_to_ts(
-                            datetime.strptime(
-                                task.finished_at, "%Y-%m-%d %H:%M:%S.%f%z"
+                    endTime=(
+                        Timestamp(
+                            datetime_to_ts(
+                                datetime.strptime(
+                                    task.finished_at, "%Y-%m-%d %H:%M:%S.%f%z"
+                                )
                             )
-                            if task.finished_at
-                            else datetime.now()
                         )
+                        if task.finished_at
+                        else None
                     ),
                 )
                 for task in self.client.get_runs(job_id=int(pipeline_details.id)) or []
@@ -294,7 +317,7 @@ class DbtcloudSource(PipelineServiceSource):
                             pipeline_details.created_at, "%Y-%m-%dT%H:%M:%S.%f%z"
                         )
                         if pipeline_details.created_at
-                        else datetime.now()
+                        else None
                     )
                 ),
             )

metadata/ingestion/source/pipeline/dbtcloud/models.py

@@ -35,7 +35,7 @@ class DBTJob(BaseModel):
 
 
 class DBTJobList(BaseModel):
-    Jobs: Optional[List[DBTJob]] = Field(None, alias="data")
+    Jobs: Optional[List[DBTJob]] = Field([], alias="data")
 
 
 class DBTRun(BaseModel):
@@ -50,7 +50,7 @@ class DBTRun(BaseModel):
 
 
 class DBTRunList(BaseModel):
-    Runs: Optional[List[DBTRun]] = Field(None, alias="data")
+    Runs: Optional[List[DBTRun]] = Field([], alias="data")
 
 
 class DBTSources(BaseModel):
@@ -60,13 +60,13 @@ class DBTSources(BaseModel):
 
 
 class DBTModel(BaseModel):
+    uniqueId: Optional[str] = None
     name: Optional[str] = None
-    alias: Optional[str] = None
     dbtschema: Optional[str] = Field(None, alias="schema")
     database: Optional[str] = None
-    rawSql: Optional[str] = None
-    parentsSources: Optional[List[DBTSources]] = None
+    dependsOn: Optional[List[str]] = None
 
 
 class DBTModelList(BaseModel):
-    models: Optional[List[DBTModel]] = None
+    models: Optional[List[DBTModel]] = []
+    seeds: Optional[List[DBTModel]] = []

metadata/ingestion/source/pipeline/dbtcloud/queries.py

@@ -12,22 +12,33 @@
 GraphQL Queries used during ingestion
 """
 
-DBT_QUERY = """
+DBT_GET_MODEL_DEPENDS_ON = """
 query Query($jobId: BigInt!, $runId: BigInt) {
   job(id: $jobId, runId: $runId) {
     models {
-      name #destinationTable
-      alias
+      name
       database
       schema
-      rawSql
-      materializedType
-      parentsSources {
-        database
-        name
-        schema
-        sourceName #sourceTable
-      }
+      dependsOn
+    }
+  }
+}
+"""
+
+DBT_GET_MODELS_SEEDS = """
+query Query($jobId: BigInt!, $runId: BigInt) {
+  job(id: $jobId, runId: $runId) {
+    models {
+      uniqueId
+      name
+      database
+      schema
+    }
+    seeds {
+      uniqueId
+      name
+      schema
+      database
     }
   }
 }

metadata/ingestion/source/pipeline/domopipeline/metadata.py

@@ -89,7 +89,7 @@ class DomopipelineSource(PipelineServiceSource):
         self, pipeline_details
     ) -> Iterable[Either[CreatePipelineRequest]]:
         try:
-            pipeline_name = pipeline_details["id"]
+            pipeline_name = str(pipeline_details["id"])
             source_url = self.get_source_url(pipeline_id=pipeline_name)
             task = Task(
                 name=pipeline_name,

metadata/ingestion/source/pipeline/fivetran/client.py

@@ -20,6 +20,7 @@ from metadata.generated.schema.entity.services.connections.pipeline.fivetranConn
     FivetranConnection,
 )
 from metadata.ingestion.ometa.client import REST, ClientConfig
+from metadata.utils.helpers import clean_uri
 
 
 class FivetranClient:
@@ -36,7 +37,7 @@ class FivetranClient:
         )
 
         client_config: ClientConfig = ClientConfig(
-            base_url=str(self.config.hostPort),
+            base_url=clean_uri(str(self.config.hostPort)),
             api_version="v1",
             auth_header="Authorization",
             auth_token=lambda: (api_token[2:-1], 0),

metadata/ingestion/source/storage/gcs/client.py

@@ -16,8 +16,6 @@ from google import auth
 from google.cloud.monitoring_v3 import MetricServiceClient
 from google.cloud.storage import Client
 
-NoProject = object()
-
 
 class MultiProjectClient:
     """Google Cloud Client does not support ad-hoc project switching. This class wraps the client and allows
@@ -37,19 +35,16 @@ class MultiProjectClient:
         project_ids: Optional[List[str]] = None,
         **client_kwargs,
     ):
+        self.default_project = None
         if project_ids:
             self.clients = {
                 project_id: client_class(project=project_id, **client_kwargs)
                 for project_id in project_ids
             }
         else:
-            self.clients = {NoProject: client_class(**client_kwargs)}
-
-    def project_ids(self):
-        if NoProject in self.clients:
             _, project_id = auth.default()
-            return [project_id]
-        return list(self.clients.keys())
+            self.default_project = project_id
+            self.clients = {project_id: client_class(**client_kwargs)}
 
     def __getattr__(self, client_method):
         """Return the underlying client method as a partial function so we can inject the project_id."""
@@ -57,7 +52,7 @@ class MultiProjectClient:
 
     def _call(self, method, project_id, *args, **kwargs):
         """Call the method on the client for the given project_id. The args and kwargs are passed through."""
-        client = self.clients.get(project_id, self.clients.get(NoProject))
+        client = self.clients.get(project_id, self.clients.get(self.default_project))
         if not client:
             raise ValueError(f"Project {project_id} not found")
         return getattr(client, method)(*args, **kwargs)

metadata/ingestion/source/storage/gcs/metadata.py

@@ -259,6 +259,10 @@ class GcsSource(StorageServiceSource):
             try:
                 bucket = client.get_bucket(bucket_name)
             except NotFound:
+                logger.warning(
+                    f"Bucket {bucket_name} not found in project {project_id}"
+                )
+                self.status.warning(f"{project_id}.{bucket_name}", "Bucket Not Found")
                 continue
             return GCSBucketResponse(
                 name=bucket.name,