openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/profiler/interface/sqlalchemy/bigquery/profiler_interface.py

@@ -43,25 +43,6 @@ class BigQueryProfilerInterface(SQAProfilerInterface):
                 columns_list.extend(col)
         return columns_list
 
-    def _get_sampler(self, **kwargs):
-        """get sampler object"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        session = kwargs.get("session")
-        table = kwargs["table"]
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            client=session or self.session,
-            table=table,
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-            table_type=self.table_entity.tableType,
-        )
-
     def get_columns(self) -> Column:
         """Get columns from table"""
         # pylint: disable=import-outside-toplevel

metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py

@@ -16,10 +16,10 @@ supporting sqlalchemy abstraction layer
 from typing import List
 
 from pyhive.sqlalchemy_hive import HiveCompiler
-from sqlalchemy import Column, inspect
+from sqlalchemy import Column
 
 from metadata.generated.schema.entity.data.table import Column as OMColumn
-from metadata.generated.schema.entity.data.table import ColumnName, DataType, TableData
+from metadata.generated.schema.entity.data.table import ColumnName, DataType
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
@@ -97,18 +97,3 @@ class DatabricksProfilerInterface(SQAProfilerInterface):
                 )
                 columns.append(col)
         return columns
-
-    def fetch_sample_data(self, table, columns) -> TableData:
-        """Fetch sample data from database
-
-        Args:
-            table: ORM declarative table
-
-        Returns:
-            TableData: sample table data
-        """
-        sampler = self._get_sampler(
-            table=table,
-        )
-
-        return sampler.fetch_sample_data(list(inspect(self.table).c))

metadata/profiler/interface/sqlalchemy/db2/profiler_interface.py

@@ -32,7 +32,7 @@ class DB2ProfilerInterface(SQAProfilerInterface):
         # pylint: disable=protected-access
         if exc.orig and "overflow" in exc.orig._message:
             logger.info(
-                f"Computing metrics without sum for {runner.table.__tablename__}.{column.name}"
+                f"Computing metrics without sum for {runner.table_name}.{column.name}"
             )
             return self._compute_static_metrics_wo_sum(metrics, runner, session, column)
         return None

metadata/profiler/interface/sqlalchemy/mariadb/profiler_interface.py

@@ -77,11 +77,11 @@ class MariaDBProfilerInterface(SQAProfilerInterface):
                 return dict(row)
         except ProgrammingError:
             logger.info(
-                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to overflow"
+                f"Skipping window metrics for {runner.table_name}.{column.name} due to overflow"
             )
             return None
 
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -21,7 +21,7 @@ import threading
 import traceback
 from collections import defaultdict
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type, Union
 
 from sqlalchemy import Column, inspect, text
 from sqlalchemy.exc import DBAPIError, ProgrammingError, ResourceClosedError
@@ -30,10 +30,18 @@ from sqlalchemy.orm import scoped_session
 from metadata.generated.schema.entity.data.table import (
     CustomMetricProfile,
     SystemProfile,
-    TableData,
+    Table,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    DatabaseServiceProfilerPipeline,
 )
 from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.mixins.sqalchemy.sqa_mixin import SQAInterfaceMixin
 from metadata.profiler.api.models import ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
@@ -47,7 +55,7 @@ from metadata.profiler.orm.functions.table_metric_computer import TableMetricCom
 from metadata.profiler.orm.registry import Dialects
 from metadata.profiler.processor.metric_filter import MetricFilter
 from metadata.profiler.processor.runner import QueryRunner
-from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.custom_thread_pool import CustomThreadPoolExecutor
 from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
@@ -75,40 +83,32 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
     """
 
     # pylint: disable=too-many-arguments
-
     def __init__(
         self,
-        service_connection_config,
-        ometa_client,
-        entity,
-        storage_config,
-        profile_sample_config,
-        source_config,
-        sample_query,
-        table_partition_config,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        source_config: DatabaseServiceProfilerPipeline,
+        sampler: SamplerInterface,
         thread_count: int = 5,
         timeout_seconds: int = 43200,
-        sqa_metadata=None,
-        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
         **kwargs,
     ):
         """Instantiate SQA Interface object"""
