AirbyteLib: Add len() support on SQL datasets and Mapping behaviors for ReadResult (#34763)

airbyte-lib/airbyte_lib/datasets/_base.py

@@ -11,11 +11,24 @@
 class DatasetBase(ABC):
     """Base implementation for all datasets."""
 
+    def __init__(self) -> None:
+        self._length: int | None = None
+
     @abstractmethod
     def __iter__(self) -> Iterator[Mapping[str, Any]]:
         """Return the iterator of records."""
         raise NotImplementedError
 
+    def __len__(self) -> int:
+        """Return the number of records in the dataset.
+
+        This method caches the length of the dataset after the first call.
+        """
+        if self._length is None:
+            self._length = sum(1 for _ in self)
+
+        return self._length
+
     def to_pandas(self) -> DataFrame:
         """Return a pandas DataFrame representation of the dataset.
 

airbyte-lib/airbyte_lib/datasets/_lazy.py

@@ -20,6 +20,7 @@ def __init__(
         iterator: Iterator[Mapping[str, Any]],
     ) -> None:
         self._iterator: Iterator[Mapping[str, Any]] = iterator
+        super().__init__()
 
     @overrides
     def __iter__(self) -> Iterator[Mapping[str, Any]]:

airbyte-lib/airbyte_lib/datasets/_sql.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from overrides import overrides
-from sqlalchemy import and_, text
+from sqlalchemy import and_, func, select, text
 
 from airbyte_lib.datasets._base import DatasetBase
 
@@ -36,6 +36,7 @@ def __init__(
         self._cache: SQLCacheBase = cache
         self._stream_name: str = stream_name
         self._query_statement: Selectable = query_statement
+        super().__init__()
 
     @property
     def stream_name(self) -> str:
@@ -48,6 +49,13 @@ def __iter__(self) -> Iterator[Mapping[str, Any]]:
                 # https://pydoc.dev/sqlalchemy/latest/sqlalchemy.engine.row.RowMapping.html
                 yield cast(Mapping[str, Any], row._mapping)  # noqa: SLF001
 
+    def __len__(self) -> int:
+        if self._length is None:
+            count_query = select([func.count()]).select_from(self._query_statement.alias())
+            with self._cache.get_sql_connection() as conn:
+                self._length = conn.execute(count_query).scalar()
+        return self._length
+
     def to_pandas(self) -> DataFrame:
         return self._cache.get_pandas_dataframe(self._stream_name)
 
@@ -85,9 +93,11 @@ class CachedDataset(SQLDataset):
     """
 
     def __init__(self, cache: SQLCacheBase, stream_name: str) -> None:
-        self._cache: SQLCacheBase = cache
-        self._stream_name: str = stream_name
-        self._query_statement: Selectable = self.to_sql_table().select()
+        super().__init__(
+            cache=cache,
+            stream_name=stream_name,
+            query_statement=self.to_sql_table().select(),
+        )
 
     @overrides
     def to_pandas(self) -> DataFrame:

airbyte-lib/airbyte_lib/results.py

@@ -1,20 +1,21 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING
 
 from airbyte_lib.datasets import CachedDataset
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Mapping
+    from collections.abc import Iterator
 
     from sqlalchemy.engine import Engine
 
     from airbyte_lib.caches import SQLCacheBase
 
 
-class ReadResult:
+class ReadResult(Mapping[str, CachedDataset]):
     def __init__(
         self, processed_records: int, cache: SQLCacheBase, processed_streams: list[str]
     ) -> None:
@@ -34,6 +35,9 @@ def __contains__(self, stream: str) -> bool:
     def __iter__(self) -> Iterator[str]:
         return self._processed_streams.__iter__()
 
+    def __len__(self) -> int:
+        return len(self._processed_streams)
+
     def get_sql_engine(self) -> Engine:
         return self._cache.get_sql_engine()