openmetadata-ingestion 1.6.0.0rc3__py3-none-any.whl → 1.6.2.0__py3-none-any.whl

metadata/ingestion/source/storage/s3/metadata.py

@@ -12,6 +12,7 @@
 import json
 import secrets
 import traceback
+from copy import deepcopy
 from datetime import datetime, timedelta
 from enum import Enum
 from typing import Dict, Iterable, List, Optional, Tuple
@@ -325,6 +326,45 @@ class S3Source(StorageServiceSource):
                 )
         return None
 
+    def _generate_structured_containers_by_depth(
+        self,
+        bucket_response: S3BucketResponse,
+        metadata_entry: MetadataEntry,
+        parent: Optional[EntityReference] = None,
+    ) -> Iterable[S3ContainerDetails]:
+        try:
+            prefix = self._get_sample_file_prefix(metadata_entry=metadata_entry)
+            if prefix:
+                response = self.s3_client.list_objects_v2(
+                    Bucket=bucket_response.name, Prefix=prefix
+                )
+                # total depth is depth of prefix + depth of the metadata entry
+                total_depth = metadata_entry.depth + len(prefix[:-1].split("/"))
+                candidate_keys = {
+                    "/".join(entry.get("Key").split("/")[:total_depth]) + "/"
+                    for entry in response[S3_CLIENT_ROOT_RESPONSE]
+                    if entry
+                    and entry.get("Key")
+                    and len(entry.get("Key").split("/")) > total_depth
+                }
+                for key in candidate_keys:
+                    metadata_entry_copy = deepcopy(metadata_entry)
+                    metadata_entry_copy.dataPath = key.strip(KEY_SEPARATOR)
+                    structured_container: Optional[
+                        S3ContainerDetails
+                    ] = self._generate_container_details(
+                        bucket_response=bucket_response,
+                        metadata_entry=metadata_entry_copy,
+                        parent=parent,
+                    )
+                    if structured_container:
+                        yield structured_container
+        except Exception as err:
+            logger.warning(
+                f"Error while generating structured containers by depth for {metadata_entry.dataPath} - {err}"
+            )
+            logger.debug(traceback.format_exc())
+
     def _generate_structured_containers(
         self,
         bucket_response: S3BucketResponse,
@@ -336,15 +376,22 @@ class S3Source(StorageServiceSource):
                 f"Extracting metadata from path {metadata_entry.dataPath.strip(KEY_SEPARATOR)} "
                 f"and generating structured container"
             )
-            structured_container: Optional[
-                S3ContainerDetails
-            ] = self._generate_container_details(
-                bucket_response=bucket_response,
-                metadata_entry=metadata_entry,
-                parent=parent,
-            )
-            if structured_container:
-                yield structured_container
+            if metadata_entry.depth == 0:
+                structured_container: Optional[
+                    S3ContainerDetails
+                ] = self._generate_container_details(
+                    bucket_response=bucket_response,
+                    metadata_entry=metadata_entry,
+                    parent=parent,
+                )
+                if structured_container:
+                    yield structured_container
+            else:
+                yield from self._generate_structured_containers_by_depth(
+                    bucket_response=bucket_response,
+                    metadata_entry=metadata_entry,
+                    parent=parent,
+                )
 
     def is_valid_unstructured_file(self, accepted_extensions: List, key: str) -> bool:
         # Split the string into a list of values

metadata/profiler/metrics/static/max.py

@@ -88,10 +88,12 @@ class Max(StaticMetric):
     def fn(self):
         """sqlalchemy function"""
         if is_concatenable(self.col.type):
-            return MaxFn(LenFn(column(self.col.name, self.col.type)))
+            return MaxFn(
+                LenFn(column(self.col.name, self.col.type)), type_=self.col.type
+            )
         if (not is_quantifiable(self.col.type)) and (not is_date_time(self.col.type)):
             return None
-        return MaxFn(column(self.col.name, self.col.type))
+        return MaxFn(column(self.col.name, self.col.type), type_=self.col.type)
 
     def df_fn(self, dfs=None):
         """pandas function"""

metadata/profiler/metrics/static/min.py