+        self.session_factory = None
+        self.session = None
 
         super().__init__(
-            service_connection_config,
-            ometa_client,
-            entity,
-            storage_config,
-            profile_sample_config,
-            source_config,
-            sample_query,
-            table_partition_config,
-            thread_count,
-            timeout_seconds,
-            sample_data_count,
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            source_config=source_config,
+            sampler=sampler,
+            thread_count=thread_count,
+            timeout_seconds=timeout_seconds,
         )
 
-        self._table = self._convert_table_to_orm_object(sqa_metadata)
+        self._table = self.sampler.raw_dataset
         self.create_session()
         self.system_metrics_computer = self.initialize_system_metrics_computer()
 
@@ -126,25 +126,6 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
     def table(self):
         return self._table
 
-    def _get_sampler(self, **kwargs):
-        """get sampler object"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        session = kwargs.get("session")
-        table = kwargs["table"]
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            client=session or self.session,
-            table=table,
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-            sample_data_count=self.sample_data_count,
-        )
-
     def _session_factory(self) -> scoped_session:
         """Create thread safe session that will be automatically
         garbage collected once the application thread ends
@@ -175,7 +156,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
             )
             return dict(row)
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -213,7 +194,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to compute profile for {runner.table.__tablename__}: {exc}"  # type: ignore
+                f"Error trying to compute profile for {runner.table_name}: {exc}"  # type: ignore
             )
             session.rollback()
             raise RuntimeError(exc)
@@ -250,7 +231,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 runner, column, exc, session, metrics
             )
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -293,10 +274,10 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 runner._session.get_bind().dialect.name
                 != Dialects.Druid
             ):
-                msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+                msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
                 handle_query_exception(msg, exc, session)
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -329,10 +310,10 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 return dict(row)
         except ProgrammingError as exc:
             logger.info(
-                f"Skipping metrics for {runner.table.__tablename__}.{column.name} due to {exc}"
+                f"Skipping metrics for {runner.table_name}.{column.name} due to {exc}"
             )
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -366,7 +347,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 )
 
             except Exception as exc:
-                msg = f"Error trying to compute profile for {runner.table.__tablename__}.{metric.columnName}: {exc}"
+                msg = f"Error trying to compute profile for {runner.table_name}.{metric.columnName}: {exc}"
                 logger.debug(traceback.format_exc())
                 logger.warning(msg)
         if custom_metrics:
@@ -390,39 +371,21 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         Returns:
             dictionnary of results
         """
-        logger.debug(f"Computing system metrics for {runner.table.__tablename__}")
-        return self.system_metrics_computer.get_system_metrics(table=runner.table)
+        logger.debug(f"Computing system metrics for {runner.table_name}")
+        return self.system_metrics_computer.get_system_metrics(table=runner.dataset)
 
-    def _create_thread_safe_sampler(
-        self,
-        session,
-        table,
-    ):
-        """Create thread safe runner"""
-        if not hasattr(thread_local, "sampler"):
-            thread_local.sampler = self._get_sampler(
-                table=table,
-                session=session,
-            )
-        return thread_local.sampler
-
-    def _create_thread_safe_runner(
-        self,
-        session,
-        table,
-        sample,
-    ):
+    def _create_thread_safe_runner(self, session, column=None):
         """Create thread safe runner"""
         if not hasattr(thread_local, "runner"):
             thread_local.runner = QueryRunner(
                 session=session,
-                table=table,
-                sample=sample,
-                partition_details=self.partition_details,
-                profile_sample_query=self.profile_query,
+                dataset=self.sampler.get_dataset(column=column),
+                raw_dataset=self.sampler.raw_dataset,
+                partition_details=self.sampler.partition_details,
+                profile_sample_query=self.sampler.sample_query,
             )
             return thread_local.runner
-        thread_local.runner._sample = sample  # pylint: disable=protected-access
+        thread_local.runner.dataset = self.sampler.get_dataset(column=column)
         return thread_local.runner
 
     def compute_metrics_in_thread(
@@ -437,16 +400,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         with Session() as session:
             self.set_session_tag(session)
             self.set_catalog(session)
-            sampler = self._create_thread_safe_sampler(
-                session,
-                metric_func.table,
-            )
-            sample = sampler.random_sample(metric_func.column)
-            runner = self._create_thread_safe_runner(
-                session,
-                metric_func.table,
-                sample,
-            )
+            runner = self._create_thread_safe_runner(session, metric_func.column)
             row = None
             try:
                 row = self._get_metric_fn[metric_func.metric_type.value](
@@ -454,7 +408,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                     runner=runner,
                     session=session,
                     column=metric_func.column,
-                    sample=sample,
+                    sample=runner.dataset,
                 )
                 if isinstance(row, dict):
                     row = self._validate_nulls(row)
@@ -546,21 +500,6 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
 
         return profile_results
 
-    def fetch_sample_data(self, table, columns) -> TableData:
-        """Fetch sample data from database
-
-        Args:
-            table: ORM declarative table
-
-        Returns:
-            TableData: sample table data
-        """
-        sampler = self._get_sampler(
-            table=table,
-        )
-
-        return sampler.fetch_sample_data(columns)
-
     def get_composed_metrics(
         self, column: Column, metric: Metrics, column_results: Dict
     ):
@@ -586,7 +525,6 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         column: Column,
         metric: Type[HybridMetric],
         column_results: Dict[str, Any],
-        **kwargs,
     ):
         """Given a list of metrics, compute the given results
         and returns the values
