fix(datalake): hybrid two-pass listing, ruff checkstyle, and Copilot findings

- Hybrid two-pass: non-Iceberg buckets (common case) skip regex/set overhead
  in pass 2; Iceberg buckets apply directory filter as before
- Move _ICEBERG_METADATA_RE and _update_iceberg_entry to DatalakeBaseClient (DRY)
- Fix _is_json_lines false-positive for minified Iceberg/Delta metadata dicts
- Fix sys.modules stub to use setdefault (avoids overwriting real google packages)
- Fix ruff: SIM103, TRY300 in json.py; RUF013, RUF059 in test file

Author: mohitjeswani01

ingestion/src/metadata/ingestion/source/database/datalake/clients/base.py

@@ -42,7 +42,6 @@ def _update_iceberg_entry(
             iceberg_tables[table_dir] = (version, name, size)
         return True
 
-
     def __init__(self, client: Any, session: Any = None, **kwargs):
         self._client = client
         self._session = session

ingestion/src/metadata/ingestion/source/database/datalake/clients/gcs.py

@@ -126,13 +126,12 @@ def get_table_names(
         files without accumulation, keeping memory overhead at O(1) per object.
         """
         bucket = self._client.get_bucket(bucket_name)
-        iceberg_tables: Dict[str, Tuple[int, str, Optional[int]]] = {}  # noqa: UP006
+        iceberg_tables: Dict[str, Tuple[int, str, int | None]] = {}  # noqa: UP006
 
         for blob in bucket.list_blobs(prefix=prefix):
             if skip_cold_storage and self._should_skip_gcs_cold_storage(blob):
                 logger.debug(
-                    f"Skipping cold storage object: {blob.name} "
-                    f"(storage_class: {getattr(blob, 'storage_class', None)})"
+                    f"Skipping cold storage object: {blob.name} (storage_class: {getattr(blob, 'storage_class', None)})"
                 )
                 continue
             self._update_iceberg_entry(iceberg_tables, blob.name, blob.size)
@@ -141,13 +140,19 @@ def get_table_names(
         for _, metadata_blob_path, size in iceberg_tables.values():
             yield metadata_blob_path, size
 
-        for blob in bucket.list_blobs(prefix=prefix):
-            if skip_cold_storage and self._should_skip_gcs_cold_storage(blob):
-                continue
-            if not self._ICEBERG_METADATA_RE.match(blob.name) and not any(
-                blob.name.startswith(d + "/") for d in iceberg_dirs
-            ):
+        if not iceberg_dirs:
+            for blob in bucket.list_blobs(prefix=prefix):
+                if skip_cold_storage and self._should_skip_gcs_cold_storage(blob):
+                    continue
                 yield blob.name, blob.size
+        else:
+            for blob in bucket.list_blobs(prefix=prefix):
+                if skip_cold_storage and self._should_skip_gcs_cold_storage(blob):
+                    continue
+                if not self._ICEBERG_METADATA_RE.match(blob.name) and not any(
+                    blob.name.startswith(d + "/") for d in iceberg_dirs
+                ):
+                    yield blob.name, blob.size
 
     def close(self, service_connection):
         os.environ.pop("GOOGLE_CLOUD_PROJECT", "")

ingestion/src/metadata/ingestion/source/database/datalake/clients/s3.py

@@ -87,7 +87,7 @@ def get_table_names(
         if prefix:
             kwargs["Prefix"] = prefix if prefix.endswith("/") else f"{prefix}/"
 
-        iceberg_tables: Dict[str, Tuple[int, str, Optional[int]]] = {}  # noqa: UP006
+        iceberg_tables: Dict[str, Tuple[int, str, int | None]] = {}  # noqa: UP006
 
         for key in list_s3_objects(self._client, **kwargs):
             key_name = key["Key"]
@@ -105,15 +105,23 @@ def get_table_names(
         for _, metadata_key, size in iceberg_tables.values():
             yield metadata_key, size
 
-        for key in list_s3_objects(self._client, **kwargs):
-            key_name = key["Key"]
-            size = key.get("Size")
-            if skip_cold_storage and self._should_skip_s3_cold_storage(key):
-                continue
-            if not self._ICEBERG_METADATA_RE.match(key_name) and not any(
-                key_name.startswith(d + "/") for d in iceberg_dirs
-            ):
+        if not iceberg_dirs:
+            for key in list_s3_objects(self._client, **kwargs):
+                key_name = key["Key"]
+                size = key.get("Size")
+                if skip_cold_storage and self._should_skip_s3_cold_storage(key):
+                    continue
                 yield key_name, size
+        else:
+            for key in list_s3_objects(self._client, **kwargs):
+                key_name = key["Key"]
+                size = key.get("Size")
+                if skip_cold_storage and self._should_skip_s3_cold_storage(key):
+                    continue
+                if not self._ICEBERG_METADATA_RE.match(key_name) and not any(
+                    key_name.startswith(d + "/") for d in iceberg_dirs
+                ):
+                    yield key_name, size
 
     def get_folders_prefix(self, bucket_name: str, prefix: Optional[str]) -> Iterable[str]:  # noqa: UP045
         for page in self._client.get_paginator("list_objects_v2").paginate(

ingestion/src/metadata/readers/dataframe/json.py

@@ -125,11 +125,9 @@ def _read_json_object(
             if isinstance(data, dict)
             and (
                 data.get("$schema") is not None  # JSON Schema files
-                or data.get("format-version")
-                is not None  # Apache Iceberg table metadata
+                or data.get("format-version") is not None  # Apache Iceberg table metadata
                 or (  # Delta Lake / Iceberg schema structure
-                    isinstance(data.get("schema"), dict)
-                    and isinstance(data.get("schema", {}).get("fields"), list)
+                    isinstance(data.get("schema"), dict) and isinstance(data.get("schema", {}).get("fields"), list)
                 )
             )
             else None
@@ -153,19 +151,16 @@ def _is_json_lines(file_obj) -> bool:
             return True
         try:
             obj = json.loads(first_line)
+        except json.JSONDecodeError:
+            return False
+        else:
             if not isinstance(obj, dict):
                 return False
             if obj.get("$schema") is not None:
                 return False
             if obj.get("format-version") is not None:
                 return False
-            if isinstance(obj.get("schema"), dict) and isinstance(
-                obj.get("schema", {}).get("fields"), list
-            ):
-                return False
-            return True
-        except json.JSONDecodeError:
-            return False
+            return not (isinstance(obj.get("schema"), dict) and isinstance(obj.get("schema", {}).get("fields"), list))
 
     def _read_json_smart(
         self,

ingestion/src/metadata/utils/datalake/datalake_utils.py

@@ -153,7 +153,7 @@ def fetch_dataframe_first_chunk(
 _ICEBERG_METADATA_PATH_RE = re.compile(r"([^/]+)/metadata/v\d+\.metadata\.json$")
 
 
-def get_iceberg_table_name_from_metadata_path(metadata_path: str) -> Optional[str]:
+def get_iceberg_table_name_from_metadata_path(metadata_path: str) -> str | None:
     """
     Extracts the Iceberg table directory name from a metadata file path.
 

ingestion/tests/unit/source/database/test_iceberg_discovery.py

@@ -12,6 +12,7 @@
 """
 Tests for Iceberg table directory detection in DatalakeGcsClient and DatalakeS3Client.
 """
+
 import sys
 import types
 from unittest.mock import MagicMock, patch
@@ -21,9 +22,7 @@
 # present, which prevents breaking other tests or masking integration issues.
 _google_mod = sys.modules.setdefault("google", types.ModuleType("google"))
 _gcloud_mod = sys.modules.setdefault("google.cloud", types.ModuleType("google.cloud"))
-_storage_mod = sys.modules.setdefault(
-    "google.cloud.storage", types.ModuleType("google.cloud.storage")
-)
+_storage_mod = sys.modules.setdefault("google.cloud.storage", types.ModuleType("google.cloud.storage"))
 if not hasattr(_storage_mod, "Client"):
     _storage_mod.Client = MagicMock
 if not hasattr(_google_mod, "cloud"):
@@ -39,9 +38,7 @@
 )
 
 
-def _make_blob(
-    name: str, size: int = 1024, storage_class: str = "STANDARD"
-) -> MagicMock:
+def _make_blob(name: str, size: int = 1024, storage_class: str = "STANDARD") -> MagicMock:
     blob = MagicMock()
     blob.name = name
     blob.size = size
@@ -54,9 +51,7 @@ def _make_gcs_client(blobs: list) -> DatalakeGcsClient:
     mock_bucket = MagicMock()
     mock_storage_client.get_bucket.return_value = mock_bucket
     mock_bucket.list_blobs.return_value = blobs
-    mock_bucket.get_blob.side_effect = lambda name: next(
-        (b for b in blobs if b.name == name), None
-    )
+    mock_bucket.get_blob.side_effect = lambda name: next((b for b in blobs if b.name == name), None)
     client = DatalakeGcsClient.__new__(DatalakeGcsClient)
     client._client = mock_storage_client
     client._temp_credentials_file_path_list = []
@@ -166,17 +161,15 @@ def test_gcs_iceberg_version_comparison_v10(self):
 
 
 class TestS3IcebergDiscovery:
-    def _make_s3_client(self, keys: list, sizes: dict = None) -> DatalakeS3Client:
+    def _make_s3_client(self, keys: list, sizes: dict | None = None) -> DatalakeS3Client:
         """Helper: create a DatalakeS3Client backed by a mocked boto3 client."""
         mock_boto_client = MagicMock()
         client = DatalakeS3Client.__new__(DatalakeS3Client)
         client._client = mock_boto_client
         client._session = None
         self._mock_boto_client = mock_boto_client
         sizes = sizes or {}
-        self._s3_objects = [
-            {"Key": k, "Size": sizes.get(k, 1024)} for k in keys
-        ]
+        self._s3_objects = [{"Key": k, "Size": sizes.get(k, 1024)} for k in keys]
         return client
 
     def test_s3_iceberg_table_detected(self):
@@ -275,39 +268,19 @@ def test_iceberg_table_name_extracted_correctly(self):
             get_iceberg_table_name_from_metadata_path,
         )
 
-        assert (
-            get_iceberg_table_name_from_metadata_path(
-                "warehouse/orders/metadata/v2.metadata.json"
-            )
-            == "orders"
-        )
-        assert (
-            get_iceberg_table_name_from_metadata_path(
-                "my_prefix/sales/metadata/v1.metadata.json"
-            )
-            == "sales"
-        )
-        assert (
-            get_iceberg_table_name_from_metadata_path(
-                "simple/metadata/v3.metadata.json"
-            )
-            == "simple"
-        )
+        assert get_iceberg_table_name_from_metadata_path("warehouse/orders/metadata/v2.metadata.json") == "orders"
+        assert get_iceberg_table_name_from_metadata_path("my_prefix/sales/metadata/v1.metadata.json") == "sales"
+        assert get_iceberg_table_name_from_metadata_path("simple/metadata/v3.metadata.json") == "simple"
 
     def test_non_iceberg_path_returns_none(self):
         from metadata.utils.datalake.datalake_utils import (
             get_iceberg_table_name_from_metadata_path,
         )
 
         assert get_iceberg_table_name_from_metadata_path("data/orders.json") is None
-        assert (
-            get_iceberg_table_name_from_metadata_path("warehouse/orders.json") is None
-        )
+        assert get_iceberg_table_name_from_metadata_path("warehouse/orders.json") is None
         assert get_iceberg_table_name_from_metadata_path("metadata/v1.json") is None
-        assert (
-            get_iceberg_table_name_from_metadata_path("orders/metadata/snapshot.avro")
-            is None
-        )
+        assert get_iceberg_table_name_from_metadata_path("orders/metadata/snapshot.avro") is None
 
     def test_table_type_iceberg_for_metadata_files(self):
         from metadata.generated.schema.entity.data.table import TableType
@@ -317,9 +290,7 @@ def test_table_type_iceberg_for_metadata_files(self):
 
         key_name = "warehouse/orders/metadata/v1.metadata.json"
         table_type = (
-            TableType.Iceberg
-            if get_iceberg_table_name_from_metadata_path(key_name) is not None
-            else TableType.Regular
+            TableType.Iceberg if get_iceberg_table_name_from_metadata_path(key_name) is not None else TableType.Regular
         )
         assert table_type == TableType.Iceberg
 
@@ -335,9 +306,7 @@ def test_table_type_regular_for_normal_files(self):
                 if get_iceberg_table_name_from_metadata_path(key_name) is not None
                 else TableType.Regular
             )
-            assert (
-                table_type == TableType.Regular
-            ), f"Expected Regular for {key_name}, got {table_type}"
+            assert table_type == TableType.Regular, f"Expected Regular for {key_name}, got {table_type}"
 
 
 class TestSlice4FetchKeyCorrectness:
@@ -375,7 +344,7 @@ def test_yield_table_uses_metadata_path_not_display_name(self):
             file_size,
             original_key,
         )
-        table_name, table_type, table_extension, t_file_size, fetch_key = tuple_5
+        table_name, _table_type, table_extension, t_file_size, fetch_key = tuple_5
 
         wrapper = DatalakeTableSchemaWrapper(
             key=fetch_key,
@@ -384,12 +353,8 @@ def test_yield_table_uses_metadata_path_not_display_name(self):
             file_size=t_file_size,
         )
 
-        assert (
-            wrapper.key == original_key
-        ), f"fetch key should be original blob path, got {wrapper.key!r}"
-        assert (
-            wrapper.key != display_name
-        ), f"fetch key must NOT be the display name '{display_name}'"
+        assert wrapper.key == original_key, f"fetch key should be original blob path, got {wrapper.key!r}"
+        assert wrapper.key != display_name, f"fetch key must NOT be the display name '{display_name}'"
         assert table_name == display_name
 
     def test_non_iceberg_fetch_key_equals_table_name(self):
@@ -406,9 +371,7 @@ def test_non_iceberg_fetch_key_equals_table_name(self):
         )
 
         key_name = "data/orders.parquet"
-        table_name = (
-            key_name  # standardize_table_name returns unchanged for non-Iceberg
-        )
+        table_name = key_name  # standardize_table_name returns unchanged for non-Iceberg
 
         assert get_iceberg_table_name_from_metadata_path(key_name) is None