openmetadata-ingestion 1.8.7.0__py3-none-any.whl → 1.8.8.0__py3-none-any.whl

metadata/ingestion/source/database/glue/metadata.py

@@ -377,12 +377,87 @@ class GlueSource(ExternalTableLineageMixin, DatabaseServiceSource):
         parsed_string["description"] = column.Comment
         return Column(**parsed_string)
 
+    # pylint: disable=too-many-locals
     def get_columns(self, column_data: StorageDetails) -> Optional[Iterable[Column]]:
+        """
+        Get columns from Glue.
+        """
+        # Check if this is an Iceberg table
+        table = self.context.get().table_data
+        is_iceberg = table.Parameters and table.Parameters.table_type == "ICEBERG"
+
+        if is_iceberg:
+            # For Iceberg tables, get the full table metadata from Glue to access column parameters
+            try:
+                schema_name = self.context.get().database_schema
+                table_name = table.Name
+
+                # Get full table metadata from Glue API
+                response = self.glue.get_table(
+                    DatabaseName=schema_name, Name=table_name
+                )
+
+                table_info = response["Table"]
+
+                # Filter out non-current Iceberg columns
+                storage_descriptor = table_info.get("StorageDescriptor", {})
+                glue_columns = storage_descriptor.get("Columns", [])
+
+                for glue_col in glue_columns:
+                    col_name = glue_col["Name"]
+                    col_type = glue_col["Type"]
+                    col_comment = glue_col.get("Comment", "")
+                    col_parameters = glue_col.get("Parameters", {})
+
+                    # Check if this is a non-current Iceberg column
+                    iceberg_current = col_parameters.get(
+                        "iceberg.field.current", "true"
+                    )
+                    is_current = iceberg_current != "false"
+
+                    if is_current:
+                        # Create a GlueColumn object for processing
+                        column_obj = GlueColumn(
+                            Name=col_name, Type=col_type, Comment=col_comment
+                        )
+                        yield self._get_column_object(column_obj)
+
+                # Process partition columns
+                partition_keys = table_info.get("PartitionKeys", [])
+                for glue_col in partition_keys:
+                    col_name = glue_col["Name"]
+                    col_type = glue_col["Type"]
+                    col_comment = glue_col.get("Comment", "")
+                    col_parameters = glue_col.get("Parameters", {})
+
+                    # Check if this is a non-current Iceberg column
+                    iceberg_current = col_parameters.get(
+                        "iceberg.field.current", "true"
+                    )
+                    is_current = iceberg_current != "false"
+
+                    if is_current:
+                        # Create a GlueColumn object for processing
+                        column_obj = GlueColumn(
+                            Name=col_name, Type=col_type, Comment=col_comment
+                        )
+                        yield self._get_column_object(column_obj)
+
+                return
+
+            except Exception as e:
+                # If we can't get Glue metadata, fall back to the original method
+                # This ensures backward compatibility
+                logger.warning(
+                    f"Failed to get Glue metadata for Iceberg table {table.Name}: {e}"
+                )
+
+        # For non-Iceberg tables or if Glue access fails, use the original method
         # process table regular columns info
         for column in column_data.Columns:
             yield self._get_column_object(column)
 
-        # process table regular columns info
+        # process table partition columns info
         for column in self.context.get().table_data.PartitionKeys:
             yield self._get_column_object(column)
 

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -187,36 +187,44 @@ class AirflowSource(PipelineServiceSource):
         """
         Return the DagRuns of given dag
         """
