openmetadata-ingestion 1.8.7.1__py3-none-any.whl → 1.8.9.0__py3-none-any.whl

metadata/ingestion/source/security/security_service.py

@@ -0,0 +1,128 @@
+#  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.
+"""
+Base class for ingesting security services
+"""
+from abc import ABC
+from typing import Set
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from metadata.generated.schema.entity.services.securityService import (
+    SecurityConnection,
+    SecurityService,
+)
+from metadata.generated.schema.metadataIngestion.securityServiceMetadataPipeline import (
+    SecurityServiceMetadataPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.steps import Source
+from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
+from metadata.ingestion.models.topology import (
+    NodeStage,
+    ServiceTopology,
+    TopologyContextManager,
+    TopologyNode,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.connections import get_connection
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class SecurityServiceTopology(ServiceTopology):
+    """
+    Defines the hierarchy in Security Services.
+
+    We could have a topology validator. We can only consume
+    data that has been produced by any parent node.
+    """
+
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
+        producer="get_services",
+        stages=[
+            NodeStage(
+                type_=SecurityService,
+                context="security_service",
+                processor="yield_create_request_security_service",
+                overwrite=False,
+                must_return=True,
+                cache_entities=True,
+            ),
+        ],
+        children=[],  # Security services typically don't have child entities like policies, roles, etc.
+        post_process=["mark_security_entities_as_deleted"],
+    )
+
+
+from metadata.utils.helpers import clean_uri
+
+
+class SecurityServiceSource(TopologyRunnerMixin, Source, ABC):
+    """
+    Base class for Security Services.
+    It implements the topology and context.
+    """
+
+    source_config: SecurityServiceMetadataPipeline
+    config: WorkflowSource
+    # Big union of types we want to fetch dynamically
+    service_connection: SecurityConnection.model_fields["config"].annotation
+
+    topology = SecurityServiceTopology()
+    context = TopologyContextManager(topology)
+    security_source_state: Set = set()
+
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
+        config.serviceConnection.root.config.hostPort = clean_uri(
+            config.serviceConnection.root.config.hostPort
+        )
+        super().__init__()
+        self.config = config
+        self.metadata = metadata
+        self.service_connection = self.config.serviceConnection.root.config
+        self.source_config: SecurityServiceMetadataPipeline = (
+            self.config.sourceConfig.config
+        )
+
+        self.connection = get_connection(self.service_connection)
+        # Flag the connection for the test connection
+        self.connection_obj = self.connection
+        self.client = self.get_client()
+        self.test_connection()
+
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
+    def close(self):
+        pass
+
+    def yield_create_request_security_service(self, config: WorkflowSource):
+        yield Either(
+            right=self.metadata.get_create_service_from_source(
+                entity=SecurityService, config=config
+            )
+        )
+
+    def test_connection(self) -> None:
+        self.client.test_connection()

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/dsv.py

@@ -54,16 +54,24 @@ class DSVDataFrameReader(DataFrameReader):
         super().__init__(config_source, client)
 
     def read_from_pandas(
-        self, path: str, storage_options: Optional[Dict[str, Any]] = None
+        self,
+        path: str,
+        storage_options: Optional[Dict[str, Any]] = None,
+        compression: Optional[str] = None,
     ) -> DatalakeColumnWrapper:
         import pandas as pd  # pylint: disable=import-outside-toplevel
 
+        # Determine compression based on file extension if not provided
+        if compression is None and path.endswith(".gz"):
+            compression = "gzip"
+
         chunk_list = []
         with pd.read_csv(
             path,
             sep=self.separator,
             chunksize=CHUNKSIZE,
             storage_options=storage_options,
+            compression=compression,
         ) as reader:
             for chunks in reader:
                 chunk_list.append(chunks)
@@ -81,16 +89,47 @@ class DSVDataFrameReader(DataFrameReader):
         """
         Read the CSV file from the gcs bucket and return a dataframe
         """
+        # Determine compression based on file extension
+        compression = None
+        if key.endswith(".gz"):
+            compression = "gzip"
+
         path = f"gs://{bucket_name}/{key}"
-        return self.read_from_pandas(path=path)
+        return self.read_from_pandas(path=path, compression=compression)
 
     @_read_dsv_dispatch.register
     def _(self, _: S3Config, key: str, bucket_name: str) -> DatalakeColumnWrapper:
-        path = self.client.get_object(Bucket=bucket_name, Key=key)["Body"]
-        return self.read_from_pandas(path=path)
+        import pandas as pd  # pylint: disable=import-outside-toplevel
+
+        # Determine compression based on file extension
+        compression = None
+        if key.endswith(".gz"):
+            compression = "gzip"
+
+        # Get the file content from S3
+        response = self.client.get_object(Bucket=bucket_name, Key=key)
+        file_content = response["Body"]
+
+        # Read the CSV data directly from the StreamingBody
+        chunk_list = []
+        with pd.read_csv(
+            file_content,
+            sep=self.separator,
+            chunksize=CHUNKSIZE,
+            compression=compression,
+        ) as reader:
+            for chunks in reader:
+                chunk_list.append(chunks)
+
+        return DatalakeColumnWrapper(dataframes=chunk_list)
 
     @_read_dsv_dispatch.register
     def _(self, _: AzureConfig, key: str, bucket_name: str) -> DatalakeColumnWrapper:
+        # Determine compression based on file extension
+        compression = None
+        if key.endswith(".gz"):
+            compression = "gzip"
+
         storage_options = return_azure_storage_options(self.config_source)
         path = AZURE_PATH.format(
             bucket_name=bucket_name,
@@ -100,13 +139,19 @@ class DSVDataFrameReader(DataFrameReader):
         return self.read_from_pandas(
             path=path,
             storage_options=storage_options,
+            compression=compression,
         )
 
     @_read_dsv_dispatch.register
     def _(  # pylint: disable=unused-argument
         self, _: LocalConfig, key: str, bucket_name: str
     ) -> DatalakeColumnWrapper:
-        return self.read_from_pandas(path=key)
+        # Determine compression based on file extension
+        compression = None
+        if key.endswith(".gz"):
+            compression = "gzip"
+
+        return self.read_from_pandas(path=key, compression=compression)
 
     def _read(self, *, key: str, bucket_name: str, **__) -> DatalakeColumnWrapper:
         return self._read_dsv_dispatch(

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

metadata/readers/dataframe/reader_factory.py

@@ -36,6 +36,7 @@ logger = utils_logger()
 
 class SupportedTypes(Enum):
     CSV = "csv"
+    CSVGZ = "csv.gz"
     TSV = "tsv"
     AVRO = "avro"
     PARQUET = "parquet"
@@ -53,6 +54,7 @@ class SupportedTypes(Enum):
 
 DF_READER_MAP = {
     SupportedTypes.CSV.value: CSVDataFrameReader,
+    SupportedTypes.CSVGZ.value: CSVDataFrameReader,
     SupportedTypes.TSV.value: TSVDataFrameReader,
     SupportedTypes.AVRO.value: AvroDataFrameReader,
     SupportedTypes.PARQUET.value: ParquetDataFrameReader,
@@ -79,7 +81,10 @@ def get_df_reader(
     Load the File Reader based on the Config Source
     """
     # If we have a DSV file, build a reader dynamically based on the received separator
-    if type_ in {SupportedTypes.CSV, SupportedTypes.TSV} and separator:
+    if (
+        type_ in {SupportedTypes.CSV, SupportedTypes.CSVGZ, SupportedTypes.TSV}
+        and separator
+    ):
         return get_dsv_reader_by_separator(separator=separator)(
             config_source=config_source, client=client
         )