openmetadata-ingestion 1.6.7.0__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/ingestion/source/search/opensearch/parser.py

@@ -0,0 +1,71 @@
+#  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.
+
+"""
+Utils module to parse the OpenSearch mapping json schema.
+"""
+
+import traceback
+from typing import List, Optional
+
+from metadata.generated.schema.entity.data.searchIndex import DataType, SearchIndexField
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+# If any type of OpenSearch field is not recognized mark it as unknown
+# pylint: disable=no-member,unused-argument,protected-access
+@classmethod
+def _missing_(cls, value):
+    return cls.UNKNOWN
+
+
+DataType._missing_ = _missing_
+
+
+def parse_os_index_mapping(mapping: dict) -> Optional[List[SearchIndexField]]:
+    """
+    Recursively convert the OpenSearch mapping into the required models.
+
+    Args:
+        mapping (dict): The OpenSearch mapping dictionary.
+
+    Returns:
+        Optional[List[SearchIndexField]]: A list of SearchIndexField objects if parsing is successful,
+        otherwise None.
+    """
+    field_models = []
+    try:
+        properties = mapping.get("properties", {})
+        for key, value in properties.items():
+            # Use the provided type if available, else default to OBJECT.
+            data_type = (
+                DataType(value.get("type").upper())
+                if value.get("type")
+                else DataType.OBJECT
+            )
+            field_models.append(
+                SearchIndexField(
+                    name=key,
+                    dataType=data_type,
+                    dataTypeDisplay=value.get("type"),
+                    description=value.get("description"),
+                    children=parse_os_index_mapping(value)
+                    if value.get("properties")
+                    else None,
+                )
+            )
+    except Exception as exc:  # pylint: disable=broad-except
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Unable to parse the index properties: {exc}")
+
+    return field_models

metadata/ingestion/source/search/opensearch/service_spec.py

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

metadata/ingestion/source/search/search_service.py

@@ -56,6 +56,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_search_index
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -138,6 +139,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     context = TopologyContextManager(topology)
     index_source_state: Set = set()
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

metadata/ingestion/source/storage/storage_service.py