@@ -598,10 +536,9 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         Returns:
             dictionnary of results
         """
-        sampler = self._get_sampler(table=kwargs.get("table"))
-        sample = sampler.random_sample(column)
+        dataset = self.sampler.get_dataset(column=column)
         try:
-            return metric(column).fn(sample, column_results, self.session)
+            return metric(column).fn(dataset, column_results, self.session)
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unexpected exception computing metrics: {exc}")

metadata/profiler/interface/sqlalchemy/single_store/profiler_interface.py

@@ -76,11 +76,11 @@ class SingleStoreProfilerInterface(SQAProfilerInterface):
                 return dict(row)
         except ProgrammingError:
             logger.info(
-                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to overflow"
+                f"Skipping window metrics for {runner.table_name}.{column.name} due to overflow"
             )
             return None
 
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None

metadata/profiler/interface/sqlalchemy/snowflake/profiler_interface.py

@@ -41,7 +41,7 @@ class SnowflakeProfilerInterface(SQAProfilerInterface):
             session.bind.dialect.name
         ):
             logger.info(
-                f"Computing metrics without sum for {runner.table.__tablename__}.{column.name}"
+                f"Computing metrics without sum for {runner.table_name}.{column.name}"
             )
             return self._compute_static_metrics_wo_sum(metrics, runner, session, column)
         return None

metadata/profiler/interface/sqlalchemy/stored_statistics_profiler.py

@@ -79,8 +79,8 @@ class ProfilerWithStatistics(SQAProfilerInterface, StoredStatisticsSource):
                 list,
                 partition(self.is_statistic_metric, metrics),
             )
-            schema = runner.table.__table_args__["schema"]
-            table_name = runner.table.__tablename__
+            schema = runner.schema_name
+            table_name = runner.table_name
             logger.debug(
                 "Getting statistics for column: %s.%s.%s",
                 schema,
@@ -118,8 +118,8 @@ class ProfilerWithStatistics(SQAProfilerInterface, StoredStatisticsSource):
                 list,
                 partition(self.is_statistic_metric, metrics),
             )
-            schema = runner.table.__table_args__["schema"]
-            table_name = runner.table.__tablename__
+            schema = runner.schema_name
+            table_name = runner.table_name
             logger.debug("Geting statistics for table: %s.%s", schema, table_name)
             result.update(
                 super().get_table_statistics(stat_metrics, schema, table_name)
@@ -135,12 +135,10 @@ class ProfilerWithStatistics(SQAProfilerInterface, StoredStatisticsSource):
             result.update(super_table_metrics)
         return result
 
-    def get_hybrid_metrics(
-        self, column: Column, metric: Metric, column_results: Dict, **kwargs
-    ):
+    def get_hybrid_metrics(self, column: Column, metric: Metric, column_results: Dict):
         # this metrics might have been computed in a previous step
         return column_results.get(metric.name()) or super().get_hybrid_metrics(
-            column, metric, column_results, **kwargs
+            column, metric, column_results
         )
 
     def is_statistic_metric(self, metric: Metric) -> bool:

metadata/profiler/interface/sqlalchemy/trino/profiler_interface.py

@@ -76,11 +76,11 @@ class TrinoProfilerInterface(ProfilerWithStatistics, TrinoStoredStatisticsSource
                 return dict(row)
         except ProgrammingError as err:
             logger.info(
-                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to {err}"
+                f"Skipping window metrics for {runner.table_name}.{column.name} due to {err}"
             )
             return None
 
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None

metadata/profiler/orm/converter/base.py

@@ -109,10 +109,29 @@ def ometa_to_sqa_orm(
     We are building the class dynamically using
     `type` and passing SQLAlchemy `Base` class
     as the bases tuple for inheritance.
+
+    Args:
+        table (Table): OpenMetadata Table instance
+        metadata (OpenMetadata): OpenMetadata connection
+        sqa_metadata_obj (MetaData): For advanced use cases, you can pass a custom MetaData object. For most cases, this
+        can be left as None so that the global_metadata object is used.
     """
