openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/profiler/source/fetcher/profiler_source_factory.py

@@ -0,0 +1,89 @@
+#  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.
+#  pylint: disable=import-outside-toplevel
+"""
+Factory class for creating profiler source objects
+"""
+
+from typing import Callable, Dict, Type
+
+from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
+    BigqueryType,
+)
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksType,
+)
+from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
+
+
+class ProfilerSourceFactory:
+    """Creational factory for profiler source objects"""
+
+    def __init__(self):
+        self._source_type: Dict[str, Callable[[], Type[ProfilerSourceInterface]]] = {
+            "base": self.base
+        }
+
+    def register_source(self, type_: str, source_fn):
+        """Register a new source type"""
+        self._source_type[type_] = source_fn
+
+    def register_many_sources(
+        self, source_dict: Dict[str, Callable[[], Type[ProfilerSourceInterface]]]
+    ):
+        """Register multiple source types at once"""
+        for type_, source_fn in source_dict.items():
+            self.register_source(type_, source_fn)
+
+    def create(self, type_: str, *args, **kwargs) -> ProfilerSourceInterface:
+        """Create source object based on source type"""
+        source_fn = self._source_type.get(type_)
+        if not source_fn:
+            source_fn = self._source_type["base"]
+
+        source_class = source_fn()
+        return source_class(*args, **kwargs)
+
+    @staticmethod
+    def base() -> Type[ProfilerSourceInterface]:
+        """Lazy loading of the base source"""
+        from metadata.profiler.source.database.base.profiler_source import (
+            ProfilerSource,
+        )
+
+        return ProfilerSource
+
+    @staticmethod
+    def bigquery() -> Type[ProfilerSourceInterface]:
+        """Lazy loading of the BigQuery source"""
+        from metadata.profiler.source.database.bigquery.profiler_source import (
+            BigQueryProfilerSource,
+        )
+
+        return BigQueryProfilerSource
+
+    @staticmethod
+    def databricks() -> Type[ProfilerSourceInterface]:
+        """Lazy loading of the Databricks source"""
+        from metadata.profiler.source.database.databricks.profiler_source import (
+            DataBricksProfilerSource,
+        )
+
+        return DataBricksProfilerSource
+
+
+source = {
+    BigqueryType.BigQuery.value.lower(): ProfilerSourceFactory.bigquery,
+    DatabricksType.Databricks.value.lower(): ProfilerSourceFactory.databricks,
+}
+
+profiler_source_factory = ProfilerSourceFactory()
+profiler_source_factory.register_many_sources(source)

metadata/profiler/source/metadata.py

@@ -11,18 +11,9 @@
 """
 OpenMetadata source for the profiler
 """
-import traceback
 from typing import Iterable, List, Optional, cast
 
-from pydantic import BaseModel, ConfigDict
-
-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 Table, TableType
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
-from metadata.generated.schema.entity.services.ingestionPipelines.status import (
-    StackTraceError,
-)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
@@ -34,32 +25,16 @@ from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.source.base.profiler_source import ProfilerSource
-from metadata.profiler.source.profiler_source_factory import profiler_source_factory
-from metadata.utils import fqn
-from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
+from metadata.profiler.source.fetcher.entity_fetcher import EntityFetcher
+from metadata.profiler.source.model import ProfilerSourceAndEntity
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
 
-
 TABLE_FIELDS = ["tableProfilerConfig", "columns", "customMetrics"]
 TAGS_FIELD = ["tags"]
 
 
-class ProfilerSourceAndEntity(BaseModel):
-    """Return class for the OpenMetadata Profiler Source"""
-
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    profiler_source: ProfilerSource
-    entity: Table
-
-    def __str__(self):
-        """Return the information of the table being profiler"""
-        return f"Table [{self.entity.name.root}]"
-
-
 class OpenMetadataSource(Source):
     """
     This source lists and filters the entities that need
@@ -70,20 +45,16 @@ class OpenMetadataSource(Source):
     We do this here as well.
     """
 
