openmetadata-ingestion 1.8.7.2__py3-none-any.whl → 1.8.8.1__py3-none-any.whl

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

@@ -17,7 +17,6 @@ from typing import Optional
 
 from databricks.sdk import WorkspaceClient
 from sqlalchemy.engine import Engine
-from sqlalchemy.exc import DatabaseError
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -35,13 +34,13 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.models import DatabricksTable
 from metadata.ingestion.source.database.unitycatalog.queries import (
     UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
     UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
     UNITY_CATALOG_GET_ALL_TABLE_TAGS,
     UNITY_CATALOG_GET_CATALOGS_TAGS,
-    UNITY_CATALOG_SQL_STATEMENT_TEST,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.db_utils import get_host_from_host_port
@@ -94,33 +93,20 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
+    client = UnityCatalogClient(service_connection)
     table_obj = DatabricksTable()
-    engine = get_sqlalchemy_connection(service_connection)
-
-    def test_database_query(engine: Engine, statement: str):
-        """
-        Method used to execute the given query and fetch a result
-        to test if user has access to the tables specified
-        in the sql statement
-        """
-        try:
-            connection = engine.connect()
-            connection.execute(statement).fetchone()
-        except DatabaseError as soe:
-            logger.debug(f"Failed to fetch catalogs due to: {soe}")
 
     def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
         for catalog in connection.catalogs.list():
-            table_obj.catalog_name = catalog.name
-            break
+            if catalog.name != "__databricks_internal":
+                table_obj.catalog_name = catalog.name
+                return
 
     def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
-        for catalog in connection.catalogs.list():
-            for schema in connection.schemas.list(catalog_name=catalog.name):
-                if schema.name:
-                    table_obj.schema_name = schema.name
-                    table_obj.catalog_name = catalog.name
-                    return
+        for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
+            if schema.name:
+                table_obj.schema_name = schema.name
+                return
 
     def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
         if table_obj.catalog_name and table_obj.schema_name:
@@ -162,11 +148,7 @@ def test_connection(
         "GetSchemas": partial(get_schemas, connection, table_obj),
         "GetTables": partial(get_tables, connection, table_obj),
         "GetViews": partial(get_tables, connection, table_obj),
-        "GetQueries": partial(
-            test_database_query,
-            engine=engine,
-            statement=UNITY_CATALOG_SQL_STATEMENT_TEST,
-        ),
+        "GetQueries": client.test_query_api_access,
         "GetTags": partial(get_tags, service_connection, table_obj),
     }
 

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

@@ -92,7 +92,7 @@ logger = ingestion_logger()
 UNITY_CATALOG_TAG = "UNITY CATALOG TAG"
 UNITY_CATALOG_TAG_CLASSIFICATION = "UNITY CATALOG TAG CLASSIFICATION"
 
-
+# pylint: disable=protected-access
 class UnitycatalogSource(
     ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
 ):
@@ -306,6 +306,8 @@ class UnitycatalogSource(
                 if table.table_type:
                     if table.table_type.value.lower() == TableType.View.value.lower():
                         table_type: TableType = TableType.View
+                    if table.table_type.value.lower() == "materialized_view":
+                        table_type: TableType = TableType.MaterializedView
                     elif (
                         table.table_type.value.lower()
                         == TableType.External.value.lower()
@@ -427,7 +429,6 @@ class UnitycatalogSource(
             )
             if referred_table_fqn:
                 for parent_column in column.parent_columns:
-                    # pylint: disable=protected-access
                     col_fqn = fqn._build(referred_table_fqn, parent_column, quote=False)
                     if col_fqn:
                         referred_column_fqns.append(FullyQualifiedEntityName(col_fqn))
@@ -563,7 +564,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,
@@ -617,7 +618,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,

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

@@ -1,19 +1,7 @@
-#  Copyright 2025 Collate
-#  Licensed under the Collate Community License, Version 1.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
-#  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.
 """
-SQL Queries used during ingestion
+SQL queries for Unity Catalog
 """
 
-import textwrap
-
 UNITY_CATALOG_GET_CATALOGS_TAGS = """
 SELECT * FROM `{database}`.information_schema.catalog_tags;
 """
@@ -29,27 +17,3 @@ SELECT * FROM `{database}`.information_schema.table_tags WHERE schema_name = '{s
 UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS = """
 SELECT * FROM `{database}`.information_schema.column_tags WHERE schema_name = '{schema}';
 """
-
-UNITY_CATALOG_SQL_STATEMENT = textwrap.dedent(
-    """
-    SELECT
-      statement_type AS query_type,
-      statement_text AS query_text,
-      executed_by AS user_name,
-      start_time AS start_time,
-      null AS database_name,
-      null AS schema_name,
-      end_time AS end_time,
-      total_duration_ms/1000 AS duration
-    from system.query.history
-    WHERE statement_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
-    AND statement_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
-    AND start_time between to_timestamp('{start_time}') and to_timestamp('{end_time}')
-    {filters}
-    LIMIT {result_limit}
-    """
-)
-
-UNITY_CATALOG_SQL_STATEMENT_TEST = """
- SELECT statement_text from  system.query.history LIMIT 1
-"""

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

@@ -27,9 +27,6 @@ from metadata.ingestion.source.database.databricks.query_parser import (
 )
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
-from metadata.ingestion.source.database.unitycatalog.connection import (
-    get_sqlalchemy_connection,
-)
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -58,7 +55,6 @@ class UnityCatalogQueryParserSource(
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         self._init_super(config=config, metadata=metadata)
         self.client = UnityCatalogClient(self.service_connection)
-        self.sql_client = get_sqlalchemy_connection(self.service_connection)
 
     @classmethod
     def create(

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

@@ -11,10 +11,12 @@
 """
 unity catalog usage module
 """
+import traceback
+from datetime import datetime
+from typing import Iterable
 
-from metadata.ingestion.source.database.unitycatalog.queries import (
-    UNITY_CATALOG_SQL_STATEMENT,
-)
+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,
 )
@@ -33,24 +35,36 @@ class UnitycatalogUsageSource(UnityCatalogQueryParserSource, UsageSource):
     the same API for fetching Usage Queries
     """
 
-    sql_stmt = UNITY_CATALOG_SQL_STATEMENT
-
-    filters = """
-        AND statement_type NOT IN ('SHOW', 'DESCRIBE', 'USE')
-    """
-
-    def get_engine(self):
-        yield self.sql_client
-
-    def get_sql_statement(self, start_time, end_time):
+    def yield_table_queries(self) -> Iterable[TableQuery]:
         """
-        returns sql statement to fetch query logs.
-
-        Override if we have specific parameters
+        Method to yield TableQueries
         """
-        return self.sql_stmt.format(
-            start_time=start_time,
-            end_time=end_time,
-            filters=self.get_filters(),
-            result_limit=self.source_config.resultLimit,
+        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/profiler/interface/sqlalchemy/databricks/profiler_interface.py

@@ -23,6 +23,7 @@ from metadata.generated.schema.entity.data.table import (
     ColumnName,
     DataType,
     SystemProfile,
+    TableType,
 )
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
@@ -51,6 +52,11 @@ class DatabricksProfilerInterface(SQAProfilerInterface):
         *args,
         **kwargs,
     ) -> List[SystemProfile]:
+        if self.table_entity.tableType in (TableType.View, TableType.MaterializedView):
+            logger.debug(
+                f"Skipping {metrics.name()} metric for view {runner.table_name}"
+            )
+            return []
         logger.debug(f"Computing {metrics.name()} metric for {runner.table_name}")
         self.system_metrics_class = cast(
             Type[DatabricksSystemMetricsComputer], self.system_metrics_class

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -16,6 +16,7 @@ supporting sqlalchemy abstraction layer
 """
 
 import concurrent.futures
+import gc
 import math
 import threading
 import time
@@ -26,7 +27,7 @@ from typing import Any, Dict, List, Optional, Type, Union
 
 from sqlalchemy import Column, inspect, text
 from sqlalchemy.exc import DBAPIError, ProgrammingError, ResourceClosedError
-from sqlalchemy.orm import Session, scoped_session
+from sqlalchemy.orm import scoped_session
 
 from metadata.generated.schema.entity.data.table import (
     CustomMetricProfile,
@@ -461,6 +462,9 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                     logger.error(error)
                     self.status.failed_profiler(error, traceback.format_exc())
                     break
+                finally:
+                    # Force garbage collection to help with memory management
+                    gc.collect()
 
         # If we've exhausted all retries without success, return a tuple of None values
         return None, None, None

metadata/readers/dataframe/base.py

@@ -22,6 +22,8 @@ from metadata.readers.file.config_source_factory import get_reader
 from metadata.readers.models import ConfigSource
 from metadata.utils.logger import ingestion_logger
 
+MAX_FILE_SIZE_FOR_PREVIEW = 50 * 1024 * 1024  # 50MB
+
 logger = ingestion_logger()
 
 

metadata/readers/dataframe/parquet.py

@@ -14,6 +14,8 @@ Generic Delimiter-Separated-Values implementation
 """
 from functools import singledispatchmethod
 
+from pyarrow.parquet import ParquetFile
+
 from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
     AzureConfig,
 )
@@ -26,11 +28,18 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     LocalConfig,
 )
-from metadata.readers.dataframe.base import DataFrameReader, FileFormatException
+from metadata.readers.dataframe.base import (
+    MAX_FILE_SIZE_FOR_PREVIEW,
+    DataFrameReader,
+    FileFormatException,
+)
 from metadata.readers.dataframe.common import dataframe_to_chunks
 from metadata.readers.dataframe.models import DatalakeColumnWrapper
 from metadata.readers.file.adls import AZURE_PATH, return_azure_storage_options
 from metadata.readers.models import ConfigSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class ParquetDataFrameReader(DataFrameReader):
@@ -39,6 +48,90 @@ class ParquetDataFrameReader(DataFrameReader):
     from any source based on its init client.
     """
 
+    def _read_parquet_in_batches(
+        self, parquet_file: ParquetFile, batch_size: int = 10000
+    ):
+        """
+        Read a large parquet file in batches to avoid memory issues.
+        Includes multiple fallback strategies for older PyArrow versions.
+
+        Args:
+            parquet_file: PyArrow ParquetFile or similar object
+            batch_size: Number of rows to read per batch
+
+        Returns:
+            List of DataFrame chunks
+        """
+        chunks = []
+        batch_count = 0
+
+        try:
+            # Method 1: iter_batches (PyArrow >= 3.0 - preferred)
+            if hasattr(parquet_file, "iter_batches"):
+                logger.info(
+                    "Reading large parquet file in batches to avoid memory issues"
+                )
+                for batch in parquet_file.iter_batches(batch_size=batch_size):
+                    df_batch = batch.to_pandas()
+                    if not df_batch.empty:
+                        chunks.extend(dataframe_to_chunks(df_batch))
+                        batch_count += 1
+
+                logger.info(
+                    f"Successfully processed {batch_count} batches from large parquet file"
+                )
+                return chunks
+
+            # Method 2: Row group reading (PyArrow >= 0.15.0)
+            elif hasattr(parquet_file, "num_row_groups") and hasattr(
+                parquet_file, "read_row_group"
+            ):
+                logger.warning(
+                    "iter_batches not available, using row group reading as fallback"
+                )
+
+                for i in range(parquet_file.num_row_groups):
+                    try:
+                        row_group_table = parquet_file.read_row_group(i)
+                        df_chunk = row_group_table.to_pandas()
+                        if not df_chunk.empty:
+                            # Further chunk if row group is still too large
+                            if len(df_chunk) > batch_size:
+                                chunks.extend(dataframe_to_chunks(df_chunk))
+                            else:
+                                chunks.append(df_chunk)
+                            batch_count += 1
+                    except Exception as row_exc:
+                        logger.warning(f"Failed to read row group {i}: {row_exc}")
+                        continue
+
+                if chunks:
+                    logger.info(
+                        f"Successfully processed {batch_count} row groups from large parquet file"
+                    )
+                    return chunks
+
+            # Method 3: Regular reading (final fallback)
+            logger.warning(
+                "No chunking methods available, falling back to regular reading"
+            )
+            df = parquet_file.read().to_pandas()
+            chunks.extend(dataframe_to_chunks(df))
+
+        except Exception as exc:
+            # If all chunking fails, try regular reading as final fallback
+            logger.warning(
+                f"Batched reading failed: {exc}. Falling back to regular reading - this may cause memory issues for large files"
+            )
+            try:
+                df = parquet_file.read().to_pandas()
+                chunks.extend(dataframe_to_chunks(df))
+            except Exception as fallback_exc:
+                logger.error(f"Failed to read parquet file: {fallback_exc}")
+                raise fallback_exc
+
+        return chunks
+
     @singledispatchmethod
     def _read_parquet_dispatch(
         self, config_source: ConfigSource, key: str, bucket_name: str
@@ -48,24 +141,46 @@ class ParquetDataFrameReader(DataFrameReader):
     @_read_parquet_dispatch.register
     def _(self, _: GCSConfig, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         """
-        Read the CSV file from the gcs bucket and return a dataframe
+        Read the Parquet file from the gcs bucket and return a dataframe
         """
         # pylint: disable=import-outside-toplevel
         from gcsfs import GCSFileSystem
-        from pyarrow.parquet import ParquetFile
 
         gcs = GCSFileSystem()
-        file = gcs.open(f"gs://{bucket_name}/{key}")
-        dataframe_response = (
-            ParquetFile(file).read().to_pandas(split_blocks=True, self_destruct=True)
-        )
-        return dataframe_to_chunks(dataframe_response)
+        file_path = f"gs://{bucket_name}/{key}"
+
+        # Check file size to determine reading strategy
+        try:
+            file_info = gcs.info(file_path)
+            file_size = file_info.get("size", 0)
+
+            file = gcs.open(file_path)
+            parquet_file = ParquetFile(file)
+
+            if self._should_use_chunking(file_size):
+                # Use batched reading for large files
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use regular reading for smaller files
+                dataframe_response = parquet_file.read().to_pandas(
+                    split_blocks=True, self_destruct=True
+                )
+                return dataframe_to_chunks(dataframe_response)
+
+        except Exception:
+            # Fallback to regular reading if size check fails
+            file = gcs.open(file_path)
+            parquet_file = ParquetFile(file)
+            dataframe_response = parquet_file.read().to_pandas(
+                split_blocks=True, self_destruct=True
+            )
+            return dataframe_to_chunks(dataframe_response)
 
     @_read_parquet_dispatch.register
     def _(self, _: S3Config, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         # pylint: disable=import-outside-toplevel
         from pyarrow.fs import S3FileSystem
-        from pyarrow.parquet import ParquetDataset
+        from pyarrow.parquet import ParquetDataset, ParquetFile
 
         client_kwargs = {
             "endpoint_override": (
@@ -90,13 +205,40 @@ class ParquetDataFrameReader(DataFrameReader):
         s3_fs = S3FileSystem(**client_kwargs)
 
         bucket_uri = f"{bucket_name}/{key}"
-        dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
 
-        return dataframe_to_chunks(dataset.read_pandas().to_pandas())
+        # Check file size to determine reading strategy
+        try:
+            file_info = s3_fs.get_file_info(bucket_uri)
+            file_size = file_info.size if hasattr(file_info, "size") else 0
+
+            if self._should_use_chunking(file_size):
+                # Use ParquetFile for batched reading of large files
+                logger.info(
+                    f"Large parquet file detected ({file_size} bytes > {MAX_FILE_SIZE_FOR_PREVIEW} bytes). "
+                    f"Using batched reading for file: {bucket_uri}"
+                )
+                parquet_file = ParquetFile(bucket_uri, filesystem=s3_fs)
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use ParquetDataset for regular reading of smaller files
+                logger.debug(
+                    f"Reading small parquet file ({file_size} bytes): {bucket_uri}"
+                )
+                dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
+                return dataframe_to_chunks(dataset.read_pandas().to_pandas())
+
+        except Exception as exc:
+            # Fallback to regular reading if size check fails
+            logger.warning(
+                f"Could not determine file size for {bucket_uri}: {exc}. Using regular reading"
+            )
+            dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
+            return dataframe_to_chunks(dataset.read_pandas().to_pandas())
 
     @_read_parquet_dispatch.register
     def _(self, _: AzureConfig, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         import pandas as pd  # pylint: disable=import-outside-toplevel
+        import pyarrow.fs as fs
 
         storage_options = return_azure_storage_options(self.config_source)
         account_url = AZURE_PATH.format(
@@ -104,8 +246,33 @@ class ParquetDataFrameReader(DataFrameReader):
             account_name=self.config_source.securityConfig.accountName,
             key=key,
         )
-        dataframe = pd.read_parquet(account_url, storage_options=storage_options)
-        return dataframe_to_chunks(dataframe)
+
+        # Check file size to determine reading strategy
+        try:
+            # Try to get file size from Azure
+            azure_fs = fs.SubTreeFileSystem(
+                account_url, fs.AzureFileSystem(**storage_options)
+            )
+            file_info = azure_fs.get_file_info("/")
+            file_size = file_info.size if hasattr(file_info, "size") else 0
+
+            if self._should_use_chunking(file_size):
+                # Use PyArrow ParquetFile for batched reading of large files
+                parquet_file = ParquetFile(
+                    account_url, filesystem=fs.AzureFileSystem(**storage_options)
+                )
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use pandas for regular reading of smaller files
+                dataframe = pd.read_parquet(
+                    account_url, storage_options=storage_options
+                )
+                return dataframe_to_chunks(dataframe)
+
+        except Exception:
+            # Fallback to regular pandas reading if size check or batching fails
+            dataframe = pd.read_parquet(account_url, storage_options=storage_options)
+            return dataframe_to_chunks(dataframe)
 
     @_read_parquet_dispatch.register
     def _(
@@ -114,10 +281,27 @@ class ParquetDataFrameReader(DataFrameReader):
         key: str,
         bucket_name: str,  # pylint: disable=unused-argument
     ) -> DatalakeColumnWrapper:
+        import os
+
         import pandas as pd  # pylint: disable=import-outside-toplevel
 
-        dataframe = pd.read_parquet(key)
-        return dataframe_to_chunks(dataframe)
+        # Check file size to determine reading strategy
+        try:
+            file_size = os.path.getsize(key)
+
+            if self._should_use_chunking(file_size):
+                # Use PyArrow ParquetFile for batched reading of large files
+                parquet_file = ParquetFile(key)
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use pandas for regular reading of smaller files
+                dataframe = pd.read_parquet(key)
+                return dataframe_to_chunks(dataframe)
+
+        except Exception:
+            # Fallback to regular pandas reading if size check fails
+            dataframe = pd.read_parquet(key)
+            return dataframe_to_chunks(dataframe)
 
     def _read(self, *, key: str, bucket_name: str, **__) -> DatalakeColumnWrapper:
         return DatalakeColumnWrapper(
@@ -125,3 +309,6 @@ class ParquetDataFrameReader(DataFrameReader):
                 self.config_source, key=key, bucket_name=bucket_name
             )
         )
+
+    def _should_use_chunking(self, file_size: int) -> bool:
+        return file_size > MAX_FILE_SIZE_FOR_PREVIEW or file_size == 0