Fix: Optimise redshift memory usage (#25904)

Author: ulixius9

ingestion/src/metadata/ingestion/source/database/redshift/metadata.py

@@ -190,27 +190,47 @@ def get_location_path(self, table_name: str, schema_name: str) -> Optional[str]:
             (self.context.get().database, schema_name, table_name)
         )
 
-    def get_partition_details(self) -> None:
+    def get_partition_details(self, schema_name: Optional[str] = None) -> None:
         """
-        Populate partition details
+        Populate partition details for the given schema (or all schemas if None).
         """
         try:
             self.partition_details.clear()
-            results = self.connection.execute(
-                statement=REDSHIFT_PARTITION_DETAILS
-            ).fetchall()
+            query = REDSHIFT_PARTITION_DETAILS
+            if schema_name:
+                query += f" AND \"schema\" = '{schema_name}'"
+            results = self.connection.execute(statement=query).fetchall()
             for row in results:
                 self.partition_details[f"{row.schema}.{row.table}"] = row.diststyle
         except Exception as exe:
             logger.debug(traceback.format_exc())
             logger.debug(f"Failed to fetch partition details due: {exe}")
 
+    def _clear_reflection_cache(self) -> None:
+        """Clear the SQLAlchemy inspector's info_cache to release
+        cached column / relation data from prior schemas.
+
+        This prevents unbounded memory growth when ingesting many
+        schemas, since _get_schema_column_info, get_columns, and
+        _get_all_relation_info all use @reflection.cache.
+        """
+        try:
+            if hasattr(self.inspector, "info_cache"):
+                self.inspector.info_cache.clear()
+        except Exception as exc:
+            logger.debug(f"Failed to clear reflection cache: {exc}")
+
     def query_table_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
         """
         Handle custom table types
         """
+        # Clear cached column / relation data from prior schemas to
+        # prevent unbounded memory growth (issue #20649)
+        self._clear_reflection_cache()
+
+        self.get_partition_details(schema_name)
         self._set_constraint_details(schema_name)
 
         result = self.connection.execute(
@@ -305,7 +325,6 @@ def set_external_location_map(self, database_name: str) -> None:
     def get_database_names(self) -> Iterable[str]:
         if not self.config.serviceConnection.root.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.root.config.database
-            self.get_partition_details()
             self._set_incremental_table_processor(configured_db)
             self.set_external_location_map(configured_db)
             yield configured_db
@@ -331,7 +350,6 @@ def get_database_names(self) -> Iterable[str]:
 
                 try:
                     self.set_inspector(database_name=new_database)
-                    self.get_partition_details()
                     self._set_incremental_table_processor(new_database)
                     self.set_external_location_map(new_database)
                     yield new_database

ingestion/src/metadata/ingestion/source/database/redshift/utils.py

@@ -51,7 +51,6 @@
 
 # pylint: disable=protected-access
 @calculate_execution_time()
-@reflection.cache
 def get_columns(self, connection, table_name, schema=None, **kw):
     """
     Return information about columns in `table_name`.
@@ -61,6 +60,10 @@ def get_columns(self, connection, table_name, schema=None, **kw):
 
     overriding the default dialect method to include the
     distkey and sortkey info
+
+    Note: @reflection.cache removed to avoid unbounded memory growth
+    across schemas (issue #20649). The underlying
+    _get_schema_column_info already caches per-schema.
     """
     cols = self._get_redshift_columns(connection, table_name, schema, **kw)
     if not self._domains:
@@ -121,22 +124,21 @@ def _get_column_info(self, *args, **kwargs):
 
 
 @calculate_execution_time()
-@reflection.cache
 def _get_schema_column_info(
     self, connection, schema=None, **kw
 ):  # pylint: disable=unused-argument
     """
     Get schema column info
 
-    Args:
-        connection:
-        schema:
-        **kw:
-    Returns:
-
-    This method is responsible for fetching all the column details like
-    name, type, constraints, distkey and sortkey etc.
+    Uses a custom single-schema cache instead of @reflection.cache
+    to prevent unbounded memory growth across schemas (issue #20649).
+    Only the most recently requested schema's data is retained.
     """
+    # Single-schema cache: invalidate when schema changes
+    cached = getattr(self, "_schema_col_cache", None)
+    if cached is not None and cached[0] == schema:
+        return cached[1]
+
     schema_clause = f"AND schema = '{schema if schema else ''}'"
     all_columns = defaultdict(list)
     result = connection.execute(
@@ -145,7 +147,10 @@ def _get_schema_column_info(
     for col in result:
         key = RelationKey(col.table_name, col.schema, connection)
         all_columns[key].append(col)
-    return dict(all_columns)
+    result.close()
+    data = dict(all_columns)
+    self._schema_col_cache = (schema, data)
+    return data
 
 
 def _handle_array_type(attype):
@@ -377,10 +382,19 @@ def get_table_comment(
 
 
 @calculate_execution_time()
-@reflection.cache
 def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-argument
+    """
+    Uses a custom single-schema cache instead of @reflection.cache
+    to prevent unbounded memory growth across schemas (issue #20649).
+    """
     # pylint: disable=consider-using-f-string
     schema = kw.get("schema", None)
+
+    # Single-schema cache: invalidate when schema changes
+    cached = getattr(self, "_relation_info_cache", None)
+    if cached is not None and cached[0] == schema:
+        return cached[1]
+
     schema_clause = "AND schema = '{schema}'".format(schema=schema) if schema else ""
 
     table_name = kw.get("table_name", None)
@@ -399,6 +413,8 @@ def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-ar
     for rel in result:
         key = RelationKey(rel.relname, rel.schema, connection)
         relations[key] = rel
+    result.close()
+    self._relation_info_cache = (schema, relations)
     return relations