-    def init_steps(self):
-        super().__init__()
-
     @property
     def name(self) -> str:
         return "OpenMetadata Service"
 
-    # pylint: disable=super-init-not-called
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
     ):
-        self.init_steps()
+        super().__init__()
 
         self.config = config
         self.metadata = metadata
@@ -107,10 +78,18 @@ class OpenMetadataSource(Source):
             f":{self.config.source.type.lower()}"
         )
 
+    def _get_fields(self) -> List[str]:
+        """Get the fields required to process the tables"""
+        return (
+            TABLE_FIELDS
+            if not self.source_config.processPiiSensitive
+            else TABLE_FIELDS + TAGS_FIELD
+        )
+
     def _validate_service_name(self):
         """Validate service name exists in OpenMetadata"""
         return self.metadata.get_by_name(
-            entity=DatabaseService, fqn=self.config.source.serviceName
+            entity=DatabaseService, fqn=self.config.source.serviceName  # type: ignore
         )
 
     def prepare(self):
@@ -125,30 +104,10 @@ class OpenMetadataSource(Source):
 
     def _iter(self, *_, **__) -> Iterable[Either[ProfilerSourceAndEntity]]:
         global_profiler_config = self.metadata.get_profiler_config_settings()
-        for database in self.get_database_entities():
-            try:
-                profiler_source = profiler_source_factory.create(
-                    self.config.source.type.lower(),
-                    self.config,
-                    database,
-                    self.metadata,
-                    global_profiler_config,
-                )
-                for entity in self.get_table_entities(database=database):
-                    yield Either(
-                        right=ProfilerSourceAndEntity(
-                            profiler_source=profiler_source,
-                            entity=entity,
-                        )
-                    )
-            except Exception as exc:
-                yield Either(
-                    left=StackTraceError(
-                        name=database.fullyQualifiedName.root,
-                        error=f"Error listing source and entities for database due to [{exc}]",
-                        stackTrace=traceback.format_exc(),
-                    )
-                )
+        entity_fetcher = EntityFetcher(
+            self.config, self.metadata, global_profiler_config, self.status
+        )
+        yield from entity_fetcher.fetch()
 
     @classmethod
     def create(
@@ -160,149 +119,5 @@ class OpenMetadataSource(Source):
         config = parse_workflow_config_gracefully(config_dict)
         return cls(config=config, metadata=metadata)
 
-    def filter_databases(self, database: Database) -> Optional[Database]:
-        """Returns filtered database entities"""
-        database_fqn = fqn.build(
-            self.metadata,
-            entity_type=Database,
-            service_name=self.config.source.serviceName,
-            database_name=database.name.root,
-        )
-        if filter_by_database(
-            self.source_config.databaseFilterPattern,
-            database_fqn
-            if self.source_config.useFqnForFiltering
-            else database.name.root,
-        ):
-            self.status.filter(database.name.root, "Database pattern not allowed")
-            return None
-        return database
-
-    def filter_entities(self, tables: Iterable[Table]) -> Iterable[Table]:
-        """
-        From a list of tables, apply the SQLSourceConfig
-        filter patterns.
-
-        We will update the status on the SQLSource Status.
-        """
-        for table in tables:
-            try:
-                schema_fqn = fqn.build(
-                    self.metadata,
-                    entity_type=DatabaseSchema,
-                    service_name=self.config.source.serviceName,
-                    database_name=table.database.name,
-                    schema_name=table.databaseSchema.name,
-                )
-                if filter_by_schema(
-                    self.source_config.schemaFilterPattern,
-                    schema_fqn
-                    if self.source_config.useFqnForFiltering
-                    else table.databaseSchema.name,  # type: ignore
-                ):
-                    self.status.filter(
-                        f"Schema pattern not allowed: {table.fullyQualifiedName.root}",
-                        "Schema pattern not allowed",
-                    )
-                    continue
-                table_fqn = fqn.build(
-                    self.metadata,
-                    entity_type=Table,
-                    service_name=self.config.source.serviceName,
-                    database_name=table.database.name,
-                    schema_name=table.databaseSchema.name,
-                    table_name=table.name.root,
-                )
-
-                if filter_by_table(
-                    self.source_config.tableFilterPattern,
-                    table_fqn
-                    if self.source_config.useFqnForFiltering
-                    else table.name.root,
-                ):
-                    self.status.filter(
-                        f"Table pattern not allowed: {table.fullyQualifiedName.root}",
-                        "Table pattern not allowed",
-                    )
-                    continue
-                if (
-                    table.tableType == TableType.View
-                    and not self.source_config.includeViews
-                ):
-                    self.status.filter(
-                        table.fullyQualifiedName.root,
-                        "View filtered out",
-                    )
-                    continue
-                yield table
-            except Exception as exc:
-                self.status.failed(
-                    StackTraceError(
-                        name=table.fullyQualifiedName.root,
-                        error=f"Unexpected error filtering entities for table [{table}]: {exc}",
-                        stackTrace=traceback.format_exc(),
-                    )
-                )
-
-    def get_database_entities(self):
-        """List all databases in service"""
-
-        databases = [
-            self.filter_databases(database)
-            for database in self.metadata.list_all_entities(
-                entity=Database,
-                params={"service": self.config.source.serviceName},
-            )
-            if self.filter_databases(database)
-        ]
-
-        if not databases:
-            raise ValueError(
-                "databaseFilterPattern returned 0 result. At least 1 database must be returned by the filter pattern."
-                f"\n\t- includes: {self.source_config.databaseFilterPattern.includes if self.source_config.databaseFilterPattern else None}"  # pylint: disable=line-too-long
-                f"\n\t- excludes: {self.source_config.databaseFilterPattern.excludes if self.source_config.databaseFilterPattern else None}"  # pylint: disable=line-too-long
-            )
-
-        return databases
-
-    def _get_fields(self) -> List[str]:
-        """Get the fields required to process the tables"""
-        return (
-            TABLE_FIELDS
-            if not self.source_config.processPiiSensitive
-            else TABLE_FIELDS + TAGS_FIELD
-        )
-
-    def get_table_entities(self, database):
-        """
-        List and filter OpenMetadata tables based on the
-        source configuration.
-
-        The listing will be based on the entities from the
-        informed service name in the source configuration.
-
-        Note that users can specify `table_filter_pattern` to
-        either be `includes` or `excludes`. This means
-        that we will either what is specified in `includes`
-        or we will use everything but the tables excluded.
-
-        Same with `schema_filter_pattern`.
-        """
-        tables = self.metadata.list_all_entities(
-            entity=Table,
-            fields=self._get_fields(),
-            params={
-                "service": self.config.source.serviceName,
-                "database": fqn.build(
-                    self.metadata,
-                    entity_type=Database,
-                    service_name=self.config.source.serviceName,
-                    database_name=database.name.root,
-                ),
-            },  # type: ignore
-        )
-
-        yield from self.filter_entities(tables)
-
     def close(self) -> None:
         """Nothing to close"""

metadata/profiler/source/metadata_ext.py

@@ -21,7 +21,7 @@ and sample data for that identified entity.
 """
 import traceback
 from copy import deepcopy
-from typing import Iterable, cast
+from typing import Iterable, Type, cast
 
 from sqlalchemy.inspection import inspect
 
@@ -29,6 +29,7 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
+from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
@@ -40,16 +41,16 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.source.metadata import (
-    OpenMetadataSource,
-    ProfilerSourceAndEntity,
-)
-from metadata.profiler.source.profiler_source_factory import profiler_source_factory
+from metadata.profiler.interface.profiler_interface import ProfilerInterface
+from metadata.profiler.source.metadata import OpenMetadataSource
+from metadata.profiler.source.model import ProfilerSourceAndEntity
 from metadata.utils import fqn
 from metadata.utils.class_helper import get_service_type_from_source_type
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
-from metadata.utils.importer import import_source_class
+from metadata.utils.importer import import_from_module
 from metadata.utils.logger import profiler_logger
+from metadata.utils.service_spec import BaseSpec
+from metadata.utils.service_spec.service_spec import import_source_class
 from metadata.utils.ssl_manager import get_ssl_connection
 
 logger = profiler_logger()
@@ -65,17 +66,12 @@ class OpenMetadataSourceExt(OpenMetadataSource):
     We do this here as well.
     """
 
-    # pylint: disable=super-init-not-called
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
     ):