@@ -62,6 +62,7 @@ from metadata.utils.datalake.datalake_utils import (
     DataFrameColumnParser,
     fetch_dataframe,
 )
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.storage_metadata_config import (
     StorageMetadataConfigException,
@@ -140,6 +141,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
 
     global_manifest: Optional[ManifestMetadataConfig]
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

metadata/ingestion/stage/table_usage.py

@@ -32,7 +32,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Stage
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import UTF_8
-from metadata.utils.helpers import init_staging_dir
+from metadata.utils.helpers import get_query_hash, init_staging_dir
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -62,8 +62,10 @@ class TableUsageStage(Stage):
         self.metadata = metadata
         self.table_usage = {}
         self.table_queries = {}
+        self.query_cost = {}
         init_staging_dir(self.config.filename)
         self.wrote_something = False
+        self.service_name = ""
 
     @property
     def name(self) -> str:
@@ -107,6 +109,7 @@ class TableUsageStage(Stage):
     def _add_sql_query(self, record, table):
         users, used_by = self._get_user_entity(record.userName)
         if self.table_queries.get((table, record.date)):
+            self.service_name = record.serviceName
             self.table_queries[(table, record.date)].append(
                 CreateQueryRequest(
                     query=record.sql,
@@ -172,6 +175,30 @@ class TableUsageStage(Stage):
             )
         yield Either(right=table)
 
+    def _handle_query_cost(self, parsed_data: ParsedData):
+        query_hash = get_query_hash(parsed_data.sql)
+        if (query_hash, parsed_data.date) in self.query_cost:
+            self.query_cost[(query_hash, parsed_data.date)].update(
+                {
+                    "cost": self.query_cost[(query_hash, parsed_data.date)]["cost"]
+                    + (parsed_data.cost or 0),
+                    "count": self.query_cost[(query_hash, parsed_data.date)]["count"]
+                    + 1,
+                    "totalDuration": self.query_cost[(query_hash, parsed_data.date)][
+                        "totalDuration"
+                    ]
+                    + (parsed_data.duration or 0),
+                }
+            )
+        else:
+            self.query_cost[(query_hash, parsed_data.date)] = {
+                "cost": parsed_data.cost or 0,
+                "count": 1,
+                "query": parsed_data.sql,
+                "dialect": parsed_data.dialect,
+                "totalDuration": parsed_data.duration or 0,
+            }
+
     def _run(self, record: QueryParserData) -> Iterable[Either[str]]:
         """
         Process the parsed data and store it in a file
@@ -187,9 +214,13 @@ class TableUsageStage(Stage):
                 yield from self._handle_table_usage(
                     parsed_data=parsed_data, table=table
                 )
+            self._handle_query_cost(parsed_data)
         self.dump_data_to_file()
 
     def dump_data_to_file(self):
+        """
+        Dump the table usage data to a file.
+        """
         for key, value in self.table_usage.items():
             if value:
                 value.sqlQueries = self.table_queries.get(key, [])
@@ -202,6 +233,27 @@ class TableUsageStage(Stage):
                     file.write(json.dumps(data))
                     file.write("\n")
 
+        for key, value in self.query_cost.items():
+            if value:
+                data = {
+                    "queryHash": key[0],
+                    "date": key[1],
+                    "cost": value["cost"],
+                    "count": value["count"],
+                    "query": value["query"],
+                    "dialect": value["dialect"],
+                    "totalDuration": value["totalDuration"],
+                }
+                with open(
+                    os.path.join(
+                        self.config.filename, f"{self.service_name}_{key[1]}_query"
+                    ),
+                    "a+",
+                    encoding=UTF_8,
+                ) as file:
+                    file.write(json.dumps(data))
+                    file.write("\n")
+
     def close(self) -> None:
         """
         Nothing to close. Data is being dumped inside a context manager

metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py

@@ -23,7 +23,6 @@ class UnityCatalogSamplerInterface(SQASampler):
     def get_client(self):
         """client is the session for SQA"""
         self.connection = databricks_get_connection(self.service_connection_config)
-        self.client = super().get_client()
-        self.set_catalog(self.client)
-
-        return self.client
+        client = super().get_client()
+        self.set_catalog(client)
+        return client

metadata/profiler/orm/converter/azuresql/converter.py

@@ -0,0 +1,39 @@
+#  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.
+
+"""
+Map Types to convert/cast mssql related data types to relevant data types
+"""
+
+
+from sqlalchemy import NVARCHAR, TEXT
+
+from metadata.profiler.orm.converter.common import CommonMapTypes
+from metadata.profiler.orm.registry import CustomImage, CustomTypes, DataType
+
+cast_dict = {
+    CustomImage: "VARBINARY(max)",
+    TEXT: "VARCHAR(max)",
+    NVARCHAR: "NVARCHAR(max)",
+}
+
+
+class AzureSqlMapTypes(CommonMapTypes):
+    """
+    AzureSql type mapper
+    """
+
+    def __init__(self) -> None:
+        self._TYPE_MAP.update(
+            {
+                DataType.TIMESTAMP: CustomTypes.TIMESTAMP.value,
+            }
+        )

metadata/profiler/orm/converter/converter_registry.py

@@ -17,6 +17,7 @@ from collections import defaultdict
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
+from metadata.profiler.orm.converter.azuresql.converter import AzureSqlMapTypes
 from metadata.profiler.orm.converter.bigquery.converter import BigqueryMapTypes
 from metadata.profiler.orm.converter.common import CommonMapTypes
 from metadata.profiler.orm.converter.mariadb.converter import MariaDBMapTypes
@@ -29,4 +30,5 @@ converter_registry[DatabaseServiceType.BigQuery] = BigqueryMapTypes
 converter_registry[DatabaseServiceType.Snowflake] = SnowflakeMapTypes
 converter_registry[DatabaseServiceType.Redshift] = RedshiftMapTypes
 converter_registry[DatabaseServiceType.Mssql] = MssqlMapTypes
+converter_registry[DatabaseServiceType.AzureSQL] = AzureSqlMapTypes
 converter_registry[DatabaseServiceType.MariaDB] = MariaDBMapTypes

metadata/profiler/orm/functions/length.py

@@ -43,6 +43,7 @@ def _(element, compiler, **kw):
 @compiles(LenFn, Dialects.MariaDB)
 @compiles(LenFn, Dialects.Athena)
 @compiles(LenFn, Dialects.Trino)
+@compiles(LenFn, Dialects.PinotDB)
 @compiles(LenFn, Dialects.Presto)
 @compiles(LenFn, Dialects.BigQuery)
 @compiles(LenFn, Dialects.Oracle)
@@ -55,6 +56,7 @@ def _(element, compiler, **kw):
     return "LENGTH(%s)" % compiler.process(element.clauses, **kw)
 
 
+@compiles(LenFn, Dialects.Cockroach)
 @compiles(LenFn, Dialects.Postgres)
 def _(element, compiler, **kw):
     return "LENGTH(CAST(%s AS text))" % compiler.process(element.clauses, **kw)

metadata/profiler/orm/functions/median.py

@@ -51,6 +51,25 @@ def _(elements, compiler, **kwargs):
     return "percentile_cont(%s , %s) OVER()" % (col, percentile)
 
 
+@compiles(MedianFn, Dialects.Databricks)
+def _(elements, compiler, **kwargs):
+    col, _, percentile = [
+        compiler.process(element, **kwargs) for element in elements.clauses
+    ]
+    return "percentile_approx(%s , %s)" % (col, percentile)
+
+
+# pylint: disable=unused-argument
+@compiles(MedianFn, Dialects.Cockroach)
+def _(elements, compiler, **kwargs):
+    col = compiler.process(elements.clauses.clauses[0])
+    percentile = elements.clauses.clauses[2].value
+    return "percentile_cont(%.2f) WITHIN GROUP (ORDER BY %s ASC)" % (
+        percentile,
+        f"(({col})::float8)",
+    )
+
+
 @compiles(MedianFn, Dialects.ClickHouse)
 def _(elements, compiler, **kwargs):
     col, _, percentile = [

metadata/profiler/orm/functions/modulo.py

@@ -60,6 +60,7 @@ def _(element, compiler, **kw):
 @compiles(ModuloFn, Dialects.Db2)
 @compiles(ModuloFn, Dialects.Vertica)
 @compiles(ModuloFn, Dialects.Hana)
+@compiles(ModuloFn, Dialects.Cockroach)
 def _(element, compiler, **kw):
     """Modulo function for specific dialect"""
     value, base = validate_and_compile(element, compiler, **kw)

metadata/profiler/orm/functions/random_num.py

@@ -85,6 +85,7 @@ def _(*_, **__):
 
 
 @compiles(RandomNumFn, Dialects.Postgres)
+@compiles(RandomNumFn, Dialects.Cockroach)
 def _(*_, **__):
     """Postgres random logic"""
     return "ABS((RANDOM() * 100)::INTEGER)"

metadata/profiler/orm/functions/sum.py

@@ -50,6 +50,7 @@ def _(element, compiler, **kw):
 
 @compiles(SumFn, Dialects.BigQuery)
 @compiles(SumFn, Dialects.Postgres)
+@compiles(SumFn, Dialects.Cockroach)
 def _(element, compiler, **kw):
     """Handle case where column type is INTEGER but SUM returns a NUMBER"""
     proc = compiler.process(element.clauses, **kw)

metadata/profiler/orm/registry.py

@@ -63,6 +63,7 @@ class PythonDialects(Enum):
     AzureSQL = "azuresql"
     BigQuery = "bigquery"
     ClickHouse = "clickhouse"
+    Cockroach = "cockroachdb"
     Databricks = "databricks"
     Db2 = "db2"
     Doris = "pydoris"
@@ -78,6 +79,7 @@ class PythonDialects(Enum):
     MSSQL = "mssql"
     MySQL = "mysql"
     Oracle = "oracle"
+    PinotDB = "pinotdb"
     Postgres = "postgresql"
     Presto = "presto"
     Redshift = "redshift"

metadata/profiler/processor/core.py

@@ -165,6 +165,7 @@ class Profiler(Generic[TMetric]):
                 column
                 for column in self.profiler_interface.get_columns()
                 if column.name in self._get_included_columns()
+                or self._get_included_columns() == {"all"}
             ]
 
         if not self._get_included_columns():

metadata/profiler/processor/metric_filter.py

@@ -230,7 +230,7 @@ class MetricFilter:
             (
                 include_columns.metrics
                 for include_columns in columns_config or []
-                if include_columns.columnName == column.name
+                if include_columns.columnName in {column.name, "all"}
             ),
             None,
         )

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

@@ -153,6 +153,7 @@ class ProfilerSource(ProfilerSourceInterface):
                 profileSample=self.source_config.profileSample,
                 profileSampleType=self.source_config.profileSampleType,
                 samplingMethodType=self.source_config.samplingMethodType,
+                randomizedSample=self.source_config.randomizedSample,
             ),
         )
 

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -220,6 +220,10 @@ class DatabaseFetcherStrategy(FetcherStrategy):
             entity=Database,
             params={"service": self.config.source.serviceName},
         )
