openmetadata-ingestion 1.5.13.1__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -0,0 +1,299 @@
+#  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.
+"""
+Entity Fetcher Strategy
+"""
+
+import traceback
+from abc import ABC, abstractmethod
+from typing import Iterable, Iterator, Optional, cast
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.generated.schema.settings.settings import Settings
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.status import Status
+from metadata.ingestion.models.entity_interface import EntityInterfaceWithTags
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.fetcher.config import EntityFilterConfigInterface
+from metadata.profiler.source.fetcher.profiler_source_factory import (
+    profiler_source_factory,
+)
+from metadata.profiler.source.model import ProfilerSourceAndEntity
+from metadata.utils.db_utils import Table
+from metadata.utils.filters import (
+    filter_by_classification,
+    filter_by_database,
+    filter_by_schema,
+    filter_by_table,
+)
+
+FIELDS = ["tableProfilerConfig", "columns", "customMetrics", "tags"]
+
+
+class FetcherStrategy(ABC):
+    """Fetcher strategy interface"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        global_profiler_config: Optional[Settings],
+        status: Status,
+    ) -> None:
+        self.config = config
+        self.source_config = config.source.sourceConfig.config
+        self.metadata = metadata
+        self.global_profiler_config = global_profiler_config
+        self.status = status
+
+    def filter_classifications(self, entity: EntityInterfaceWithTags) -> bool:
+        """Given a list of entities, filter out entities that do not match the classification filter pattern
+
+        Args:
+            entities (List[EntityInterfaceWithTags]): List of entities to filter implemnenting the `tags` attribute
+
+        Raises:
+            NotImplementedError: Must be implemented by subclass
+        """
+        classification_filter_pattern = getattr(
+            self.source_config, "classificationFilterPattern", None
+        )
+        if not classification_filter_pattern:
+            return False
+
+        use_fqn_for_filtering = getattr(self.source_config, "useFqnForFiltering", False)
+
+        if not entity.tags:
+            # if we are not explicitly including entities with tags we'll add the ones without tags
+            if not classification_filter_pattern.includes:
+                return False
+            return True
+
+        for tag in entity.tags:
+            tag_name = tag.tagFQN.root if use_fqn_for_filtering else tag.name
+            if not tag_name:
+                continue
+            if filter_by_classification(classification_filter_pattern, tag_name):
+                self.status.filter(
+                    tag_name,
+                    f"Classification pattern not allowed for entity {entity.fullyQualifiedName.root}",
+                )  # type: ignore
+                return True
+
+        return False
+
+    @abstractmethod
+    def fetch(self) -> Iterator[Either[ProfilerSourceAndEntity]]:
+        """Fetch entity"""
+        raise NotImplementedError
+
+
+class DatabaseFetcherStrategy(FetcherStrategy):
+    """Database fetcher strategy"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        global_profiler_config: Optional[Settings],
+        status: Status,
+    ) -> None:
+        super().__init__(config, metadata, global_profiler_config, status)
+        self.source_config = cast(
+            EntityFilterConfigInterface, self.source_config
+        )  # Satisfy typchecker
+
+    def _filter_databases(self, databases: Iterable[Database]) -> Iterable[Database]:
+        """Filter databases based on the filter pattern
+
+        Args:
+            database (Database): Database to filter
+
+        Returns:
+            bool
+        """
+        filtered_databases = []
+
+        for database in databases:
+            database_name = database.name.root
+            if database.fullyQualifiedName and self.source_config.useFqnForFiltering:
+                database_name = database.fullyQualifiedName.root
+            if filter_by_database(
+                self.source_config.databaseFilterPattern, database_name
+            ):
+                self.status.filter(
+                    database_name,
+                    f"Database pattern not allowed for database {database_name}",
+                )
+                continue
+            filtered_databases.append(database)
+
+        return filtered_databases
+
+    def _filter_schemas(self, table: Table) -> bool:
+        """Filter tables based on the schema filter pattern
+
+        Args:
+            tables (List[Table]): Tables to filter
+
+        Returns:
+            List[Table]
+        """
+
+        if not table.databaseSchema:
+            return False
+        schema_name = (
+            table.databaseSchema.fullyQualifiedName
+            if self.source_config.useFqnForFiltering
+            else table.databaseSchema.name
+        )
+        if schema_name and filter_by_schema(
+            self.source_config.schemaFilterPattern, schema_name
+        ):
+            self.status.filter(
+                schema_name, f"Schema pattern not allowed for schema {schema_name}"
+            )
+            return True
+
+        return False
+
+    def _filter_tables(self, table: Table) -> bool:
+        """Filter tables based on the table filter pattern
+
+        Args:
+            tables (Iterable[Table]):
+
+        Returns:
+            Iterable[Table]:
+        """
+        table_name = table.name.root
+        if table.fullyQualifiedName and self.source_config.useFqnForFiltering:
+            table_name = table.fullyQualifiedName.root
+
+        if filter_by_table(self.source_config.tableFilterPattern, table_name):
+            self.status.filter(
+                table_name, f"Table pattern not allowed for table {table_name}"
+            )
+            return True
+
+        return False
+
+    def _filter_column_metrics_computation(self):
+        """Filter"""
+
+    def _get_database_entities(self) -> Iterable[Database]:
+        """Get database entities"""
+        if not self.config.source.serviceName:
+            raise ValueError("serviceName must be provided in the source configuration")
+
+        databases = self.metadata.list_all_entities(
+            entity=Database,
+            params={"service": self.config.source.serviceName},
+        )
+        databases = cast(Iterable[Database], databases)
+
+        if self.source_config.databaseFilterPattern:
+            databases = self._filter_databases(databases)
+
+        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 cast(Iterable[Database], databases)
+
+    def _filter_entities(self, tables: Iterable[Table]) -> Iterable[Table]:
+        """Filter entities based on the filter pattern
+
+        Args:
+            entities (Iterable[EntityInterfaceWithTags]): Entities to filter
+
+        Returns:
+            Iterable[EntityInterfaceWithTags]
+        """
+        tables = [
+            table
+            for table in tables
+            if (
+                not self.source_config.schemaFilterPattern
+                or not self._filter_schemas(table)
+            )
+            and (
+                not self.source_config.tableFilterPattern
+                or not self._filter_tables(table)
+            )
+            and (
+                not self.source_config.classificationFilterPattern
+                or not self.filter_classifications(table)
+            )
+        ]
+
+        return tables
+
+    def _get_table_entities(self, database: Database) -> Iterable[Table]:
+        """Given a database, get all table entities
+
+        Args:
+            database (Database): Database to get tables from
+
+        Returns:
+            Iterable[Table]
+        """
+        tables = self.metadata.list_all_entities(
+            entity=Table,
+            fields=FIELDS,
+            params={
+                "service": self.config.source.serviceName,
+                "database": database.fullyQualifiedName.root,  # type: ignore
+            },  # type: ignore
+        )
+        tables = cast(Iterable[Table], tables)
+        tables = self._filter_entities(tables)
+
+        return cast(Iterable[Table], tables)
+
+    def fetch(self) -> Iterator[Either[ProfilerSourceAndEntity]]:
+        """Fetch database entity"""
+        for database in self._get_database_entities():
+            try:
+                profiler_source = profiler_source_factory.create(
+                    self.config.source.type.lower(),
+                    self.config,
+                    database,
+                    self.metadata,
+                    self.global_profiler_config,
+                )
+
+                for table in self._get_table_entities(database):
+                    yield Either(
+                        left=None,
+                        right=ProfilerSourceAndEntity(
+                            profiler_source=profiler_source,
+                            entity=table,
+                        ),
+                    )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name=database.fullyQualifiedName.root,  # type: ignore
+                        error=f"Error listing source and entities for database due to [{exc}]",
+                        stackTrace=traceback.format_exc(),
+                    ),
+                    right=None,
+                )

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"""