-        self.init_steps()
-
-        self.config = config
-        self.metadata = metadata
-        self.test_connection()
+        super().__init__(config, metadata)
 
         # Init and type the source config
         self.service_connection = self.config.source.serviceConnection.root.config
@@ -145,8 +141,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                             )
                             continue
 
-                        profiler_source = profiler_source_factory.create(
-                            self.config.source.type.lower(),
+                        profiler_source = self.import_profiler_interface()(
                             self.config,
                             database_entity,
                             self.metadata,
@@ -174,6 +169,14 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                 continue
             yield table_name
 
+    def import_profiler_interface(self) -> Type[ProfilerInterface]:
+        class_path = BaseSpec.get_for_source(
+            ServiceType.Database,
+            source_type=self.config.source.type.lower(),
+        ).profiler_class
+        profiler_source_class = import_from_module(class_path)
+        return cast(Type[ProfilerInterface], profiler_source_class)
+
     def get_schema_names(self) -> Iterable[str]:
         if self.service_connection.__dict__.get("databaseSchema"):
             yield self.service_connection.databaseSchema
@@ -206,9 +209,11 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                         )
                         if filter_by_database(
                             self.source_config.databaseFilterPattern,
-                            database_fqn
-                            if self.source_config.useFqnForFiltering
-                            else database,
+                            (
+                                database_fqn
+                                if self.source_config.useFqnForFiltering
+                                else database
+                            ),
                         ):
                             self.status.filter(database, "Database pattern not allowed")
                             continue

metadata/profiler/source/model.py

@@ -0,0 +1,29 @@
+#  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.
+"""Model for the OpenMetadata Profiler Source"""
+from pydantic import ConfigDict
+
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.models.entity_interface import EntityInterface
+from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
+
+
+class ProfilerSourceAndEntity(BaseModel):
+    """Return class for the OpenMetadata Profiler Source"""
+
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    profiler_source: ProfilerSourceInterface
+    entity: EntityInterface
+
+    def __str__(self):
+        """Return the information of the table being profiler"""
+        return f"Table [{self.entity.name.root}]"

metadata/utils/class_helper.py

@@ -22,6 +22,9 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
     PipelineType,
 )
 from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.apiServiceMetadataPipeline import (
+    ApiServiceMetadataPipeline,
+)
 from metadata.generated.schema.metadataIngestion.dashboardServiceMetadataPipeline import (
     DashboardServiceMetadataPipeline,
 )
