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

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

@@ -433,18 +433,41 @@ class AirflowSource(PipelineServiceSource):
         - `owners`: Applied at the tasks. In Airflow's source code, DAG ownership is then a
           list joined with the owners of all the tasks.
 
-        We will pick the owner from the tasks that appears in most tasks.
+        We will pick the owner from the tasks that appears in most tasks,
+        or fall back to the default_args owner if available.
         """
         try:
-            if self.source_config.includeOwners:
-                task_owners = [
-                    task.get("owner")
-                    for task in data.get("tasks", [])
-                    if task.get("owner") is not None
-                ]
-                if task_owners:
-                    most_common_owner, _ = Counter(task_owners).most_common(1)[0]
-                    return most_common_owner
+            if not self.source_config.includeOwners:
+                return None
+
+            tasks = data.get("tasks", [])
+            task_owners = []
+
+            # Handle default_args.owner (wrapped or not)
+            default_args = data.get("default_args", {})
+            if isinstance(default_args, dict) and "__var" in default_args:
+                default_args = default_args["__var"]
+            default_owner = default_args.get("owner")
+
+            for task in tasks:
+                # Flatten serialized task
+                task_data = (
+                    task.get("__var")
+                    if isinstance(task, dict) and "__var" in task
+                    else task
+                )
+
+                owner = task_data.get("owner") or default_owner
+
+                if owner:
+                    task_owners.append(owner)
+
+            if task_owners:
+                most_common_owner, _ = Counter(task_owners).most_common(1)[0]
+                return most_common_owner
+
+            return default_owner
+
         except Exception as exc:
             self.status.warning(
                 data.get("dag_id"), f"Could not extract owner information due to {exc}"

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

@@ -9,7 +9,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """
-Airbyte source to extract metadata
+Fivetran source to extract metadata
 """
 
 import traceback
@@ -166,7 +166,7 @@ class FivetranSource(PipelineServiceSource):
     ) -> Iterable[Either[AddLineageRequest]]:
         """
         Parse all the stream available in the connection and create a lineage between them
-        :param pipeline_details: pipeline_details object from airbyte
+        :param pipeline_details: pipeline_details object from fivetran
         :return: Lineage from inlets and outlets
         """
         self.client = cast(FivetranClient, self.client)

metadata/profiler/processor/processor.py

@@ -12,7 +12,7 @@
 Profiler Processor Step
 """
 import traceback
-from typing import Optional, cast
+from typing import Optional, Type, cast
 
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -31,6 +31,11 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.api.models import ProfilerProcessorConfig, ProfilerResponse
 from metadata.profiler.processor.core import Profiler
 from metadata.profiler.source.model import ProfilerSourceAndEntity
+from metadata.utils.dependency_injector.dependency_injector import (
+    DependencyNotFoundError,
+    Inject,
+    inject,
+)
 
 
 class ProfilerProcessor(Processor):
@@ -39,11 +44,21 @@ class ProfilerProcessor(Processor):
     the OpenMetadataSource and compute the metrics.
     """
 
-    def __init__(self, config: OpenMetadataWorkflowConfig):
+    @inject
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        profiler_config_class: Inject[Type[ProfilerProcessorConfig]] = None,
+    ):
+        if profiler_config_class is None:
+            raise DependencyNotFoundError(
+                "ProfilerProcessorConfig class not found. Please ensure the ProfilerProcessorConfig is properly registered."
+            )
+
         super().__init__()
 
         self.config = config
-        self.profiler_config = ProfilerProcessorConfig.model_validate(
+        self.profiler_config = profiler_config_class.model_validate(
             self.config.processor.model_dump().get("config")
         )
         self.source_config: DatabaseServiceProfilerPipeline = cast(

metadata/profiler/source/database/base/profiler_resolver.py

@@ -2,6 +2,9 @@ from abc import ABC, abstractmethod
 from typing import Tuple, Type
 
 from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    ProcessingEngine,
+)
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.service_spec.service_spec import (
@@ -16,7 +19,7 @@ class ProfilerResolver(ABC):
     @staticmethod
     @abstractmethod
     def resolve(
-        processing_engine: str, service_type: ServiceType, source_type: str
+        processing_engine: ProcessingEngine, service_type: ServiceType, source_type: str
     ) -> Tuple[Type[SamplerInterface], Type[ProfilerInterface]]:
         """Resolve the sampler and profiler based on the processing engine."""
         raise NotImplementedError
@@ -27,7 +30,7 @@ class DefaultProfilerResolver(ProfilerResolver):
 
     @staticmethod
     def resolve(
-        processing_engine: str, service_type: ServiceType, source_type: str
+        processing_engine: ProcessingEngine, service_type: ServiceType, source_type: str
     ) -> Tuple[Type[SamplerInterface], Type[ProfilerInterface]]:
         """Resolve the sampler and profiler based on the processing engine."""
         sampler_class = import_sampler_class(service_type, source_type=source_type)

metadata/profiler/source/database/base/profiler_source.py

@@ -61,16 +61,23 @@ class ProfilerSource(ProfilerSourceInterface):
     Base class for the profiler source
     """
 
