openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/trino/profiler/system_tables_profiler.py

@@ -0,0 +1,167 @@
+#  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.
+"""
+System table profiler
+"""
+from datetime import datetime
+from decimal import Decimal
+from typing import Any, Dict, List, Optional, Set, Union
+
+from more_itertools import partition
+from pydantic import field_validator
+from sqlalchemy import Table, text
+from sqlalchemy.engine import Engine
+
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.profiler.interface.sqlalchemy.stored_statistics_profiler import (
+    StoredStatisticsSource,
+)
+from metadata.profiler.metrics.core import Metric
+from metadata.profiler.metrics.registry import Metrics
+from metadata.utils.logger import profiler_logger
+from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
+from metadata.utils.ssl_manager import get_ssl_connection
+
+logger = profiler_logger()
+
+
+class ColumnStats(BaseModel):
+    """Based on https://trino.io/docs/current/sql/show-stats.html"""
+
+    column_name: Optional[str] = None
+    data_size: Optional[int] = None
+    distinct_values_count: Optional[int] = None
+    nulls_fraction: Optional[float] = None
+    low_value: Optional[Union[int, float, datetime, Decimal]] = None
+    high_value: Optional[Union[int, float, datetime, Decimal]] = None
+
+    @field_validator("data_size", mode="before")
+    @classmethod
+    def data_size_validator(cls, value):
+        """Data size validator
+
+        Args:
+            value: value
+        """
+        if value is None:
+            return None
+        return int(value)
+
+
+class TableStats(BaseModel):
+    row_count: Optional[int] = None
+    columns: Dict[str, ColumnStats] = {}
+
+
+class TrinoStoredStatisticsSource(StoredStatisticsSource):
+    """Trino system profile source"""
+
+    metric_stats_map: Dict[Metrics, str] = {
+        Metrics.NULL_RATIO: "nulls_fractions",
+        Metrics.DISTINCT_COUNT: "distinct_values_count",
+        Metrics.ROW_COUNT: "row_count",
+        Metrics.MAX: "high_value",
+        Metrics.MIN: "low_value",
+    }
+
+    metric_stats_by_name: Dict[str, str] = {
+        k.name: v for k, v in metric_stats_map.items()
+    }
+
+    def get_statistics_metrics(self) -> Set[Metrics]:
+        return set(self.metric_stats_map.keys())
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        service_connection_config = kwargs["service_connection_config"]
+        self.session: Engine = get_ssl_connection(service_connection_config)
+        self.stats_cache = LRUCache(capacity=LRU_CACHE_SIZE)
+
+    def get_column_statistics(
+        self, metric: List[Metric], schema: str, table_name: Table, column: str
+    ) -> Dict[str, Any]:
+        table_stats = self._get_cached_stats(schema, table_name)
+        try:
+            column_stats = table_stats.columns[column]
+        except KeyError:
+            raise RuntimeError(
+                f"Column {column} not found in table {table_name}. Statistics might be stale or missing."
+            )
+        result = {
+            m.name(): getattr(column_stats, self.metric_stats_by_name[m.name()])
+            for m in metric
+        }
+        result.update(self.get_hybrid_statistics(table_stats, column_stats))
+        self.warn_for_missing_stats(schema, table_name, column_stats)
+        return result
+
+    def get_table_statistics(
+        self, metric: List[Metric], schema: str, table_name: Table
+    ) -> dict:
+        table_stats = self._get_cached_stats(schema, table_name)
+        return {
+            m.name(): getattr(table_stats, self.metric_stats_by_name[m.name()])
+            for m in metric
+        }
+
+    def warn_for_missing_stats(self, schema: str, table: str, stats: BaseModel):
+        if (
+            isinstance(stats, ColumnStats)
+            and all(map(lambda x: x is None, stats.model_dump().values()))
+        ) or (
+            isinstance(stats, TableStats)
+            and all(
+                map(
+                    lambda x: x is None,
+                    [v for k, v in stats.model_dump().items() if k != "columns"],
+                )
+            )
+        ):
+            logger.warning(
+                'Statistics are missing for table "{schema}.{table}". Profiling might be inaccurate.\n'
+                "Gather statistics for the table by running:"
+                f"  ANALYZE {schema}.{table}"
+            )
+
+    def _get_cached_stats(self, schema: str, table: str) -> TableStats:
+        path = f"{schema}.{table}"
+        if path in self.stats_cache:
+            return self.stats_cache.get(path)
+        stats = self._get_db_stats(schema, table)
+        self.stats_cache.put(path, stats)
+        return stats
+
+    def _get_db_stats(self, schema, table) -> TableStats:
+        rows = self.session.execute(text(f'SHOW STATS FOR "{schema}"."{table}"'))
+        table_rows, column_rows = map(
+            list, partition(lambda row: row.get("column_name"), map(dict, rows))
+        )
+        if len(table_rows) != 1:
+            raise RuntimeError(
+                f"Expected one row for table {table}, got {len(table_rows)}"
+            )
+        table = table_rows[0]
+        columns_dict = {
+            row.get("column_name"): ColumnStats(**row) for row in column_rows
+        }
+        return TableStats(row_count=table["row_count"], columns=columns_dict)
+
+    def get_hybrid_statistics(
+        self, table_stats: TableStats, column_stats: ColumnStats
+    ) -> Dict[str, Any]:
+        return {
+            # trino stats are in fractions, so we need to convert them to counts (unlike our default profiler)
+            Metrics.NULL_COUNT.name: (
+                int(table_stats.row_count * column_stats.nulls_fraction)
+                if None not in [table_stats.row_count, column_stats.nulls_fraction]
+                else None
+            ),
+        }