@@ -65,6 +68,7 @@ from metadata.generated.schema.metadataIngestion.testSuitePipeline import (
 from metadata.generated.schema.metadataIngestion.workflow import SourceConfig
 
 SERVICE_TYPE_REF = {
+    ServiceType.Api.value: "apiService",
     ServiceType.Database.value: "databaseService",
     ServiceType.Dashboard.value: "dashboardService",
     ServiceType.Pipeline.value: "pipelineService",
@@ -76,6 +80,7 @@ SERVICE_TYPE_REF = {
 }
 
 SOURCE_CONFIG_TYPE_INGESTION = {
+    ApiServiceMetadataPipeline.__name__: PipelineType.metadata,
     DatabaseServiceMetadataPipeline.__name__: PipelineType.metadata,
     DatabaseServiceQueryUsagePipeline.__name__: PipelineType.usage,
     DatabaseServiceQueryLineagePipeline.__name__: PipelineType.lineage,

metadata/utils/collaborative_super.py

@@ -0,0 +1,28 @@
+"""
+This module provides utilties for building a collaborative constructor library.
+"""
+
+from abc import ABC
+
+
+class Root:
+    """Root class for any class that needs to implement a colllaborative constructor but
+    might end up at the end of the inheritance chain. Since python's object.__init__ is not
+    a collaborative constructor, we need to have a root class that has a collaborative constructor.
+    """
+
+    __terminal__ = {object, ABC}
+
+    def __init__(self, *args, **kwargs):
+        """Collaborative constructor"""
+        super_class = None
+        for cls, super_class in zip(
+            self.__class__.mro()[:-1], self.__class__.mro()[1:]
+        ):
+            if cls is Root:
+                break
+        for cls in self.__terminal__:
+            if super_class is cls:
+                break
+        else:
+            super().__init__(*args, **kwargs)

metadata/utils/constants.py

@@ -26,6 +26,7 @@ from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
 from metadata.generated.schema.entity.domains.dataProduct import DataProduct
 from metadata.generated.schema.entity.domains.domain import Domain
+from metadata.generated.schema.entity.services.apiService import ApiService
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.messagingService import MessagingService
@@ -77,6 +78,7 @@ ENTITY_REFERENCE_CLASS_MAP = {
     # Service Entities
     "databaseService": DatabaseService,
     "messagingService": MessagingService,
+    "apiService": ApiService,
     "dashboardService": DashboardService,
     "pipelineService": PipelineService,
     "storageService": StorageService,

metadata/utils/constraints.py

@@ -0,0 +1,65 @@
+#  Copyright 2024 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.
+
+"""
+Define constraints helper methods useful for the metadata ingestion
+"""
+
+from typing import Dict, List
+
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    ConstraintType,
+    RelationshipType,
+)
+from metadata.ingestion.ometa.utils import model_str
+
+
+def _is_column_unique(column: Dict, columns: List[Column]) -> bool:
+    """
+    Method to check if the column in unique in the table
+    """
+    if column and len(column) > 0:
+        constrained_column = column[0]
+        for col in columns or []:
+            if model_str(col.name) == constrained_column:
+                if col.constraint and col.constraint.value in {
+                    ConstraintType.UNIQUE.value,
+                    ConstraintType.PRIMARY_KEY.value,
+                }:
+                    return True
+                break
+    return False
+
+
+def get_relationship_type(
+    column: Dict, referred_table_columns: List[Column], columns: List[Column]
+) -> str:
+    """
+    Determine the type of relationship (one-to-one, one-to-many, etc.)
+    """
+    # Check if the column is unique in the current table
+    is_unique_in_current_table = _is_column_unique(
+        column.get("constrained_columns"), columns
+    )
+
+    # Check if the referred column is unique in the referred table
+    is_unique_in_referred_table = _is_column_unique(
+        column.get("referred_columns"), referred_table_columns
+    )
+
+    if is_unique_in_current_table and is_unique_in_referred_table:
+        return RelationshipType.ONE_TO_ONE
+    if is_unique_in_current_table:
+        return RelationshipType.ONE_TO_MANY
+    if is_unique_in_referred_table:
+        return RelationshipType.MANY_TO_ONE
+    return RelationshipType.MANY_TO_MANY

metadata/utils/dict.py

@@ -0,0 +1,8 @@
+"""
+A custom dictionary class that extends functionality.
+"""
+
+
+class ExtendedDict(dict):
+    def lower_case_keys(self):
+        return {k.lower(): v for k, v in self.items()}