@@ -88,11 +88,13 @@ class Min(StaticMetric):
     def fn(self):
         """sqlalchemy function"""
         if is_concatenable(self.col.type):
-            return MinFn(LenFn(column(self.col.name, self.col.type)))
+            return MinFn(
+                LenFn(column(self.col.name, self.col.type)), type_=self.col.type
+            )
 
         if (not is_quantifiable(self.col.type)) and (not is_date_time(self.col.type)):
             return None
-        return MinFn(column(self.col.name, self.col.type))
+        return MinFn(column(self.col.name, self.col.type), type_=self.col.type)
 
     def df_fn(self, dfs=None):
         """pandas function"""

metadata/profiler/orm/converter/common.py

@@ -76,7 +76,9 @@ class CommonMapTypes:
         """returns an ORM type"""
 
         if col.arrayDataType:
-            return self._TYPE_MAP.get(col.dataType)(item_type=col.arrayDataType)
+            return self._TYPE_MAP.get(col.dataType)(
+                item_type=self._TYPE_MAP.get(col.arrayDataType)
+            )
         return self.return_custom_type(col, table_service_type)
 
     def return_custom_type(self, col: Column, _):

metadata/profiler/orm/converter/converter_registry.py

@@ -19,6 +19,7 @@ from metadata.generated.schema.entity.services.databaseService import (
 )
 from metadata.profiler.orm.converter.bigquery.converter import BigqueryMapTypes
 from metadata.profiler.orm.converter.common import CommonMapTypes
+from metadata.profiler.orm.converter.mssql.converter import MssqlMapTypes
 from metadata.profiler.orm.converter.redshift.converter import RedshiftMapTypes
 from metadata.profiler.orm.converter.snowflake.converter import SnowflakeMapTypes
 
@@ -26,3 +27,4 @@ converter_registry = defaultdict(lambda: CommonMapTypes)
 converter_registry[DatabaseServiceType.BigQuery] = BigqueryMapTypes
 converter_registry[DatabaseServiceType.Snowflake] = SnowflakeMapTypes
 converter_registry[DatabaseServiceType.Redshift] = RedshiftMapTypes
+converter_registry[DatabaseServiceType.Mssql] = MssqlMapTypes

metadata/profiler/orm/converter/mssql/converter.py

@@ -16,10 +16,24 @@ Map Types to convert/cast mssql related data types to relevant data types
 
 from sqlalchemy import NVARCHAR, TEXT
 
-from metadata.profiler.orm.registry import CustomImage
+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 MssqlMapTypes(CommonMapTypes):
+    """
+    Mssql type mapper
+    """
+
+    def __init__(self) -> None:
+        self._TYPE_MAP.update(
+            {
+                DataType.TIMESTAMP: CustomTypes.TIMESTAMP.value,
+            }
+        )

metadata/profiler/orm/functions/md5.py

@@ -10,15 +10,13 @@
 #  limitations under the License.
 
 """
-Define Concat function
+Define MD5 hashing function
 """
-# Keep SQA docs style defining custom constructs
-
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.functions import FunctionElement
 
 from metadata.profiler.metrics.core import CACHE
-from metadata.profiler.orm.registry import Dialects
+from metadata.profiler.orm.registry import PythonDialects
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -33,7 +31,12 @@ def _(element, compiler, **kw):
     return f"MD5({compiler.process(element.clauses, **kw)})"
 
 
-@compiles(MD5, Dialects.MSSQL)
+@compiles(MD5, PythonDialects.MSSQL.value)
 def _(element, compiler, **kw):
     # TODO requires separate where clauses for each table
     return f"CONVERT(VARCHAR(8), HashBytes('MD5', {compiler.process(element.clauses, **kw)}), 2)"
+
+
+@compiles(MD5, PythonDialects.BigQuery.value)
+def _(element, compiler, **kw):
+    return f"TO_HEX(MD5(CAST({compiler.process(element.clauses, **kw)} AS STRING)))"

metadata/profiler/orm/types/custom_timestamp.py

@@ -45,6 +45,8 @@ class CustomTimestamp(TypeDecorator):
         import struct  # pylint: disable=import-outside-toplevel
 
         if dialect.name == "mssql" and isinstance(value, bytes):
-            bytes_to_int = struct.unpack(">Q", value)[0]
-            return bytes_to_int
+            unpacked_value = struct.unpack("@Q", value)
+            if unpacked_value:
+                return unpacked_value[0]
+            return None
         return value

metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -0,0 +1,50 @@
+#  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.
+"""
+Helper module to handle data sampling
+for the profiler
+"""
+
+
+from sqlalchemy import Table, text
+from sqlalchemy.sql.selectable import CTE
+
+from metadata.generated.schema.entity.data.table import ProfileSampleType
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+
+
+class MssqlSampler(SQASampler):
+    """
+    Generates a sample of the data to not
+    run the query in the whole table.
+    """
+
+    def set_tablesample(self, selectable: Table):
+        """Set the TABLESAMPLE clause for MSSQL
+        Args:
+            selectable (Table): _description_
+        """
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            return selectable.tablesample(
+                text(f"{self.sample_config.profile_sample or 100} PERCENT")
+            )
+
+        return selectable.tablesample(
+            text(f"{int(self.sample_config.profile_sample or 100)} ROWS")
+        )
+
+    def get_sample_query(self, *, column=None) -> CTE:
+        """get query for sample data"""
+        rnd = self._base_sample_query(column).cte(
+            f"{self.raw_dataset.__tablename__}_rnd"
+        )
+        query = self.client.query(rnd)
+        return query.cte(f"{self.raw_dataset.__tablename__}_sample")

metadata/utils/constants.py

@@ -12,12 +12,17 @@
 """
 Define constants useful for the metadata ingestion
 """
+from metadata.generated.schema.entity.data.apiCollection import APICollection
+from metadata.generated.schema.entity.data.apiEndpoint import APIEndpoint
 from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
 from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.glossary import Glossary
+from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
+from metadata.generated.schema.entity.data.metric import Metric
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.searchIndex import SearchIndex
@@ -42,12 +47,6 @@ from metadata.generated.schema.entity.services.connections.database.deltaLakeCon
 from metadata.generated.schema.entity.services.connections.database.domoDatabaseConnection import (
     DomoDatabaseType,
 )
-from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
-    DorisType,
-)
-from metadata.generated.schema.entity.services.connections.database.druidConnection import (
-    DruidType,
-)
 from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
     DynamoDBType,
 )
@@ -69,9 +68,6 @@ from metadata.generated.schema.entity.services.connections.database.sapErpConnec
 from metadata.generated.schema.entity.services.connections.database.sasConnection import (
     SasType,
 )
-from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
-    DatabricksType,
-)
 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
@@ -131,6 +127,8 @@ ENTITY_REFERENCE_CLASS_MAP = {
     "metadataService": MetadataService,
     "searchService": SearchService,
     # Data Asset Entities
+    "apiCollection": APICollection,
+    "apiEndpoint": APIEndpoint,
     "table": Table,
     "storedProcedure": StoredProcedure,
     "database": Database,
@@ -149,6 +147,10 @@ ENTITY_REFERENCE_CLASS_MAP = {
     # Domain
     "domain": Domain,
     "dataProduct": DataProduct,
+    # Governance
+    "metric": Metric,
+    "glossary": Glossary,
+    "glossaryTerm": GlossaryTerm,
 }
 
 ENTITY_REFERENCE_TYPE_MAP = {
@@ -161,11 +163,8 @@ NON_SQA_DATABASE_CONNECTIONS = (
     DatalakeType.Datalake.value,
     BigtableType.BigTable.value,
     CouchbaseType.Couchbase.value,
-    DatabricksType.UnityCatalog.value,
     DeltaLakeType.DeltaLake.value,
     DomoDatabaseType.DomoDatabase.value,
-    DorisType.Doris.value,
-    DruidType.Druid.value,
     DynamoDBType.DynamoDB.value,
     GlueType.Glue.value,
     IcebergType.Iceberg.value,

metadata/utils/elasticsearch.py

@@ -18,6 +18,8 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.analytics.reportData import ReportData
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.data.apiCollection import APICollection
+from metadata.generated.schema.entity.data.apiEndpoint import APIEndpoint
 from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
@@ -26,6 +28,7 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.glossary import Glossary
 from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
+from metadata.generated.schema.entity.data.metric import Metric
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
@@ -33,6 +36,7 @@ from metadata.generated.schema.entity.data.searchIndex import SearchIndex
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
+from metadata.generated.schema.entity.services.apiService import ApiService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import User
@@ -42,6 +46,9 @@ logger = utils_logger()
 T = TypeVar("T", bound=BaseModel)
 
 ES_INDEX_MAP = {
+    ApiService.__name__: "api_service_search_index",
+    APICollection.__name__: "api_collection_search_index",
+    APIEndpoint.__name__: "api_endpoint_search_index",
     Table.__name__: "table_search_index",
     Database.__name__: "database_search_index",
     DatabaseSchema.__name__: "database_schema_search_index",
@@ -56,12 +63,13 @@ ES_INDEX_MAP = {
     Topic.__name__: "topic_search_index",
     Pipeline.__name__: "pipeline_search_index",
     Glossary.__name__: "glossary_search_index",
-    GlossaryTerm.__name__: "glossary_search_index",
+    GlossaryTerm.__name__: "glossary_term_search_index",
     MlModel.__name__: "mlmodel_search_index",
     Tag.__name__: "tag_search_index",
     Container.__name__: "container_search_index",
     Query.__name__: "query_search_index",
     ReportData.__name__: "entity_report_data_index",
+    Metric.__name__: "metric_search_index",
     "web_analytic_user_activity_report": "web_analytic_user_activity_report_data_index",
     "web_analytic_entity_view_report": "web_analytic_entity_view_report_data_index",
 }

metadata/workflow/data_quality.py

@@ -11,15 +11,10 @@
 """
 Workflow definition for the Data Quality
 """
-import traceback
 from typing import Optional
 
 from metadata.data_quality.processor.test_case_runner import TestCaseRunner
 from metadata.data_quality.source.test_suite import TestSuiteSource
-from metadata.generated.schema.entity.services.connections.serviceConnection import (
-    ServiceConnection,
-)
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.tests.testSuite import ServiceType, TestSuite
 from metadata.ingestion.api.steps import Processor, Sink
 from metadata.ingestion.ometa.utils import model_str
@@ -63,44 +58,11 @@ class TestSuiteWorkflow(IngestionWorkflow):
     def _get_test_runner_processor(self) -> Processor:
         return TestCaseRunner.create(self.config.model_dump(), self.metadata)
 
-    def _retrieve_service_connection_if_needed(self, service_type: ServiceType) -> None:
-        """Get service object from source config `entityFullyQualifiedName`"""
-        if (
-            not self.config.source.serviceConnection
-            and not self.metadata.config.forceEntityOverwriting
-        ):
-            fully_qualified_name = (
-                self.config.source.sourceConfig.config.entityFullyQualifiedName.root
-            )
-            try:
-                service_name = fqn.split(fully_qualified_name)[0]
-            except IndexError as exc:
-                logger.debug(traceback.format_exc())
-                raise IndexError(
-                    f"Could not retrieve service name from entity fully qualified name {fully_qualified_name}: {exc}"
-                )
-            try:
-                service: DatabaseService = self.metadata.get_by_name(
-                    DatabaseService, service_name
-                )
-                if not service:
-                    raise ConnectionError(
-                        f"Could not retrieve service with name `{service_name}`. "
-                        "Typically caused by the `entityFullyQualifiedName` does not exists in OpenMetadata "
-                        "or the JWT Token is invalid."
-                    )
-
-                self.config.source.serviceConnection = ServiceConnection(
-                    service.connection
-                )
-
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.error(
-                    f"Error getting service connection for service name [{service_name}]"
-                    f" using the secrets manager provider [{self.metadata.config.secretsManagerProvider}]: {exc}"
-                )
-                raise exc
+    def _retrieve_service_connection_if_needed(self, _: ServiceType) -> None:
+        """A test suite might require multiple connections (e.g., for logical test suites)
+        We'll skip this step and get the connections at runtime if they are not informed
+        in the YAML already.
+        """
 
     def _get_ingestion_pipeline_service(self) -> Optional[T]:
         """