openmetadata-ingestion 1.6.3.2.dev1__py3-none-any.whl → 1.6.5.0__py3-none-any.whl

metadata/ingestion/source/database/unitycatalog/usage.py

@@ -11,17 +11,22 @@
 """
 unity catalog usage module
 """
+import traceback
+from datetime import datetime
+from typing import Iterable
 
-from metadata.ingestion.source.database.databricks.usage import DatabricksUsageSource
+from metadata.generated.schema.type.basic import DateTime
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.database.unitycatalog.query_parser import (
     UnityCatalogQueryParserSource,
 )
+from metadata.ingestion.source.database.usage_source import UsageSource
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSource):
+class UnitycatalogUsageSource(UnityCatalogQueryParserSource, UsageSource):
     """
     UnityCatalog Usage Source
 
@@ -29,3 +34,37 @@ class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSour
     DatabricksUsageSource as both the sources would call
     the same API for fetching Usage Queries
     """
+
+    def yield_table_queries(self) -> Iterable[TableQuery]:
+        """
+        Method to yield TableQueries
+        """
+        queries = []
+        data = self.client.list_query_history(
+            start_date=self.start,
+            end_date=self.end,
+        )
+        for row in data or []:
+            try:
+                if self.client.is_query_valid(row):
+                    queries.append(
+                        TableQuery(
+                            dialect=self.dialect.value,
+                            query=row.get("query_text"),
+                            userName=row.get("user_name"),
+                            startTime=str(row.get("query_start_time_ms")),
+                            endTime=str(row.get("execution_end_time_ms")),
+                            analysisDate=DateTime(datetime.now()),
+                            serviceName=self.config.serviceName,
+                            duration=row.get("duration")
+                            if row.get("duration")
+                            else None,
+                        )
+                    )
+            except Exception as err:
+                logger.debug(traceback.format_exc())
+                logger.warning(
+                    f"Failed to process query {row.get('query_text')} due to: {err}"
+                )
+
+        yield TableQueries(queries=queries)

metadata/ingestion/source/database/usage_source.py

@@ -121,6 +121,7 @@ class UsageSource(QueryParserSource, ABC):
                         for row in rows:
                             row = dict(row)
                             try:
+                                row.update({k.lower(): v for k, v in row.items()})
                                 logger.debug(f"Processing row: {query}")
                                 query_type = row.get("query_type")
                                 query = self.format_query(row["query_text"])
@@ -153,7 +154,7 @@ class UsageSource(QueryParserSource, ABC):
                 if query:
                     logger.debug(
                         (
-                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}"
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
                             "\n##########################"
                         )
                     )

metadata/ingestion/source/metadata/atlas/client.py

@@ -18,6 +18,7 @@ from metadata.generated.schema.entity.services.connections.metadata.atlasConnect
     AtlasConnection,
 )
 from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
+from metadata.utils.helpers import clean_uri
 
 
 class AtlasClient:
@@ -31,7 +32,7 @@ class AtlasClient:
             config.username, config.password.get_secret_value()
         )
         client_config: ClientConfig = ClientConfig(
-            base_url=str(config.hostPort),
+            base_url=clean_uri(config.hostPort),
             auth_header="Authorization",
             api_version="api",
             auth_token=self.get_auth_token,

metadata/ingestion/source/pipeline/nifi/connection.py

@@ -17,8 +17,10 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.pipeline.nifi.basicAuth import (
+    NifiBasicAuth,
+)
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
-    BasicAuthentication,
     NifiConnection,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
@@ -34,7 +36,7 @@ def get_connection(connection: NifiConnection) -> NifiClient:
     """
     Create connection
     """
-    if isinstance(connection.nifiConfig, BasicAuthentication):
+    if isinstance(connection.nifiConfig, NifiBasicAuth):
         return NifiClient(
             host_port=connection.hostPort,
             username=connection.nifiConfig.username,

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -24,6 +24,12 @@ from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
     SnowflakeType,
 )
@@ -79,21 +85,26 @@ class SQAInterfaceMixin(Root):
             )
 
     def set_catalog(self, session) -> None:
-        """Set catalog for the session. Right now only databricks and unity catalog requires it
+        """Set the catalog or database for the session.
 
         Args:
             session (Session): sqa session object
         """
-        if not isinstance(
+        if isinstance(
             self.service_connection_config,
             (UnityCatalogConnection, DatabricksConnection),
         ):
-            return
-        bind = session.get_bind()
-        bind.execute(
-            "USE CATALOG %(catalog)s;",
-            {"catalog": self.service_connection_config.catalog},
-        ).first()
+            session.get_bind().execute(
+                "USE CATALOG %(catalog)s;",
+                {"catalog": self.service_connection_config.catalog},
+            ).first()
+
+        if isinstance(
+            self.service_connection_config, (MysqlConnection, MariaDBConnection)
+        ):
+            session.get_bind().execute(
+                f"USE {self.table_entity.databaseSchema.name};",
+            )
 
     def close(self):
         """close session"""

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -297,7 +297,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
             column: the column to compute the metrics against
             metrics: list of metrics to compute
         Returns:
-            dictionnary of results
+            dictionary of results
         """
 
         if not metrics:

metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.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.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+from metadata.ingestion.source.database.databricks.connection import (
+    get_connection as databricks_get_connection,
+)
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+
+
+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

metadata/profiler/metrics/window/first_quartile.py

@@ -88,7 +88,7 @@ class FirstQuartile(StaticMetric, PercentilMixin):
             except MemoryError:
                 logger.error(
                     f"Unable to compute Median for {self.col.name} due to memory constraints."
-                    f"We recommend using a smaller sample size or partitionning."
+                    f"We recommend using a smaller sample size or partitioning."
                 )
                 return None
             # check if nan

metadata/profiler/metrics/window/median.py

@@ -87,7 +87,7 @@ class Median(StaticMetric, PercentilMixin):
             except MemoryError:
                 logger.error(
                     f"Unable to compute Median for {self.col.name} due to memory constraints."
-                    f"We recommend using a smaller sample size or partitionning."
+                    f"We recommend using a smaller sample size or partitioning."
                 )
                 return None
             try:

metadata/profiler/metrics/window/third_quartile.py

@@ -88,7 +88,7 @@ class ThirdQuartile(StaticMetric, PercentilMixin):
             except MemoryError:
                 logger.error(
                     f"Unable to compute Median for {self.col.name} due to memory constraints."
-                    f"We recommend using a smaller sample size or partitionning."
+                    f"We recommend using a smaller sample size or partitioning."
                 )
                 return None
             # check if nan

metadata/profiler/source/fetcher/config.py

@@ -42,3 +42,7 @@ class EntityFilterConfigInterface(Protocol):
     @property
     def useFqnForFiltering(self) -> Optional[bool]:
         ...
+
+    @property
+    def includeViews(self) -> Optional[bool]:
+        ...

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -17,6 +17,7 @@ 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.data.table import TableType
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -115,13 +116,13 @@ class DatabaseFetcherStrategy(FetcherStrategy):
         super().__init__(config, metadata, global_profiler_config, status)
         self.source_config = cast(
             EntityFilterConfigInterface, self.source_config
-        )  # Satisfy typchecker
+        )  # Satisfy typechecker
 
     def _filter_databases(self, databases: Iterable[Database]) -> Iterable[Database]:
         """Filter databases based on the filter pattern
 
         Args:
-            database (Database): Database to filter
+            databases (Database): Database to filter
 
         Returns:
             bool
@@ -192,6 +193,21 @@ class DatabaseFetcherStrategy(FetcherStrategy):
 
         return False
 
+    def _filter_views(self, table: Table) -> bool:
+        """Filter the tables based on include views configuration"""
+        # If we include views, nothing to filter
+        if self.source_config.includeViews:
+            return False
+
+        # Otherwise, filter out views
+        if table.tableType == TableType.View:
+            self.status.filter(
+                table.name.root, f"We are not including views {table.name.root}"
+            )
+            return True
+
+        return False
+
     def _filter_column_metrics_computation(self):
         """Filter"""
 
@@ -242,6 +258,7 @@ class DatabaseFetcherStrategy(FetcherStrategy):
                 not self.source_config.classificationFilterPattern
                 or not self.filter_classifications(table)
             )
+            and not self._filter_views(table)
         ]
 
         return tables

metadata/sampler/sampler_interface.py

@@ -76,9 +76,6 @@ class SamplerInterface(ABC):
         self._columns: Optional[List[SQALikeColumn]] = None
         self.sample_config = sample_config
 
-        if not self.sample_config.profileSample:
-            self.sample_config.profileSample = 100
-
         self.entity = entity
         self.include_columns = include_columns
         self.exclude_columns = exclude_columns

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -54,7 +54,6 @@ class BigQuerySampler(SQASampler):
         sample_query: Optional[str] = None,
         storage_config: DataStorageConfig = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
-        table_type: TableType = None,
         **kwargs,
     ):
         super().__init__(
@@ -68,7 +67,7 @@ class BigQuerySampler(SQASampler):
             sample_data_count=sample_data_count,
             **kwargs,
         )
-        self.raw_dataset_type: TableType = table_type
+        self.raw_dataset_type: Optional[TableType] = entity.tableType
 
     def set_tablesample(self, selectable: SqaTable):
         """Set the TABLESAMPLE clause for BigQuery
@@ -120,7 +119,7 @@ class BigQuerySampler(SQASampler):
             and self.raw_dataset_type != TableType.View
         ):
             return self._base_sample_query(column).cte(
-                f"{self.raw_dataset.__tablename__}_sample"
+                f"{self.get_sampler_table_name()}_sample"
             )
 
         return super().get_sample_query(column=column)

metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -46,7 +46,7 @@ class MssqlSampler(SQASampler):
     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"
+            f"{self.get_sampler_table_name()}_rnd"
         )
         query = self.client.query(rnd)
-        return query.cte(f"{self.raw_dataset.__tablename__}_sample")
+        return query.cte(f"{self.get_sampler_table_name()}_sample")

metadata/sampler/sqlalchemy/postgres/sampler.py

@@ -85,7 +85,7 @@ class PostgresSampler(SQASampler):
     def get_sample_query(self, *, column=None) -> Query:
         if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
             return self._base_sample_query(column).cte(
-                f"{self.raw_dataset.__tablename__}_rnd"
+                f"{self.get_sampler_table_name()}_rnd"
             )
 
         return super().get_sample_query(column=column)

metadata/sampler/sqlalchemy/sampler.py

@@ -12,6 +12,7 @@
 Helper module to handle data sampling
 for the profiler
 """
+import hashlib
 import traceback
 from typing import List, Optional, Union, cast
 
@@ -32,6 +33,7 @@ from metadata.profiler.orm.functions.modulo import ModuloFn
 from metadata.profiler.orm.functions.random_num import RandomNumFn
 from metadata.profiler.processor.handle_partition import build_partition_predicate
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.constants import UTF_8
 from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
 
@@ -109,17 +111,28 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             query = self.get_partitioned_query(query)
         return query
 
+    def get_sampler_table_name(self) -> str:
+        """Get the base name of the SQA table for sampling.
+        We use MD5 as a hashing algorithm to generate a unique name for the table
+        keeping its length controlled. Otherwise, we ended up having issues
+        with names getting truncated when we add the suffixes to the identifiers
+        such as _sample, or _rnd.
+        """
+        encoded_name = self.raw_dataset.__tablename__.encode(UTF_8)
+        hash_object = hashlib.md5(encoded_name)
+        return hash_object.hexdigest()
+
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
         if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
             rnd = self._base_sample_query(
                 column,
                 (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
-            ).cte(f"{self.raw_dataset.__tablename__}_rnd")
+            ).cte(f"{self.get_sampler_table_name()}_rnd")
             session_query = self.client.query(rnd)
             return session_query.where(
                 rnd.c.random <= self.sample_config.profileSample
-            ).cte(f"{self.raw_dataset.__tablename__}_sample")
+            ).cte(f"{self.get_sampler_table_name()}_sample")
 
         table_query = self.client.query(self.raw_dataset)
         session_query = self._base_sample_query(
@@ -129,7 +142,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         return (
             session_query.order_by(RANDOM_LABEL)
             .limit(self.sample_config.profileSample)
-            .cte(f"{self.raw_dataset.__tablename__}_rnd")
+            .cte(f"{self.get_sampler_table_name()}_rnd")
         )
 
     def get_dataset(self, column=None, **__) -> Union[DeclarativeMeta, AliasedClass]:
@@ -143,7 +156,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not self.sample_config.profileSample:
             if self.partition_details:
                 partitioned = self._partitioned_table()
-                return partitioned.cte(f"{self.raw_dataset.__tablename__}_partitioned")
+                return partitioned.cte(f"{self.get_sampler_table_name()}_partitioned")
 
             return self.raw_dataset
 
@@ -162,23 +175,23 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             return self._fetch_sample_data_from_user_query()
 
         # Add new RandomNumFn column
-        rnd = self.get_sample_query()
+        ds = self.get_dataset()
         if not columns:
-            sqa_columns = [col for col in inspect(rnd).c if col.name != RANDOM_LABEL]
+            sqa_columns = [col for col in inspect(ds).c if col.name != RANDOM_LABEL]
         else:
             # we can't directly use columns as it is bound to self.raw_dataset and not the rnd table.
             # If we use it, it will result in a cross join between self.raw_dataset and rnd table
             names = [col.name for col in columns]
             sqa_columns = [
                 col
-                for col in inspect(rnd).c
+                for col in inspect(ds).c
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
         try:
             sqa_sample = (
                 self.client.query(*sqa_columns)
-                .select_from(rnd)
+                .select_from(ds)
                 .limit(self.sample_limit)
                 .all()
             )
@@ -224,7 +237,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
 
         return self.client.query(stmt.subquery()).cte(
-            f"{self.raw_dataset.__tablename__}_user_sampled"
+            f"{self.get_sampler_table_name()}_user_sampled"
         )
 
     def _partitioned_table(self) -> Query:

metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -89,7 +89,7 @@ class SnowflakeSampler(SQASampler):
     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"
+            f"{self.get_sampler_table_name()}_rnd"
         )
         query = self.client.query(rnd)
-        return query.cte(f"{self.raw_dataset.__tablename__}_sample")
+        return query.cte(f"{self.get_sampler_table_name()}_sample")