-        dag_run_list = (
-            self.session.query(
-                DagRun.dag_id,
-                DagRun.run_id,
-                DagRun.queued_at,
-                DagRun.execution_date,
-                DagRun.start_date,
-                DagRun.state,
+        try:
+            dag_run_list = (
+                self.session.query(
+                    DagRun.dag_id,
+                    DagRun.run_id,
+                    DagRun.queued_at,
+                    DagRun.execution_date,
+                    DagRun.start_date,
+                    DagRun.state,
+                )
+                .filter(DagRun.dag_id == dag_id)
+                .order_by(DagRun.execution_date.desc())
+                .limit(self.config.serviceConnection.root.config.numberOfStatus)
+                .all()
             )
-            .filter(DagRun.dag_id == dag_id)
-            .order_by(DagRun.execution_date.desc())
-            .limit(self.config.serviceConnection.root.config.numberOfStatus)
-            .all()
-        )
 
-        dag_run_dict = [dict(elem) for elem in dag_run_list]
-
-        # Build DagRun manually to not fall into new/old columns from
-        # different Airflow versions
-        return [
-            DagRun(
-                dag_id=elem.get("dag_id"),
-                run_id=elem.get("run_id"),
-                queued_at=elem.get("queued_at"),
-                execution_date=elem.get("execution_date"),
-                start_date=elem.get("start_date"),
-                state=elem.get("state"),
+            dag_run_dict = [dict(elem) for elem in dag_run_list]
+
+            # Build DagRun manually to not fall into new/old columns from
+            # different Airflow versions
+            return [
+                DagRun(
+                    dag_id=elem.get("dag_id"),
+                    run_id=elem.get("run_id"),
+                    queued_at=elem.get("queued_at"),
+                    execution_date=elem.get("execution_date"),
+                    start_date=elem.get("start_date"),
+                    state=elem.get("state"),
+                )
+                for elem in dag_run_dict
+            ]
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Could not get pipeline status for {dag_id}. "
+                f"This might be due to Airflow version incompatibility - {exc}"
             )
-            for elem in dag_run_dict
-        ]
+            return []
 
     def get_task_instances(
         self, dag_id: str, run_id: str, serialized_tasks: List[AirflowTask]
@@ -369,16 +377,27 @@ class AirflowSource(PipelineServiceSource):
                 break
             for serialized_dag in results:
                 try:
-                    dag_model = (
-                        self.session.query(DagModel)
-                        .filter(DagModel.dag_id == serialized_dag[0])
-                        .one_or_none()
-                    )
-                    pipeline_state = (
-                        PipelineState.Active.value
-                        if dag_model and not dag_model.is_paused
-                        else PipelineState.Inactive.value
-                    )
+                    # Query only the is_paused column from DagModel
+                    try:
+                        is_paused_result = (
+                            self.session.query(DagModel.is_paused)
+                            .filter(DagModel.dag_id == serialized_dag[0])
+                            .scalar()
+                        )
+                        pipeline_state = (
+                            PipelineState.Active.value
+                            if not is_paused_result
+                            else PipelineState.Inactive.value
+                        )
+                    except Exception as exc:
+                        logger.debug(traceback.format_exc())
+                        logger.warning(
+                            f"Could not query DagModel.is_paused for {serialized_dag[0]}. "
+                            f"Using default pipeline state - {exc}"
+                        )
+                        # If we can't query is_paused, assume the pipeline is active
+                        pipeline_state = PipelineState.Active.value
+
                     data = serialized_dag[1]["dag"]
                     dag = AirflowDagDetails(
                         dag_id=serialized_dag[0],

metadata/ingestion/source/pipeline/airflow/utils.py

@@ -23,6 +23,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+# pylint: disable=too-many-branches,too-many-return-statements,too-many-nested-blocks
 def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
     """
     Fetch Schedule Intervals from Airflow Dags
@@ -40,7 +41,49 @@ def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
 
             expression_class = timetable.get("__type")
             if expression_class:
-                return import_from_module(expression_class)().summary
+                try:
+                    # Try to instantiate the timetable class safely
+                    timetable_class = import_from_module(expression_class)
+
+                    # Handle special cases for classes that require arguments
+                    if "DatasetTriggeredTimetable" in expression_class:
+                        # DatasetTriggeredTimetable requires datasets argument
+                        # For now, return a descriptive string since we can't instantiate it properly
+                        return "Dataset Triggered"
+                    if "CronDataIntervalTimetable" in expression_class:
+                        # Handle cron-based timetables
+                        try:
+                            return timetable_class().summary
+                        except (TypeError, AttributeError):
+                            return "Cron Based"
+                    else:
+                        # Try to instantiate with no arguments
+                        try:
+                            return timetable_class().summary
+                        except (TypeError, AttributeError):
+                            # If summary attribute doesn't exist, try to get a string representation
+                            try:
+                                instance = timetable_class()
+                                return str(instance)
+                            except TypeError:
+                                # If instantiation fails, return the class name
+                                return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except ImportError as import_error:
+                    logger.debug(
+                        f"Could not import timetable class {expression_class}: {import_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except TypeError as type_error:
+                    # If instantiation fails due to missing arguments, log and continue
+                    logger.debug(
+                        f"Could not instantiate timetable class {expression_class}: {type_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except Exception as inst_error:
+                    logger.debug(
+                        f"Error instantiating timetable class {expression_class}: {inst_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
 
         if schedule:
             if isinstance(schedule, str):
@@ -57,7 +100,6 @@ def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
 
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        logger.warning(
-            f"Couldn't fetch schedule interval for dag {pipeline_data.get('_dag_id'): [{exc}]}"
-        )
+        dag_id = pipeline_data.get("_dag_id", "unknown")
+        logger.warning(f"Couldn't fetch schedule interval for dag {dag_id}: {exc}")
     return None

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