+    _metadata = sqa_metadata_obj or Base.metadata
     table.serviceType = cast(
         databaseService.DatabaseServiceType, table.serviceType
     )  # satisfy mypy
+
+    orm_database_name = get_orm_database(table, metadata)
+    # SQLite does not support schemas
+    orm_schema_name = (
+        get_orm_schema(table, metadata)
+        if table.serviceType != databaseService.DatabaseServiceType.SQLite
+        else None
+    )
+    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.root}".replace(
+        ".", "_"
+    )
+
     cols = {
         (
             col.name.root + "_"
@@ -122,12 +141,6 @@ def ometa_to_sqa_orm(
         for idx, col in enumerate(table.columns)
     }
 
-    orm_database_name = get_orm_database(table, metadata)
-    orm_schema_name = get_orm_schema(table, metadata)
-    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.root}".replace(
-        ".", "_"
-    )
-
     # Type takes positional arguments in the form of (name, bases, dict)
     orm = type(
         orm_name,  # Output class name
@@ -135,10 +148,7 @@ def ometa_to_sqa_orm(
         {
             "__tablename__": str(table.name.root),
             "__table_args__": {
-                # SQLite does not support schemas
-                "schema": orm_schema_name
-                if table.serviceType != databaseService.DatabaseServiceType.SQLite
-                else None,
+                "schema": orm_schema_name,
                 "extend_existing": True,  # Recreates the table ORM object if it already exists. Useful for testing
                 "quote": check_snowflake_case_sensitive(
                     table.serviceType, table.name.root
@@ -146,13 +156,12 @@ def ometa_to_sqa_orm(
                 or None,
             },
             **cols,
-            "metadata": sqa_metadata_obj or Base.metadata,
+            "metadata": _metadata,
         },
     )
 
     if not isinstance(orm, DeclarativeMeta):
         raise ValueError("OMeta to ORM did not create a DeclarativeMeta")
-
     return orm
 
 

metadata/profiler/orm/functions/table_metric_computer.py

@@ -53,7 +53,7 @@ class AbstractTableMetricComputer(ABC):
         self._metrics = metrics
         self._conn_config = conn_config
         self._database = self._runner._session.get_bind().url.database
-        self._table = self._runner.table
+        self._table = self._runner.dataset
         self._entity = entity
 
     @property
@@ -91,8 +91,8 @@ class AbstractTableMetricComputer(ABC):
             table (DeclarativeMeta): _description_
         """
         try:
-            self._schema_name = self.table.__table_args__.get("schema")
-            self._table_name = self.table.__tablename__
+            self._schema_name = self.table.schema
+            self._table_name = self.table.name
         except AttributeError:
             raise AttributeError(ERROR_MSG)
 
@@ -147,7 +147,8 @@ class BaseTableMetricComputer(AbstractTableMetricComputer):
     """Base table computer"""
 
     def compute(self):
-        """Default compute behavior for table metrics"""
+        """Default compute behavior for table metrics. This method will use the raw table
+        to compute metrics and omit any sampling or partitioning logic."""
         return self.runner.select_first_from_table(
             *[metric().fn() for metric in self.metrics]
         )