metadata/ingestion/source/database/trino/service_spec.py

@@ -0,0 +1,16 @@
+from metadata.ingestion.source.database.trino.lineage import TrinoLineageSource
+from metadata.ingestion.source.database.trino.metadata import TrinoSource
+from metadata.ingestion.source.database.trino.usage import TrinoUsageSource
+from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import (
+    TrinoProfilerInterface,
+)
+from metadata.sampler.sqlalchemy.trino.sampler import TrinoSampler
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=TrinoSource,
+    lineage_source_class=TrinoLineageSource,
+    usage_source_class=TrinoUsageSource,
+    profiler_class=TrinoProfilerInterface,
+    sampler_class=TrinoSampler,
+)

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

@@ -13,7 +13,7 @@ Databricks Unity Catalog Source source methods.
 """
 import json
 import traceback
-from typing import Any, Iterable, List, Optional, Tuple, Union
+from typing import Any, Iterable, List, Optional, Tuple
 
 from databricks.sdk.service.catalog import ColumnInfo
 from databricks.sdk.service.catalog import TableConstraint as DBTableConstraint
@@ -22,12 +22,10 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
-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 (
@@ -76,7 +74,6 @@ from metadata.ingestion.source.database.unitycatalog.models import (
 )
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
-from metadata.utils.db_utils import get_view_lineage
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.logger import ingestion_logger
 
@@ -325,7 +322,7 @@ class UnitycatalogSource(
             ) = self.get_table_constraints(table.table_constraints)
 
             table_constraints = self.update_table_constraints(
-                primary_constraints, foreign_constraints
+                primary_constraints, foreign_constraints, columns
             )
 
             table_request = CreateTableRequest(
@@ -436,7 +433,7 @@ class UnitycatalogSource(
         return table_constraints
 
     def update_table_constraints(
-        self, table_constraints, foreign_columns
+        self, table_constraints, foreign_columns, columns
     ) -> List[TableConstraint]:
         """
         From topology.
@@ -521,18 +518,6 @@ class UnitycatalogSource(
             )
             yield parsed_column
 
