openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

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

@@ -13,7 +13,7 @@
 Databricks pipeline Source Model module
 """
 
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, Field
 
@@ -27,11 +27,19 @@ class DependentTask(BaseModel):
     name: Optional[str] = Field(None, alias="task_key")
 
 
+class PipelineTask(BaseModel):
+    pipeline_id: Optional[str] = None
+    full_refresh: Optional[bool] = None
+
+
 class DBTasks(BaseModel):
     name: Optional[str] = Field(None, alias="task_key")
     description: Optional[str] = None
     depends_on: Optional[List[DependentTask]] = None
     run_page_url: Optional[str] = None
+    pipeline_task: Optional[PipelineTask] = None
+    notebook_task: Optional[Dict[str, Any]] = None
+    spark_python_task: Optional[Dict[str, Any]] = None
 
 
 class DBSettings(BaseModel):
@@ -41,13 +49,21 @@ class DBSettings(BaseModel):
     description: Optional[str] = None
     schedule: Optional[DBRunSchedule] = None
     task_type: Optional[str] = Field(None, alias="format")
+    tasks: Optional[List[DBTasks]] = None
 
 
 class DataBrickPipelineDetails(BaseModel):
-    job_id: int
+    job_id: Optional[int] = None
+    pipeline_id: Optional[str] = None
     creator_user_name: Optional[str] = None
     settings: Optional[DBSettings] = None
-    created_time: int
+    created_time: Optional[int] = None
+    name: Optional[str] = None
+    pipeline_type: Optional[str] = None
+
+    @property
+    def id(self) -> str:
+        return str(self.pipeline_id) if self.pipeline_id else str(self.job_id)
 
 
 class DBRunState(BaseModel):

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

@@ -13,7 +13,7 @@ Client to interact with DBT Cloud REST APIs
 """
 
 import traceback
-from typing import List, Optional
+from typing import Iterable, List, Optional, Tuple
 
 from metadata.generated.schema.entity.services.connections.pipeline.dbtCloudConnection import (
     DBTCloudConnection,
@@ -22,13 +22,13 @@ from metadata.ingestion.ometa.client import REST, ClientConfig
 from metadata.ingestion.source.pipeline.dbtcloud.models import (
     DBTJob,
     DBTJobList,
+    DBTModel,
     DBTModelList,
     DBTRun,
     DBTRunList,
 )
 from metadata.ingestion.source.pipeline.dbtcloud.queries import (
-    DBT_GET_MODEL_DEPENDS_ON,
-    DBT_GET_MODELS_SEEDS,
+    DBT_GET_MODELS_WITH_LINEAGE,
 )
 from metadata.utils.constants import AUTHORIZATION_HEADER
 from metadata.utils.helpers import clean_uri
@@ -69,14 +69,15 @@ class DBTCloudClient:
         self.graphql_client = REST(graphql_client_config)
 
     def _get_jobs(
-        self, job_id: str = None, project_id: str = None
-    ) -> Optional[List[DBTJob]]:
+        self,
+        job_id: str = None,
+        project_id: str = None,
+    ) -> Iterable[DBTJob]:
         """
-        fetch jobs for an account in dbt cloud
+        Fetch jobs for an account in dbt cloud
         """
-        job_list = []
-        # we will get 100 jobs at a time
         query_params = {"offset": 0, "limit": 100}
+
         try:
             job_path = f"{job_id}/" if job_id else ""
             project_path = f"?project_id={project_id}" if project_id else ""
@@ -86,10 +87,10 @@ class DBTCloudClient:
             )
 
             if job_id:
-                job_list = [DBTJob.model_validate(result["data"])]
+                yield DBTJob.model_validate(result["data"])
             else:
                 job_list_response = DBTJobList.model_validate(result)
-                job_list = job_list_response.Jobs
+                yield from job_list_response.Jobs
 
                 while job_list_response.extra and job_list_response.extra.pagination:
                     total_count = job_list_response.extra.pagination.total_count
@@ -104,14 +105,13 @@ class DBTCloudClient:
                         data=query_params,
                     )
                     job_list_response = DBTJobList.model_validate(result)
-                    job_list.extend(job_list_response.Jobs)
+                    yield from job_list_response.Jobs
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(
                 f"Failed to get job info for project_id: `{project_id}` or job_id: `{job_id}` : {exc}"
             )
-        return job_list
 
     def test_get_jobs(self) -> List[DBTJob]:
         """
@@ -128,47 +128,68 @@ class DBTCloudClient:
         run_list = DBTRunList.model_validate(result).Runs
         return run_list
 
-    def get_jobs(self) -> Optional[List[DBTJob]]:
+    def get_jobs(self) -> Iterable[DBTJob]:
         """
-        list jobs for an account in dbt cloud
+        List jobs for an account in dbt cloud using generator pattern.
+        yields job one at a time for memory efficiency.
         """
         try:
-            jobs = []
             # case when job_ids are specified and project_ids are not
             if self.job_ids and not self.project_ids:
                 for job_id in self.job_ids:
-                    jobs.extend(self._get_jobs(job_id=job_id))
+                    yield from self._get_jobs(job_id=job_id)
             # case when project_ids are specified or both are specified
             elif self.project_ids:
                 for project_id in self.project_ids:
-                    results = self._get_jobs(project_id=project_id)
-                    if self.job_ids:
-                        jobs.extend(
-                            [
-                                result
-                                for result in results
-                                if str(result.id) in self.job_ids
-                            ]
-                        )
-                    else:
-                        jobs.extend(results)
+                    for job in self._get_jobs(project_id=project_id):
+                        if self.job_ids:
+                            if str(job.id) in self.job_ids:
+                                yield job
+                        else:
+                            yield job
             else:
-                results = self._get_jobs()
-                jobs.extend(results)
-            return jobs
+                yield from self._get_jobs()
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Unable to get job info :{exc}")
 
-        return None
+    def get_latest_successful_run_id(self, job_id: int) -> Optional[int]:
+        """
+        Get the latest successful run ID for a given job.
+        """
+        try:
+            query_params = {
+                "job_definition_id": job_id,
+                "order_by": "-finished_at",
+                "limit": "1",
+                "status": "10",  # 10 = Success in dbt Cloud API
+            }
 
-    def get_runs(self, job_id: int) -> Optional[List[DBTRun]]:
+            result = self.client.get(
+                f"/accounts/{self.config.accountId}/runs/", data=query_params
+            )
+            run_list_response = DBTRunList.model_validate(result)
+
+            if run_list_response.Runs:
+                return run_list_response.Runs[0].id
+
+            return None
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unable to get latest successful run for job {job_id}: {exc}"
+            )
+            return None
+
+    def get_runs(self, job_id: int) -> Iterable[DBTRun]:
         """
-        list runs for a job in dbt cloud
+        List runs for a job in dbt cloud using generator pattern.
+        yields run one at a time for memory efficiency.
         """
         try:
             number_of_runs = self.config.numberOfRuns
-            runs = []
+            yielded_count = 0
             # we will get 100 runs at a time and order by created_at in descending order
             query_params = {
                 "job_definition_id": job_id,
@@ -181,13 +202,17 @@ class DBTCloudClient:
                 f"/accounts/{self.config.accountId}/runs/", data=query_params
             )
             run_list_response = DBTRunList.model_validate(result)
-            runs.extend(run_list_response.Runs)
 
-            while (
-                (number_of_runs is None or len(runs) < number_of_runs)
-                and run_list_response.extra
-                and run_list_response.extra.pagination
-            ):
+            for run in run_list_response.Runs or []:
+                if number_of_runs is not None and yielded_count >= number_of_runs:
+                    return
+                yield run
+                yielded_count += 1
+
+            while run_list_response.extra and run_list_response.extra.pagination:
+                if number_of_runs is not None and yielded_count >= number_of_runs:
+                    return
+
                 total_count = run_list_response.extra.pagination.total_count
                 current_count = run_list_response.extra.pagination.count
 
@@ -200,60 +225,44 @@ class DBTCloudClient:
                     data=query_params,
                 )
                 run_list_response = DBTRunList.model_validate(result)
-                runs.extend(run_list_response.Runs)
 
-            return runs[:number_of_runs] if number_of_runs is not None else runs
+                for run in run_list_response.Runs or []:
+                    if number_of_runs is not None and yielded_count >= number_of_runs:
+                        return
+                    yield run
+                    yielded_count += 1
+
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unable to get run info :{exc}")
 
-        return None
-
-    def get_model_details(self, job_id: int, run_id: int):
+    def get_models_with_lineage(
+        self, job_id: int, run_id: int
+    ) -> Tuple[
+        Optional[List[DBTModel]], Optional[List[DBTModel]], Optional[List[DBTModel]]
+    ]:
         """
-        get model details for a job in dbt cloud for lineage
+        Get models with dependsOn and seeds in a single GraphQL call.
         """
         try:
-            query_params = {
-                "query": DBT_GET_MODEL_DEPENDS_ON,
-                "variables": {"jobId": job_id, "runId": run_id},
-            }
-
-            result = self.graphql_client.post("", json=query_params)
-
-            if result.get("data") and result["data"].get("job"):
-                model_list = DBTModelList.model_validate(result["data"]["job"]).models
-                logger.debug(
-                    f"Successfully fetched models from dbt for job_id:{job_id} run_id:{run_id}: {model_list}"
-                )
-                return model_list
 
-        except Exception as exc:
-            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 parent model details for a job in dbt cloud for lineage
-        """
-        try:
             query_params = {
-                "query": DBT_GET_MODELS_SEEDS,
+                "query": DBT_GET_MODELS_WITH_LINEAGE,
                 "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.model_validate(result["data"]["job"])
-                parents_list = result.models + result.seeds
+                model_list = DBTModelList.model_validate(result["data"]["job"])
                 logger.debug(
-                    f"Successfully fetched parent models from dbt for job_id:{job_id} run_id:{run_id}: {parents_list}"
+                    f"Successfully fetched models and seeds from dbt for "
+                    f"job_id:{job_id} run_id:{run_id}: "
+                    f"models={len(model_list.models or [])}, seeds={len(model_list.seeds or [])}, sources={len(model_list.sources or [])}"
                 )
-                return parents_list
+                return model_list.models, model_list.seeds, model_list.sources
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Unable to get parents model info :{exc}")
-        return None
+            logger.warning(f"Unable to get models with lineage info: {exc}")
+        return None, None, None

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

@@ -13,7 +13,7 @@ DBTcloud source to extract metadata from OM UI
 """
 import traceback
 from datetime import datetime
-from typing import Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional
 
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
@@ -87,6 +87,21 @@ class DbtcloudSource(PipelineServiceSource):
             )
         return cls(config, metadata)
 
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        # Cache for table entity lookups to avoid redundant API calls
+        self._table_entity_cache: Dict[str, Optional[Table]] = {}
+
+    def _get_table_entity(self, table_fqn: str) -> Optional[Table]:
+        """
+        Cached table entity lookup to avoid redundant API calls.
+        """
+        if table_fqn not in self._table_entity_cache:
+            self._table_entity_cache[table_fqn] = self.metadata.get_by_name(
+                entity=Table, fqn=table_fqn
+            )
+        return self._table_entity_cache[table_fqn]
+
     def _get_task_list(self, job_id: int) -> Optional[List[Task]]:
         """
         Method to collect all the tasks from dbt cloud job and return it in a task list
@@ -94,19 +109,21 @@ class DbtcloudSource(PipelineServiceSource):
         self.context.get().latest_run_id = None
         try:
             task_list: List[Task] = []
-            runs = self.client.get_runs(job_id=job_id)
-            if runs:
-                for run in runs or []:
-                    task = Task(
-                        name=str(run.id),
-                        sourceUrl=SourceUrl(run.href),
-                        startDate=str(run.started_at),
-                        endDate=str(run.finished_at),
-                    )
-                    task_list.append(task)
-                self.context.get().latest_run_id = (
-                    task_list[0].name if task_list else None
+            runs_list: List = []
+            # Consume generator and store runs for later use
+            for run in self.client.get_runs(job_id=job_id):
+                runs_list.append(run)
+                task = Task(
+                    name=str(run.id),
+                    sourceUrl=SourceUrl(run.href),
+                    startDate=str(run.started_at),
+                    endDate=str(run.finished_at),
                 )
+                task_list.append(task)
+
+            if task_list:
+                self.context.get().latest_run_id = runs_list[0].id
+
             return task_list or None
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -153,90 +170,124 @@ class DbtcloudSource(PipelineServiceSource):
         self, pipeline_details: DBTJob
     ) -> Iterable[Either[AddLineageRequest]]:
         """
-        Get lineage between pipeline and data sources
+        Get lineage between pipeline and data sources.
+        Uses combined GraphQL call for models and seeds, with optimized caching.
         """
         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,
-                )
+            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
-                )
+            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
-                )
+            if not pipeline_entity:
+                logger.warning(f"Pipeline entity not found for FQN: {pipeline_fqn}")
+                return
 
-                dbt_parents = self.client.get_models_and_seeds_details(
-                    job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
-                )
+            # Use combined GraphQL call instead of two separate calls
+            (dbt_models, dbt_seeds, dbt_sources,) = self.client.get_models_with_lineage(
+                job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
+            )
 
-                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.name,
-                                database_name=model.database,
-                                schema_name=model.dbtschema,
-                                service_name=dbservicename,
-                            ),
-                        )
+            # Combine models and seeds for parent lookup
+            dbt_parents = (dbt_models or []) + (dbt_seeds or []) + (dbt_sources or [])
 
-                        if to_entity is None:
+            # Build parent lookup dict for O(1) access instead of O(n) list search
+            parent_by_unique_id = {p.uniqueId: p for p in dbt_parents if p.uniqueId}
+
+            for model in dbt_models or []:
+                if not model.runGeneratedAt:
+                    logger.debug(
+                        f"Skipping model with missing runGeneratedAt: name={getattr(model, 'name', None)}"
+                    )
+                    continue
+                if not all([model.name, model.database, model.dbtschema]):
+                    logger.debug(
+                        f"Skipping model with missing attributes: name={getattr(model, 'name', None)}, "
+                        f"database={getattr(model, 'database', None)}, schema={getattr(model, 'dbtschema', None)}"
+                    )
+                    continue
+
+                for dbservicename in self.get_db_service_names() or ["*"]:
+                    to_entity_fqn = fqn.build(
+                        metadata=self.metadata,
+                        entity_type=Table,
+                        table_name=model.name,
+                        database_name=model.database,
+                        schema_name=model.dbtschema,
+                        service_name=dbservicename,
+                    )
+
+                    # Use cached table entity lookup
+                    to_entity = self._get_table_entity(to_entity_fqn)
+
+                    if to_entity is None:
+                        continue
+
+                    for unique_id in model.dependsOn or []:
+                        # Use dict lookup instead of list comprehension
+                        parent = parent_by_unique_id.get(unique_id)
+                        if not parent:
                             continue
 
-                        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,
-                                    ),
-                                )
+                        # Check runGeneratedAt for models and seeds (not sources)
+                        # Sources are auto-generated and don't have runGeneratedAt
+                        is_source = unique_id.startswith("source.")
+                        if not is_source and not parent.runGeneratedAt:
+                            logger.debug(
+                                f"Skipping parent with missing runGeneratedAt: uniqueId={unique_id}"
+                            )
+                            continue
 
-                                if from_entity is None:
-                                    continue
+                        if not all([parent.name, parent.database, parent.dbtschema]):
+                            logger.debug(
+                                f"Skipping parent with missing attributes: name={getattr(parent, 'name', None)}, "
+                                f"database={getattr(parent, 'database', None)}, schema={getattr(parent, 'dbtschema', None)}"
+                            )
+                            continue
 
-                                lineage_details = LineageDetails(
-                                    pipeline=EntityReference(
-                                        id=pipeline_entity.id.root, type="pipeline"
-                                    ),
-                                    source=LineageSource.PipelineLineage,
-                                )
+                        from_entity_fqn = fqn.build(
+                            metadata=self.metadata,
+                            entity_type=Table,
+                            table_name=parent.name,
+                            database_name=parent.database,
+                            schema_name=parent.dbtschema,
+                            service_name=dbservicename,
+                        )
+
+                        # Use cached table entity lookup
+                        from_entity = self._get_table_entity(from_entity_fqn)
+
+                        if from_entity is None:
+                            continue
+
+                        lineage_details = LineageDetails(
+                            pipeline=EntityReference(
+                                id=pipeline_entity.id.root, type="pipeline"
+                            ),
+                            source=LineageSource.PipelineLineage,
+                        )
 
-                                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,
-                                        )
-                                    )
+                        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(

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

@@ -65,9 +65,11 @@ class DBTRunList(BaseModel):
 
 
 class DBTSources(BaseModel):
+    uniqueId: Optional[str] = None
     name: Optional[str] = None
     dbtschema: Optional[str] = Field(None, alias="schema")
     database: Optional[str] = None
+    runGeneratedAt: Optional[str] = None
     extra: Optional[Extra] = None
 
 
@@ -76,10 +78,12 @@ class DBTModel(BaseModel):
     name: Optional[str] = None
     dbtschema: Optional[str] = Field(None, alias="schema")
     database: Optional[str] = None
+    runGeneratedAt: Optional[str] = None
     dependsOn: Optional[List[str]] = None
 
 
 class DBTModelList(BaseModel):
     models: Optional[List[DBTModel]] = []
     seeds: Optional[List[DBTModel]] = []
+    sources: Optional[List[DBTModel]] = []
     extra: Optional[Extra] = None

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

@@ -12,29 +12,25 @@
 GraphQL Queries used during ingestion
 """
 
-DBT_GET_MODEL_DEPENDS_ON = """
+DBT_GET_MODELS_WITH_LINEAGE = """
 query Query($jobId: BigInt!, $runId: BigInt) {
   job(id: $jobId, runId: $runId) {
     models {
+      uniqueId
       name
       database
       schema
       dependsOn
+      runGeneratedAt
     }
-  }
-}
-"""
-
-DBT_GET_MODELS_SEEDS = """
-query Query($jobId: BigInt!, $runId: BigInt) {
-  job(id: $jobId, runId: $runId) {
-    models {
+    seeds {
       uniqueId
       name
-      database
       schema
+      database
+      runGeneratedAt
     }
-    seeds {
+    sources {
       uniqueId
       name
       schema