+        if not databases:
+            raise ValueError(
+                f"No databases found for service {self.config.source.serviceName}"
+            )
         databases = cast(Iterable[Database], databases)
 
         if self.source_config.databaseFilterPattern:

metadata/sampler/config.py

@@ -122,6 +122,7 @@ def get_profile_sample_config(
                     profileSample=config.profileSample,
                     profileSampleType=config.profileSampleType,
                     samplingMethodType=config.samplingMethodType,
+                    randomizedSample=config.randomizedSample,
                 )
         except AttributeError:
             pass

metadata/sampler/models.py

@@ -41,6 +41,7 @@ class BaseProfileConfig(ConfigModel):
     profileSampleType: Optional[ProfileSampleType] = None
     samplingMethodType: Optional[SamplingMethodType] = None
     sampleDataCount: Optional[int] = 100
+    randomizedSample: Optional[bool] = True
 
 
 class ColumnConfig(ConfigModel):
@@ -56,6 +57,7 @@ class TableConfig(BaseProfileConfig):
     profileQuery: Optional[str] = None
     partitionConfig: Optional[PartitionProfilerConfig] = None
     columnConfig: Optional[ColumnConfig] = None
+    randomizedSample: Optional[bool] = True
 
     @classmethod
     def from_database_and_schema_config(
@@ -104,3 +106,4 @@ class SampleConfig(ConfigModel):
     profileSample: Optional[Union[float, int]] = None
     profileSampleType: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
     samplingMethodType: Optional[SamplingMethodType] = None
+    randomizedSample: Optional[bool] = True

metadata/sampler/nosql/sampler.py

@@ -24,6 +24,10 @@ class NoSQLSampler(SamplerInterface):
 
     client: NoSQLAdaptor
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = self.get_client()
+
     @property
     def raw_dataset(self):
         return self.entity

metadata/sampler/pandas/sampler.py

@@ -46,6 +46,7 @@ class DatalakeSampler(SamplerInterface, PandasInterfaceMixin):
         super().__init__(*args, **kwargs)
         self.partition_details = cast(PartitionProfilerConfig, self.partition_details)
         self._table = None
+        self.client = self.get_client()
 
     @property
     def raw_dataset(self):

metadata/sampler/partition.py

@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
 from metadata.sampler.models import TableConfig
+from metadata.utils.logger import sampler_logger
+
+logger = sampler_logger()
 
 
 def validate_athena_injected_partitioning(
@@ -164,8 +167,12 @@ def _handle_bigquery_partition(
                 partitionIntegerRangeEnd=10000,
             )
         # TODO: Allow External Hive Partitioning for profiler
-        raise TypeError(
-            f"Unsupported partition type {partition.intervalType}. Skipping table"
+        logger.info(
+            f"Unknown partition type {partition.intervalType}. Cannot infer partition. Partitionning "
+            "will be set to None. If partitionning is required for your table, you can set it "
+            "directly in the table config. More information can be found at "
+            "https://docs.open-metadata.org/latest/how-to-guides/data-quality-observability/profiler/"
+            "workflow#4.-updating-profiler-setting-at-the-table-level"
         )
 
     return None

metadata/sampler/sampler_interface.py

@@ -86,7 +86,6 @@ class SamplerInterface(ABC):
 
         self.service_connection_config = service_connection_config
         self.connection = get_ssl_connection(self.service_connection_config)
-        self.client = self.get_client()
 
     # pylint: disable=too-many-arguments, too-many-locals
     @classmethod

metadata/sampler/sqlalchemy/azuresql/sampler.py

@@ -14,10 +14,11 @@ for the profiler
 """
 from typing import List, Optional
 
-from sqlalchemy import Column
+from sqlalchemy import Column, Table, text
+from sqlalchemy.orm import Query
 
-from metadata.generated.schema.entity.data.table import TableData
-from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.generated.schema.entity.data.table import TableData, TableType
+from metadata.sampler.sqlalchemy.sampler import ProfileSampleType, SQASampler
 
 
 class AzureSQLSampler(SQASampler):
@@ -31,6 +32,32 @@ class AzureSQLSampler(SQASampler):
     # pyodbc.ProgrammingError: ('ODBC SQL type -151 is not yet supported.  column-index=x  type=-151', 'HY106')
     NOT_COMPUTE_PYODBC = {"SQASGeography", "UndeterminedType"}
 
+    def set_tablesample(self, selectable: Table):
+        """Set the TABLESAMPLE clause for MSSQL
+        Args:
+            selectable (Table): _description_
+        """
+        if self.entity.tableType != TableType.View:
+            if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
+                return selectable.tablesample(
+                    text(f"{self.sample_config.profileSample or 100} PERCENT")
+                )
+
+            return selectable.tablesample(
+                text(f"{int(self.sample_config.profileSample or 100)} ROWS")
+            )
+
+        return selectable
+
+    def get_sample_query(self, *, column=None) -> Query:
+        """get query for sample data"""
+        rnd = self._base_sample_query(column).cte(
+            f"{self.get_sampler_table_name()}_rnd"
+        )
+        with self.get_client() as client:
+            query = client.query(rnd)
+        return query.cte(f"{self.get_sampler_table_name()}_sample")
+
     def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:
         sqa_columns = []
         if columns:

metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -48,5 +48,6 @@ class MssqlSampler(SQASampler):
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        query = self.client.query(rnd)
+        with self.get_client() as client:
+            query = client.query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")