-    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
-        logger.info("Processing Lineage for Views")
-        for view in [
-            v for v in self.context.get().table_views if v.view_definition is not None
-        ]:
-            yield from get_view_lineage(
-                view=view,
-                metadata=self.metadata,
-                service_name=self.context.get().database_service,
-                connection_type=self.service_connection.type.value,
-            )
-
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
@@ -549,12 +534,6 @@ class UnitycatalogSource(
     def get_stored_procedure_queries(self) -> Iterable[QueryByProcedure]:
         """Not Implemented"""
 
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Not Implemented"""
-        yield from []
-
     def close(self):
         """Nothing to close"""
 

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

@@ -0,0 +1,22 @@
+from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
+    UnityCatalogTestSuiteInterface,
+)
+from metadata.ingestion.source.database.unitycatalog.lineage import (
+    UnitycatalogLineageSource,
+)
+from metadata.ingestion.source.database.unitycatalog.metadata import UnitycatalogSource
+from metadata.ingestion.source.database.unitycatalog.usage import (
+    UnitycatalogUsageSource,
+)
+from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
+    UnityCatalogProfilerInterface,
+)
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=UnitycatalogSource,
+    lineage_source_class=UnitycatalogLineageSource,
+    usage_source_class=UnitycatalogUsageSource,
+    profiler_class=UnityCatalogProfilerInterface,
+    test_suite_class=UnityCatalogTestSuiteInterface,
+)

metadata/ingestion/source/database/usage_source.py

@@ -152,7 +152,10 @@ class UsageSource(QueryParserSource, ABC):
             except Exception as exc:
                 if query:
                     logger.debug(
-                        f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}\n##########################"
+                        (
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}"
+                            "\n##########################"
+                        )
                     )
                 logger.debug(traceback.format_exc())
                 logger.error(f"Source usage processing error: {exc}")

metadata/ingestion/source/database/vertica/service_spec.py

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.vertica.lineage import VerticaLineageSource
+from metadata.ingestion.source.database.vertica.metadata import VerticaSource
+from metadata.ingestion.source.database.vertica.usage import VerticaUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=VerticaSource,
+    lineage_source_class=VerticaLineageSource,
+    usage_source_class=VerticaUsageSource,
+)

metadata/ingestion/source/messaging/kafka/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kafka.metadata import KafkaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KafkaSource)

metadata/ingestion/source/messaging/kinesis/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kinesis.metadata import KinesisSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KinesisSource)

metadata/ingestion/source/messaging/redpanda/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.redpanda.metadata import RedpandaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=RedpandaSource)

metadata/ingestion/source/metadata/alationsink/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.alationsink.metadata import AlationsinkSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AlationsinkSource)

metadata/ingestion/source/metadata/amundsen/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.amundsen.metadata import AmundsenSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AmundsenSource)

metadata/ingestion/source/metadata/atlas/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.atlas.metadata import AtlasSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AtlasSource)

metadata/ingestion/source/mlmodel/mlflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.mlflow.metadata import MlflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=MlflowSource)

metadata/ingestion/source/mlmodel/sagemaker/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.sagemaker.metadata import SagemakerSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SagemakerSource)

metadata/ingestion/source/pipeline/airbyte/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airbyte.metadata import AirbyteSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirbyteSource)

metadata/ingestion/source/pipeline/airflow/connection.py

@@ -16,7 +16,9 @@ from functools import partial, singledispatch
 from typing import Optional
 
 from airflow import settings
+from airflow.models.serialized_dag import SerializedDagModel
 from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -102,6 +104,18 @@ def get_connection(connection: AirflowConnection) -> Engine:
         raise SourceConnectionException(msg) from exc
 
 
+class AirflowPipelineDetailsAccessError(Exception):
+    """
+    Raise when Pipeline information is not retrieved
+    """
+
+
+class AirflowTaskDetailsAccessError(Exception):
+    """
+    Raise when Task detail information is not retrieved
+    """
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,
@@ -114,8 +128,37 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"CheckAccess": partial(test_connection_engine_step, engine)}
-
+    session_maker = sessionmaker(bind=engine)
+    session = session_maker()
+
+    def test_pipeline_details_access(session):
+        try:
+            result = session.query(SerializedDagModel).first()
+            return result
+        except Exception as e:
+            raise AirflowPipelineDetailsAccessError(
+                f"Pipeline details access error: {e}"
+            )
+
+    def test_task_detail_access(session):
+        try:
+            json_data_column = (
+                SerializedDagModel._data  # For 2.3.0 onwards # pylint: disable=protected-access
+                if hasattr(SerializedDagModel, "_data")
+                else SerializedDagModel.data  # For 2.2.5 and 2.1.4
+            )
+            result = session.query(json_data_column).first()
+
+            retrieved_tasks = result[0]["dag"]["tasks"]
+            return retrieved_tasks
+        except Exception as e:
+            raise AirflowTaskDetailsAccessError(f"Task details access error : {e}")
+
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, engine),
+        "PipelineDetailsAccess": partial(test_pipeline_details_access, session),
+        "TaskDetailAccess": partial(test_task_detail_access, session),
+    }
     return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from enum import Enum
 from typing import Dict, Iterable, List, Optional, cast
 
-from airflow.models import BaseOperator, DagRun, TaskInstance
+from airflow.models import BaseOperator, DagRun, DagTag, TaskInstance
 from airflow.models.dag import DagModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.serialization.serialized_objects import SerializedDAG
@@ -57,6 +57,7 @@ from metadata.generated.schema.type.entityReferenceList import EntityReferenceLi
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.connections.session import create_and_bind_session
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.airflow.lineage_parser import (
@@ -74,9 +75,12 @@ from metadata.utils import fqn
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 from metadata.utils.helpers import clean_uri, datetime_to_ts
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
 
 logger = ingestion_logger()
 
+AIRFLOW_TAG_CATEGORY = "AirflowTags"
+
 
 class AirflowTaskStatus(Enum):
     SUCCESS = "success"
@@ -152,6 +156,31 @@ class AirflowSource(PipelineServiceSource):
             return task["__var"]
         return task
 
+    def get_all_tags(self, dag_id: str) -> List[str]:
+        try:
+            tag_query = (
+                self.session.query(DagTag.name)
+                .filter(DagTag.dag_id == dag_id)
+                .distinct()
+                .all()
+            )
+            return [tag[0] for tag in tag_query]
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not extract tags details due to {exc}")
+        return []
+
+    def yield_tag(
+        self, pipeline_details: AirflowDagDetails
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        yield from get_ometa_tag_and_classification(
+            tags=self.get_all_tags(dag_id=pipeline_details.dag_id),
+            classification_name=AIRFLOW_TAG_CATEGORY,
+            tag_description="Airflow Tag",
+            classification_description="Tags associated with airflow entities.",
+            include_tags=self.source_config.includeTags,
+        )
+
     def get_pipeline_status(self, dag_id: str) -> List[DagRun]:
         """
         Return the DagRuns of given dag
@@ -328,32 +357,42 @@ class AirflowSource(PipelineServiceSource):
             ).filter(
                 DagModel.is_paused == False  # pylint: disable=singleton-comparison
             )