+    @inject
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
         database: Database,
         ometa_client: OpenMetadata,
         global_profiler_configuration: ProfilerConfiguration,
+        profiler_config_class: Inject[Type[ProfilerProcessorConfig]] = None,
     ):
+        if profiler_config_class is None:
+            raise DependencyNotFoundError(
+                "ProfilerProcessorConfig class not found. Please ensure the ProfilerProcessorConfig is properly registered."
+            )
+
         self.config = config
         self.service_conn_config = self._copy_service_config(config, database)
-        self.profiler_config = ProfilerProcessorConfig.model_validate(
+        self.profiler_config = profiler_config_class.model_validate(
             config.processor.model_dump().get("config")
         )
         self.ometa_client = ometa_client

metadata/profiler/source/profiler_source_interface.py

@@ -18,6 +18,11 @@ from typing import Optional
 
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
+    ProcessingEngine,
+)
+from metadata.generated.schema.metadataIngestion.engine.nativeEngineConfig import (
+    NativeEngineConfiguration,
+    Type,
 )
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 
@@ -43,6 +48,10 @@ class ProfilerSourceInterface(ABC):
         raise NotImplementedError
 
     @staticmethod
-    def get_processing_engine(config: DatabaseServiceProfilerPipeline):
+    def get_processing_engine(
+        config: DatabaseServiceProfilerPipeline,
+    ) -> ProcessingEngine:
         """Get the processing engine based on the configuration."""
-        return "Native"
+        return config.processingEngine or ProcessingEngine(
+            root=NativeEngineConfiguration(type=Type.Native)
+        )

metadata/sampler/pandas/sampler.py

@@ -60,7 +60,7 @@ class DatalakeSampler(SamplerInterface, PandasInterfaceMixin):
         return self._table
 
     def get_client(self):
-        return self.connection.client
+        return self.connection
 
     def _partitioned_table(self):
         """Get partitioned table"""

metadata/sampler/processor.py

@@ -13,7 +13,7 @@ Data Sampler for the PII Workflow
 """
 import traceback
 from copy import deepcopy
-from typing import Optional, cast
+from typing import Optional, Type, cast
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.table import Table
@@ -42,6 +42,11 @@ from metadata.sampler.config import get_config_for_table
 from metadata.sampler.models import SampleConfig, SampleData, SamplerResponse
 from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.bigquery_utils import copy_service_config
+from metadata.utils.dependency_injector.dependency_injector import (
+    DependencyNotFoundError,
+    Inject,
+    inject,
+)
 from metadata.utils.profiler_utils import get_context_entities
 from metadata.utils.service_spec.service_spec import import_sampler_class
 
@@ -49,7 +54,18 @@ from metadata.utils.service_spec.service_spec import import_sampler_class
 class SamplerProcessor(Processor):
     """Use the profiler interface to fetch the sample data"""
 
-    def __init__(self, config: OpenMetadataWorkflowConfig, metadata: OpenMetadata):
+    @inject
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        profiler_config_class: Inject[Type[ProfilerProcessorConfig]] = None,
+    ):
+        if profiler_config_class is None:
+            raise DependencyNotFoundError(
+                "ProfilerProcessorConfig class not found. Please ensure the ProfilerProcessorConfig is properly registered."
+            )
+
         super().__init__()
 
         self.config = config
@@ -60,7 +76,7 @@ class SamplerProcessor(Processor):
             self.config.source.sourceConfig.config,
         )  # Used to satisfy type checked
         # We still rely on the orm-processor. We should decouple this in the future
-        self.profiler_config = ProfilerProcessorConfig.model_validate(
+        self.profiler_config = profiler_config_class.model_validate(
             self.config.processor.model_dump().get("config")
         )
 

metadata/sampler/sampler_interface.py

@@ -30,6 +30,9 @@ from metadata.generated.schema.entity.services.connections.database.datalakeConn
     DatalakeConnection,
 )
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    ProcessingEngine,
+)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.api.models import TableConfig
 from metadata.profiler.processor.sample_data_handler import upload_sample_data
@@ -70,6 +73,7 @@ class SamplerInterface(ABC):
         sample_query: Optional[str] = None,
         storage_config: Optional[DataStorageConfig] = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        processing_engine: Optional[ProcessingEngine] = None,
         **__,
     ):
         self.ometa_client = ometa_client
@@ -84,6 +88,7 @@ class SamplerInterface(ABC):
         self.sample_limit = sample_data_count
         self.partition_details = partition_details
         self.storage_config = storage_config
+        self.processing_engine = processing_engine
 
         self.service_connection_config = service_connection_config
         self.connection = get_ssl_connection(self.service_connection_config)
@@ -101,6 +106,7 @@ class SamplerInterface(ABC):
         storage_config: Optional[DataStorageConfig] = None,
         default_sample_config: Optional[SampleConfig] = None,
         default_sample_data_count: int = SAMPLE_DATA_DEFAULT_COUNT,
+        processing_engine: Optional[ProcessingEngine] = None,
         **kwargs,
     ) -> "SamplerInterface":
         """Create sampler"""
@@ -137,6 +143,7 @@ class SamplerInterface(ABC):
             sample_query=sample_query,
             storage_config=storage_config,
             sample_data_count=sample_data_count,
+            processing_engine=processing_engine,
             **kwargs,
         )
 

metadata/utils/credentials.py

@@ -21,6 +21,10 @@ from cryptography.hazmat.primitives import serialization
 from google import auth
 from google.auth import impersonated_credentials
 
+from metadata.clients.azure_client import AzureClient
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
 from metadata.generated.schema.security.credentials.gcpCredentials import (
     GcpADC,
     GCPCredentials,
@@ -237,3 +241,34 @@ def get_gcp_impersonate_credentials(
         target_scopes=scopes,
         lifetime=lifetime,
     )
+
+
+def get_azure_access_token(azure_config: AzureConfigurationSource) -> str:
+    """
+    Get Azure access token using the provided Azure configuration.
+
+    Args:
+        azure_config: Azure configuration containing the necessary credentials and scopes
+
+    Returns:
+        str: The access token
+
+    Raises:
+        ValueError: If Azure config is missing or scopes are not provided
+    """
+    if not azure_config.azureConfig:
+        raise ValueError("Azure Config is missing")
+
+    if not azure_config.azureConfig.scopes:
+        raise ValueError(
+            "Azure Scopes are missing, please refer https://learn.microsoft.com/"
+            "en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token "
+            "and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
+        )
+
+    azure_client = AzureClient(azure_config.azureConfig).create_client()
+    access_token_obj = azure_client.get_token(
+        *azure_config.azureConfig.scopes.split(",")
+    )
+
+    return access_token_obj.token

metadata/utils/db_utils.py

@@ -60,6 +60,7 @@ def get_view_lineage(
     table_name = view.table_name
     schema_name = view.schema_name
     db_name = view.db_name
+    schema_fallback = False
     view_definition = view.view_definition
     table_fqn = fqn.build(
         metadata,
@@ -88,6 +89,7 @@ def get_view_lineage(
         if table_entity.serviceType == DatabaseServiceType.Postgres:
             # For Postgres, if schema is not defined, we need to use the public schema
             schema_name = PUBLIC_SCHEMA
+            schema_fallback = True
 
         if lineage_parser.source_tables and lineage_parser.target_tables:
             yield from get_lineage_by_query(
@@ -99,6 +101,7 @@ def get_view_lineage(
                 dialect=dialect,
                 timeout_seconds=timeout_seconds,
                 lineage_source=LineageSource.ViewLineage,
+                schema_fallback=schema_fallback,
             ) or []
 
         else:
@@ -112,6 +115,7 @@ def get_view_lineage(
                 dialect=dialect,
                 timeout_seconds=timeout_seconds,
                 lineage_source=LineageSource.ViewLineage,
+                schema_fallback=schema_fallback,
             ) or []
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/utils/filters.py

@@ -338,3 +338,63 @@ def filter_by_tag(tag_pattern: Optional[FilterPattern], tag_name: str) -> bool:
     :return: True for filtering, False otherwise
     """
     return _filter(tag_pattern, tag_name)
+
+
+def filter_by_spreadsheet(
+    spreadsheet_filter_pattern: Optional[FilterPattern], spreadsheet_name: str
+) -> bool:
+    """
+    Return True if the spreadsheet needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param spreadsheet_filter_pattern: Model defining spreadsheet filtering logic
+    :param spreadsheet_name: spreadsheet name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(spreadsheet_filter_pattern, spreadsheet_name)
+
+
+def filter_by_directory(
+    directory_filter_pattern: Optional[FilterPattern], directory_name: str
+) -> bool:
+    """
+    Return True if the directory needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param directory_filter_pattern: Model defining directory filtering logic
+    :param directory_name: directory name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(directory_filter_pattern, directory_name)
+
+
+def filter_by_file(
+    file_filter_pattern: Optional[FilterPattern], file_name: str
+) -> bool:
+    """
+    Return True if the file needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param file_filter_pattern: Model defining file filtering logic
+    :param file_name: file name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(file_filter_pattern, file_name)
+
+
+def filter_by_worksheet(
+    worksheet_filter_pattern: Optional[FilterPattern], worksheet_name: str
+) -> bool:
+    """
+    Return True if the worksheet needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param worksheet_filter_pattern: Model defining worksheet filtering logic
+    :param worksheet_name: worksheet name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(worksheet_filter_pattern, worksheet_name)

metadata/utils/fqn.py

@@ -34,13 +34,18 @@ from metadata.generated.schema.entity.data.dashboard import Dashboard
 from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.directory import Directory
+from metadata.generated.schema.entity.data.file import File
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
 from metadata.generated.schema.entity.data.searchIndex import SearchIndex
+from metadata.generated.schema.entity.data.spreadsheet import Spreadsheet
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.generated.schema.entity.data.table import Column, DataModel, Table
 from metadata.generated.schema.entity.data.topic import Topic
+from metadata.generated.schema.entity.data.worksheet import Worksheet
+from metadata.generated.schema.entity.services.driveService import DriveService
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.tests.testCase import TestCase
@@ -545,6 +550,80 @@ def _(
     return _build(service_name, query_checksum)
 
 
+@fqn_build_registry.add(DriveService)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for dashboard FQN building
+    *,
+    service_name: str,
+) -> str:
+    return _build(service_name)
+
+
+@fqn_build_registry.add(Directory)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for directory FQN building
+    *,
+    service_name: str,
+    directory_path: List[str],
+) -> str:
+    if not service_name:
+        raise FQNBuildingException(
+            f"Service name should be informed, but got service=`{service_name}`"
+        )
+
+    if not directory_path:
+        raise FQNBuildingException("Directory path should not be empty")
+
+    return _build(service_name, *directory_path)
+
+
+@fqn_build_registry.add(File)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for file FQN building
+    *,
+    service_name: str,
+    directory_path: List[str],
+    file_name: str,
+) -> str:
+    if not service_name or not file_name:
+        raise FQNBuildingException(
+            f"Args should be informed, but got service=`{service_name}`, file=`{file_name}`"
+        )
+    if not directory_path:
+        raise FQNBuildingException("Directory path should not be empty")
+    return _build(service_name, *directory_path, file_name)
+
+
+@fqn_build_registry.add(Worksheet)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for dashboard FQN building
+    *,
+    service_name: str,
+    spreadsheet_name: str,
+    worksheet_name: str,
+) -> str:
+    if not service_name or not spreadsheet_name or not worksheet_name:
+        raise FQNBuildingException(
+            f"Args should be informed, but got service=`{service_name}`, "
+            f"spreadsheet=`{spreadsheet_name}`, worksheet=`{worksheet_name}``"
+        )
+    return _build(service_name, spreadsheet_name, worksheet_name)
+
+
+@fqn_build_registry.add(Spreadsheet)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for dashboard FQN building
+    *,
+    service_name: str,
+    spreadsheet_name: str,
+) -> str:
+    if not service_name or not spreadsheet_name:
+        raise FQNBuildingException(
+            f"Args should be informed, but got service=`{service_name}`, spreadsheet=`{spreadsheet_name}``"
+        )
+    return _build(service_name, spreadsheet_name)
+
+
 def split_table_name(table_name: str) -> Dict[str, Optional[str]]:
     """
     Given a table name, try to extract database, schema and

metadata/utils/logger.py

@@ -25,6 +25,9 @@ from metadata.data_quality.api.models import (
     TestCaseResults,
 )
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.datacontract.dataContractResult import (
+    DataContractResult,
+)
 from metadata.generated.schema.type.queryParserData import QueryParserData
 from metadata.generated.schema.type.tableQuery import TableQueries
 from metadata.ingestion.api.models import Entity
@@ -301,6 +304,12 @@ def _(record: QueryParserData) -> str:
     return f"Usage ParsedData [{len(record.parsedData)}]"
 
 
+@get_log_name.register
+def _(record: DataContractResult) -> str:
+    """Get the log of the DataContractResult"""
+    return f"DataContractResult for [{record.dataContractFQN.root}]; status: {record.contractExecutionStatus.value}]"
+
+
 def redacted_config(config: Dict[str, Union[str, dict]]) -> Dict[str, Union[str, dict]]:
     config_copy = deepcopy(config)
 

metadata/utils/metadata_service_helper.py

@@ -83,7 +83,7 @@ SERVICE_TYPE_MAPPER = {
         "connection": {"config": {"awsConfig": "aws_config"}},
     },
     "snowflake": {
-        "service_nmae": "Snowflake",
+        "service_name": "Snowflake",
         "connection": {
             "config": {
                 "username": "randomName",