-        for serialized_dag in session_query.yield_per(100):
-            try:
-                data = serialized_dag[1]["dag"]
-                dag = AirflowDagDetails(
-                    dag_id=serialized_dag[0],
-                    fileloc=serialized_dag[2],
-                    data=AirflowDag.model_validate(serialized_dag[1]),
-                    max_active_runs=data.get("max_active_runs", None),
-                    description=data.get("_description", None),
-                    start_date=data.get("start_date", None),
-                    tasks=list(
-                        map(self._extract_serialized_task, data.get("tasks", []))
-                    ),
-                    schedule_interval=get_schedule_interval(data),
-                    owner=self.fetch_dag_owners(data),
-                )
+        limit = 100  # Number of records per batch
+        offset = 0  # Start
+
+        while True:
+            paginated_query = session_query.limit(limit).offset(offset)
+            results = paginated_query.all()
+            if not results:
+                break
+            for serialized_dag in results:
+                try:
+                    data = serialized_dag[1]["dag"]
+                    dag = AirflowDagDetails(
+                        dag_id=serialized_dag[0],
+                        fileloc=serialized_dag[2],
+                        data=AirflowDag.model_validate(serialized_dag[1]),
+                        max_active_runs=data.get("max_active_runs", None),
+                        description=data.get("_description", None),
+                        start_date=data.get("start_date", None),
+                        tasks=list(
+                            map(self._extract_serialized_task, data.get("tasks", []))
+                        ),
+                        schedule_interval=get_schedule_interval(data),
+                        owner=self.fetch_dag_owners(data),
+                    )
 
-                yield dag
-            except ValidationError as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Error building pydantic model for {serialized_dag} - {err}"
-                )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+                    yield dag
+                except ValidationError as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Error building pydantic model for {serialized_dag} - {err}"
+                    )
+                except Exception as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+
+            offset += limit
 
     def fetch_dag_owners(self, data) -> Optional[str]:
         """
@@ -441,7 +480,6 @@ class AirflowSource(PipelineServiceSource):
         try:
             # Airflow uses /dags/dag_id/grid to show pipeline / dag
             source_url = f"{clean_uri(self.service_connection.hostPort)}/dags/{pipeline_details.dag_id}/grid"
-
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.dag_id),
                 description=Markdown(pipeline_details.description)
@@ -459,6 +497,12 @@ class AirflowSource(PipelineServiceSource):
                 service=FullyQualifiedEntityName(self.context.get().pipeline_service),
                 owners=self.get_owner(pipeline_details.owner),
                 scheduleInterval=pipeline_details.schedule_interval,
+                tags=get_tag_labels(
+                    metadata=self.metadata,
+                    tags=pipeline_details.data.dag.tags,
+                    classification_name=AIRFLOW_TAG_CATEGORY,
+                    include_tags=self.source_config.includeTags,
+                ),
             )
             yield Either(right=pipeline_request)
             self.register_record(pipeline_request=pipeline_request)

metadata/ingestion/source/pipeline/airflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airflow.metadata import AirflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirflowSource)

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

@@ -161,7 +161,9 @@ class DagsterSource(PipelineServiceSource):
                 )
             )
 
-    def yield_tag(self, *_, **__) -> Iterable[Either[OMetaTagAndClassification]]:
+    def yield_tag(
+        self, pipeline_details: DagsterPipeline
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
         yield from get_ometa_tag_and_classification(
             tags=[self.context.get().repository_name],
             classification_name=DAGSTER_TAG_CATEGORY,

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

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.dagster.metadata import DagsterSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DagsterSource)

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

@@ -0,0 +1,6 @@
+from metadata.ingestion.source.pipeline.databrickspipeline.metadata import (
+    DatabrickspipelineSource,
+)
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DatabrickspipelineSource)