ingestr 0.13.41__py3-none-any.whl → 0.13.43__py3-none-any.whl

ingestr/src/attio/__init__.py

@@ -0,0 +1,100 @@
+from typing import Iterable, Iterator
+
+import dlt
+from dlt.sources import DltResource
+
+from .helpers import AttioClient
+
+
+@dlt.source(max_table_nesting=0)
+def attio_source(
+    api_key: str,
+    params: list[str],
+) -> Iterable[DltResource]:
+    
+    attio_client = AttioClient(api_key)
+
+    @dlt.resource(
+        name="objects",
+        write_disposition="replace",
+        columns={
+            "created_at": {"data_type": "timestamp", "partition": True},
+        },
+    )
+    def fetch_objects() -> Iterator[dict]:
+        if len(params) != 0:
+            raise ValueError("Objects table must be in the format `objects`")
+
+        path = "objects"
+        yield attio_client.fetch_data(path, "get")
+
+    @dlt.resource(
+        name="records",
+        write_disposition="replace",
+        columns={
+            "created_at": {"data_type": "timestamp", "partition": True},
+        },
+    )
+    def fetch_records() -> Iterator[dict]:
+        if len(params) != 1:
+            raise ValueError(
+                "Records table must be in the format `records:{object_api_slug}`"
+            )
+
+        object_id = params[0]
+        path = f"objects/{object_id}/records/query"
+
+        yield attio_client.fetch_data(path, "post")
+
+    @dlt.resource(
+        name="lists",
+        write_disposition="replace",
+        columns={
+            "created_at": {"data_type": "timestamp", "partition": True},
+        },
+    )
+    def fetch_lists() -> Iterator[dict]:
+        path = "lists"
+        yield attio_client.fetch_data(path, "get")
+
+    @dlt.resource(
+        name="list_entries",
+        write_disposition="replace",
+        columns={
+            "created_at": {"data_type": "timestamp", "partition": True},
+        },
+    )
+    def fetch_list_entries() -> Iterator[dict]:
+        if len(params) != 1:
+            raise ValueError(
+                "List entries table must be in the format `list_entries:{list_id}`"
+            )
+        path = f"lists/{params[0]}/entries/query"
+
+        yield attio_client.fetch_data(path, "post")
+
+    @dlt.resource(
+        name="all_list_entries",
+        write_disposition="replace",
+        columns={
+            "created_at": {"data_type": "timestamp", "partition": True},
+        },
+    )
+    def fetch_all_list_entries() -> Iterator[dict]:
+        if len(params) != 1:
+            raise ValueError(
+                "All list entries table must be in the format `all_list_entries:{object_api_slug}`"
+            )
+        path = "lists"
+        for lst in attio_client.fetch_data(path, "get"):
+            if params[0] in lst["parent_object"]:
+                path = f"lists/{lst['id']['list_id']}/entries/query"
+                yield from attio_client.fetch_data(path, "post")
+
+    return (
+        fetch_objects,
+        fetch_records,
+        fetch_lists,
+        fetch_list_entries,
+        fetch_all_list_entries,
+    )

ingestr/src/attio/helpers.py

@@ -0,0 +1,54 @@
+from ingestr.src.http_client import create_client
+
+
+class AttioClient:
+    def __init__(self, api_key: str):
+        self.base_url = "https://api.attio.com/v2"
+        self.headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+        self.client = create_client()
+
+    def fetch_data(self, path: str, method: str, limit: int = 1000, params=None):
+        url = f"{self.base_url}/{path}"
+        if params is None:
+            params = {}
+        offset = 0
+        while True:
+            query_params = {**params, "limit": limit, "offset": offset}
+            if method == "get":
+                response = self.client.get(
+                    url, headers=self.headers, params=query_params
+                )
+            else:
+                response = self.client.post(
+                    url, headers=self.headers, params=query_params
+                )
+
+            if response.status_code != 200:
+                raise Exception(f"HTTP {response.status_code} error: {response.text}")
+
+            response_data = response.json()
+            if "data" not in response_data:
+                print(f"API Response: {response_data}")
+                raise Exception(
+                    "Attio API returned a response without the expected data"
+                )
+
+            data = response_data["data"]
+
+            for item in data:
+                flat_item = flatten_item(item)
+                yield flat_item
+
+            if len(data) < limit:
+                break
+            offset += limit
+
+
+def flatten_item(item: dict) -> dict:
+    if "id" in item:
+        for key, value in item["id"].items():
+            item[key] = value
+    return item

ingestr/src/buildinfo.py

@@ -1 +1 @@
-version = "v0.13.41"
+version = "v0.13.43"

ingestr/src/factory.py

@@ -26,6 +26,7 @@ from ingestr.src.sources import (
     AppsflyerSource,
     ArrowMemoryMappedSource,
     AsanaSource,
+    AttioSource,
     ChessSource,
     DynamoDBSource,
     ElasticsearchSource,
@@ -120,6 +121,7 @@ class SourceDestinationFactory:
     sources: Dict[str, Type[SourceProtocol]] = {
         "csv": LocalCsvSource,
         "mongodb": MongoDbSource,
+        "mongodb+srv": MongoDbSource,
         "notion": NotionSource,
         "gsheets": GoogleSheetsSource,
         "shopify": ShopifySource,
@@ -156,6 +158,7 @@ class SourceDestinationFactory:
         "freshdesk": FreshdeskSource,
         "phantombuster": PhantombusterSource,
         "elasticsearch": ElasticsearchSource,
+        "attio": AttioSource,
     }
     destinations: Dict[str, Type[DestinationProtocol]] = {
         "bigquery": BigQueryDestination,

ingestr/src/http_client.py

@@ -0,0 +1,17 @@
+import requests
+from dlt.sources.helpers.requests import Client
+
+def create_client() -> requests.Session:
+    return Client(
+        raise_for_status=False,
+        retry_condition=retry_on_limit,
+        request_max_attempts=12,
+    ).session
+
+
+def retry_on_limit(
+    response: requests.Response | None, exception: BaseException | None
+) -> bool:
+    if response is None:
+        return False
+    return response.status_code == 502

ingestr/src/mongodb/__init__.py

@@ -1,8 +1,9 @@
 """Source that loads collections form any a mongo database, supports incremental loads."""
 
-from typing import Any, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Union
 
 import dlt
+from dlt.common.data_writers import TDataItemFormat
 from dlt.sources import DltResource
 
 from .helpers import (
@@ -21,6 +22,10 @@ def mongodb(
     incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
     write_disposition: Optional[str] = dlt.config.value,
     parallel: Optional[bool] = dlt.config.value,
+    limit: Optional[int] = None,
+    filter_: Optional[Dict[str, Any]] = None,
+    projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    pymongoarrow_schema: Optional[Any] = None,
 ) -> Iterable[DltResource]:
     """
     A DLT source which loads data from a mongo database using PyMongo.
@@ -34,6 +39,18 @@ def mongodb(
             E.g., `incremental=dlt.sources.incremental('updated_at', pendulum.parse('2022-01-01T00:00:00Z'))`
         write_disposition (str): Write disposition of the resource.
         parallel (Optional[bool]): Option to enable parallel loading for the collection. Default is False.
+        limit (Optional[int]):
+            The maximum number of documents to load. The limit is
+            applied to each requested collection separately.
+        filter_ (Optional[Dict[str, Any]]): The filter to apply to the collection.
+        projection: (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields of a collection
+            when loading the collection. Supported inputs:
+                include (list) - ["year", "title"]
+                include (dict) - {"year": True, "title": True}
+                exclude (dict) - {"released": False, "runtime": False}
+            Note: Can't mix include and exclude statements '{"title": True, "released": False}`
+        pymongoarrow_schema (pymongoarrow.schema.Schema): Mapping of expected field types of a collection to convert BSON to Arrow
+
     Returns:
         Iterable[DltResource]: A list of DLT resources for each collection to be loaded.
     """
@@ -58,19 +75,36 @@ def mongodb(
             primary_key="_id",
             write_disposition=write_disposition,
             spec=MongoDbCollectionConfiguration,
-        )(client, collection, incremental=incremental, parallel=parallel)
+        )(
+            client,
+            collection,
+            incremental=incremental,
+            parallel=parallel,
+            limit=limit,
+            filter_=filter_ or {},
+            projection=projection,
+            pymongoarrow_schema=pymongoarrow_schema,
+        )
 
 
-@dlt.common.configuration.with_config(
-    sections=("sources", "mongodb"), spec=MongoDbCollectionResourceConfiguration
+@dlt.resource(
+    name=lambda args: args["collection"],
+    standalone=True,
+    spec=MongoDbCollectionResourceConfiguration,
 )
 def mongodb_collection(
-    connection_url: str = dlt.config.value,
+    connection_url: str = dlt.secrets.value,
     database: Optional[str] = dlt.config.value,
     collection: str = dlt.config.value,
     incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
     write_disposition: Optional[str] = dlt.config.value,
     parallel: Optional[bool] = False,
+    limit: Optional[int] = None,
+    chunk_size: Optional[int] = 10000,
+    data_item_format: Optional[TDataItemFormat] = "object",
+    filter_: Optional[Dict[str, Any]] = None,
+    projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = dlt.config.value,
+    pymongoarrow_schema: Optional[Any] = None,
 ) -> Any:
     """
     A DLT source which loads a collection from a mongo database using PyMongo.
@@ -83,6 +117,21 @@ def mongodb_collection(
             E.g., `incremental=dlt.sources.incremental('updated_at', pendulum.parse('2022-01-01T00:00:00Z'))`
         write_disposition (str): Write disposition of the resource.
         parallel (Optional[bool]): Option to enable parallel loading for the collection. Default is False.
+        limit (Optional[int]): The number of documents load.
+        chunk_size (Optional[int]): The number of documents load in each batch.
+        data_item_format (Optional[TDataItemFormat]): The data format to use for loading.
+            Supported formats:
+                object - Python objects (dicts, lists).
+                arrow - Apache Arrow tables.
+        filter_ (Optional[Dict[str, Any]]): The filter to apply to the collection.
+        projection: (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields
+            when loading the collection. Supported inputs:
+                include (list) - ["year", "title"]
+                include (dict) - {"year": True, "title": True}
+                exclude (dict) - {"released": False, "runtime": False}
+            Note: Can't mix include and exclude statements '{"title": True, "released": False}`
+        pymongoarrow_schema (pymongoarrow.schema.Schema): Mapping of expected field types to convert BSON to Arrow
+
     Returns:
         Iterable[DltResource]: A list of DLT resources for each collection to be loaded.
     """
@@ -100,4 +149,15 @@ def mongodb_collection(
         name=collection_obj.name,
         primary_key="_id",
         write_disposition=write_disposition,
-    )(client, collection_obj, incremental=incremental, parallel=parallel)
+    )(
+        client,
+        collection_obj,
+        incremental=incremental,
+        parallel=parallel,
+        limit=limit,
+        chunk_size=chunk_size,
+        data_item_format=data_item_format,
+        filter_=filter_ or {},
+        projection=projection,
+        pymongoarrow_schema=pymongoarrow_schema,
+    )

ingestr/src/mongodb/helpers.py

@@ -1,12 +1,27 @@
 """Mongo database source helpers"""
 
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import dlt
 from bson.decimal128 import Decimal128
 from bson.objectid import ObjectId
+from bson.regex import Regex
+from bson.timestamp import Timestamp
+from dlt.common import logger
 from dlt.common.configuration.specs import BaseConfiguration, configspec
+from dlt.common.data_writers import TDataItemFormat
 from dlt.common.time import ensure_pendulum_datetime
 from dlt.common.typing import TDataItem
 from dlt.common.utils import map_nested_in_place
@@ -14,17 +29,23 @@ from pendulum import _datetime
 from pymongo import ASCENDING, DESCENDING, MongoClient
 from pymongo.collection import Collection
 from pymongo.cursor import Cursor
+from pymongo.helpers_shared import _fields_list_to_dict
 
 if TYPE_CHECKING:
     TMongoClient = MongoClient[Any]
-    TCollection = Collection[Any]  # type: ignore
+    TCollection = Collection[Any]
     TCursor = Cursor[Any]
 else:
     TMongoClient = Any
     TCollection = Any
     TCursor = Any
 
-CHUNK_SIZE = 10000
+try:
+    import pymongoarrow  # type: ignore
+
+    PYMONGOARROW_AVAILABLE = True
+except ImportError:
+    PYMONGOARROW_AVAILABLE = False
 
 
 class CollectionLoader:
@@ -32,11 +53,14 @@ class CollectionLoader:
         self,
         client: TMongoClient,
         collection: TCollection,
+        chunk_size: int,
         incremental: Optional[dlt.sources.incremental[Any]] = None,
     ) -> None:
         self.client = client
         self.collection = collection
         self.incremental = incremental
+        self.chunk_size = chunk_size
+
         if incremental:
             self.cursor_field = incremental.cursor_path
             self.last_value = incremental.last_value
@@ -45,45 +69,186 @@ class CollectionLoader:
             self.last_value = None
 
     @property
-    def _filter_op(self) -> Dict[str, Any]:
+    def _sort_op(self) -> List[Optional[Tuple[str, int]]]:
         if not self.incremental or not self.last_value:
-            return {}
-        if self.incremental.last_value_func is max:
-            return {self.cursor_field: {"$gte": self.last_value}}
-        elif self.incremental.last_value_func is min:
-            return {self.cursor_field: {"$lt": self.last_value}}
-        return {}
+            return []
 
-    def load_documents(self) -> Iterator[TDataItem]:
-        cursor = self.collection.find(self._filter_op)
-        while docs_slice := list(islice(cursor, CHUNK_SIZE)):
-            yield map_nested_in_place(convert_mongo_objs, docs_slice)
+        if (
+            self.incremental.row_order == "asc"
+            and self.incremental.last_value_func is max
+        ) or (
+            self.incremental.row_order == "desc"
+            and self.incremental.last_value_func is min
+        ):
+            return [(self.cursor_field, ASCENDING)]
 
+        elif (
+            self.incremental.row_order == "asc"
+            and self.incremental.last_value_func is min
+        ) or (
+            self.incremental.row_order == "desc"
+            and self.incremental.last_value_func is max
+        ):
+            return [(self.cursor_field, DESCENDING)]
+
+        return []
 
-class CollectionLoaderParallell(CollectionLoader):
     @property
-    def _sort_op(self) -> List[Optional[Tuple[str, int]]]:
-        if not self.incremental or not self.last_value:
-            return []
+    def _filter_op(self) -> Dict[str, Any]:
+        """Build a filtering operator.
+
+        Includes a field and the filtering condition for it.
+
+        Returns:
+            Dict[str, Any]: A dictionary with the filter operator.
+        """
+        if not (self.incremental and self.last_value):
+            return {}
+
+        filt = {}
         if self.incremental.last_value_func is max:
-            return [(self.cursor_field, ASCENDING)]
+            filt = {self.cursor_field: {"$gte": self.last_value}}
+            if self.incremental.end_value:
+                filt[self.cursor_field]["$lt"] = self.incremental.end_value
+
         elif self.incremental.last_value_func is min:
-            return [(self.cursor_field, DESCENDING)]
-        return []
+            filt = {self.cursor_field: {"$lte": self.last_value}}
+            if self.incremental.end_value:
+                filt[self.cursor_field]["$gt"] = self.incremental.end_value
+
+        return filt
+
+    def _projection_op(
+        self, projection: Optional[Union[Mapping[str, Any], Iterable[str]]]
+    ) -> Optional[Dict[str, Any]]:
+        """Build a projection operator.
+
+        Args:
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): A tuple of fields to include or a dict specifying fields to include or exclude.
+            The incremental `primary_key` needs to be handle differently for inclusion
+            and exclusion projections.
+
+        Returns:
+            Tuple[str, ...] | Dict[str, Any]: A tuple or dictionary with the projection operator.
+        """
+        if projection is None:
+            return None
+
+        projection_dict = dict(_fields_list_to_dict(projection, "projection"))
+
+        if self.incremental:
+            # this is an inclusion projection
+            if any(v == 1 for v in projection_dict.values()):
+                # ensure primary_key is included
+                projection_dict.update(m={self.incremental.primary_key: 1})
+            # this is an exclusion projection
+            else:
+                try:
+                    # ensure primary_key isn't excluded
+                    projection_dict.pop(self.incremental.primary_key)  # type: ignore
+                except KeyError:
+                    pass  # primary_key was properly not included in exclusion projection
+                else:
+                    dlt.common.logger.warn(
+                        f"Primary key `{self.incremental.primary_key}` was removed from exclusion projection"
+                    )
+
+        return projection_dict
+
+    def _limit(self, cursor: Cursor, limit: Optional[int] = None) -> TCursor:  # type: ignore
+        """Apply a limit to the cursor, if needed.
+
+        Args:
+            cursor (Cursor): The cursor to apply the limit.
+            limit (Optional[int]): The number of documents to load.
+
+        Returns:
+            Cursor: The cursor with the limit applied (if given).
+        """
+        if limit not in (0, None):
+            if self.incremental is None or self.incremental.last_value_func is None:
+                logger.warning(
+                    "Using limit without ordering - results may be inconsistent."
+                )
+
+            cursor = cursor.limit(abs(limit))
 
+        return cursor
+
+    def load_documents(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    ) -> Iterator[TDataItem]:
+        """Construct the query and load the documents from the collection.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+
+        Yields:
+            Iterator[TDataItem]: An iterator of the loaded documents.
+        """
+        filter_op = self._filter_op
+        _raise_if_intersection(filter_op, filter_)
+        filter_op.update(filter_)
+
+        projection_op = self._projection_op(projection)
+
+        cursor = self.collection.find(filter=filter_op, projection=projection_op)
+        if self._sort_op:
+            cursor = cursor.sort(self._sort_op)
+
+        cursor = self._limit(cursor, limit)
+
+        while docs_slice := list(islice(cursor, self.chunk_size)):
+            yield map_nested_in_place(convert_mongo_objs, docs_slice)
+
+
+class CollectionLoaderParallel(CollectionLoader):
     def _get_document_count(self) -> int:
         return self.collection.count_documents(filter=self._filter_op)
 
-    def _create_batches(self) -> List[Dict[str, int]]:
+    def _create_batches(self, limit: Optional[int] = None) -> List[Dict[str, int]]:
         doc_count = self._get_document_count()
-        return [
-            dict(skip=sk, limit=CHUNK_SIZE) for sk in range(0, doc_count, CHUNK_SIZE)
-        ]
+        if limit:
+            doc_count = min(doc_count, abs(limit))
+
+        batches = []
+        left_to_load = doc_count
+
+        for sk in range(0, doc_count, self.chunk_size):
+            batches.append(dict(skip=sk, limit=min(self.chunk_size, left_to_load)))
+            left_to_load -= self.chunk_size
+
+        return batches
 
-    def _get_cursor(self) -> TCursor:
-        cursor = self.collection.find(filter=self._filter_op)
+    def _get_cursor(
+        self,
+        filter_: Dict[str, Any],
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    ) -> TCursor:
+        """Get a reading cursor for the collection.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+
+        Returns:
+            Cursor: The cursor for the collection.
+        """
+        filter_op = self._filter_op
+        _raise_if_intersection(filter_op, filter_)
+        filter_op.update(filter_)
+
+        projection_op = self._projection_op(projection)
+
+        cursor = self.collection.find(filter=filter_op, projection=projection_op)
         if self._sort_op:
             cursor = cursor.sort(self._sort_op)
+
         return cursor
 
     @dlt.defer
@@ -93,25 +258,223 @@ class CollectionLoaderParallell(CollectionLoader):
         data = []
         for document in cursor.skip(batch["skip"]).limit(batch["limit"]):
             data.append(map_nested_in_place(convert_mongo_objs, document))
+
         return data
 
-    def _get_all_batches(self) -> Iterator[TDataItem]:
-        batches = self._create_batches()
-        cursor = self._get_cursor()
+    def _get_all_batches(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    ) -> Iterator[TDataItem]:
+        """Load all documents from the collection in parallel batches.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The maximum number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+
+        Yields:
+            Iterator[TDataItem]: An iterator of the loaded documents.
+        """
+        batches = self._create_batches(limit=limit)
+        cursor = self._get_cursor(filter_=filter_, projection=projection)
 
         for batch in batches:
             yield self._run_batch(cursor=cursor, batch=batch)
 
-    def load_documents(self) -> Iterator[TDataItem]:
-        for document in self._get_all_batches():
+    def load_documents(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    ) -> Iterator[TDataItem]:
+        """Load documents from the collection in parallel.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+
+        Yields:
+            Iterator[TDataItem]: An iterator of the loaded documents.
+        """
+        for document in self._get_all_batches(
+            limit=limit, filter_=filter_, projection=projection
+        ):
             yield document
 
 
+class CollectionArrowLoader(CollectionLoader):
+    """
+    Mongo DB collection loader, which uses
+    Apache Arrow for data processing.
+    """
+
+    def load_documents(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+        pymongoarrow_schema: Any = None,
+    ) -> Iterator[Any]:
+        """
+        Load documents from the collection in Apache Arrow format.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+            pymongoarrow_schema (Any): The mapping of field types to convert BSON to Arrow.
+
+        Yields:
+            Iterator[Any]: An iterator of the loaded documents.
+        """
+        from pymongoarrow.context import PyMongoArrowContext  # type: ignore
+        from pymongoarrow.lib import process_bson_stream  # type: ignore
+
+        filter_op = self._filter_op
+        _raise_if_intersection(filter_op, filter_)
+        filter_op.update(filter_)
+
+        projection_op = self._projection_op(projection)
+
+        # NOTE the `filter_op` isn't passed
+        cursor = self.collection.find_raw_batches(
+            filter_, batch_size=self.chunk_size, projection=projection_op
+        )
+        if self._sort_op:
+            cursor = cursor.sort(self._sort_op)  # type: ignore
+
+        cursor = self._limit(cursor, limit)  # type: ignore
+
+        context = PyMongoArrowContext.from_schema(
+            schema=pymongoarrow_schema, codec_options=self.collection.codec_options
+        )
+        for batch in cursor:
+            process_bson_stream(batch, context)
+            table = context.finish()
+            yield convert_arrow_columns(table)
+
+
+class CollectionArrowLoaderParallel(CollectionLoaderParallel):
+    """
+    Mongo DB collection parallel loader, which uses
+    Apache Arrow for data processing.
+    """
+
+    def load_documents(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+        pymongoarrow_schema: Any = None,
+    ) -> Iterator[TDataItem]:
+        """Load documents from the collection in parallel.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+            pymongoarrow_schema (Any): The mapping of field types to convert BSON to Arrow.
+
+        Yields:
+            Iterator[TDataItem]: An iterator of the loaded documents.
+        """
+        yield from self._get_all_batches(
+            limit=limit,
+            filter_=filter_,
+            projection=projection,
+            pymongoarrow_schema=pymongoarrow_schema,
+        )
+
+    def _get_all_batches(
+        self,
+        filter_: Dict[str, Any],
+        limit: Optional[int] = None,
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+        pymongoarrow_schema: Any = None,
+    ) -> Iterator[TDataItem]:
+        """Load all documents from the collection in parallel batches.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            limit (Optional[int]): The maximum number of documents to load.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+            pymongoarrow_schema (Any): The mapping of field types to convert BSON to Arrow.
+
+        Yields:
+            Iterator[TDataItem]: An iterator of the loaded documents.
+        """
+        batches = self._create_batches(limit=limit)
+        cursor = self._get_cursor(filter_=filter_, projection=projection)
+        for batch in batches:
+            yield self._run_batch(
+                cursor=cursor,
+                batch=batch,
+                pymongoarrow_schema=pymongoarrow_schema,
+            )
+
+    def _get_cursor(
+        self,
+        filter_: Dict[str, Any],
+        projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None,
+    ) -> TCursor:
+        """Get a reading cursor for the collection.
+
+        Args:
+            filter_ (Dict[str, Any]): The filter to apply to the collection.
+            projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+
+        Returns:
+            Cursor: The cursor for the collection.
+        """
+        filter_op = self._filter_op
+        _raise_if_intersection(filter_op, filter_)
+        filter_op.update(filter_)
+
+        projection_op = self._projection_op(projection)
+
+        cursor = self.collection.find_raw_batches(
+            filter=filter_op, batch_size=self.chunk_size, projection=projection_op
+        )
+        if self._sort_op:
+            cursor = cursor.sort(self._sort_op)  # type: ignore
+
+        return cursor
+
+    @dlt.defer
+    def _run_batch(
+        self,
+        cursor: TCursor,
+        batch: Dict[str, int],
+        pymongoarrow_schema: Any = None,
+    ) -> TDataItem:
+        from pymongoarrow.context import PyMongoArrowContext
+        from pymongoarrow.lib import process_bson_stream
+
+        cursor = cursor.clone()
+
+        context = PyMongoArrowContext.from_schema(
+            schema=pymongoarrow_schema, codec_options=self.collection.codec_options
+        )
+        for chunk in cursor.skip(batch["skip"]).limit(batch["limit"]):
+            process_bson_stream(chunk, context)
+            table = context.finish()
+            yield convert_arrow_columns(table)
+
+
 def collection_documents(
     client: TMongoClient,
     collection: TCollection,
+    filter_: Dict[str, Any],
+    projection: Union[Dict[str, Any], List[str]],
+    pymongoarrow_schema: "pymongoarrow.schema.Schema",
     incremental: Optional[dlt.sources.incremental[Any]] = None,
     parallel: bool = False,
+    limit: Optional[int] = None,
+    chunk_size: Optional[int] = 10000,
+    data_item_format: Optional[TDataItemFormat] = "object",
 ) -> Iterator[TDataItem]:
     """
     A DLT source which loads data from a Mongo database using PyMongo.
@@ -120,27 +483,145 @@ def collection_documents(
     Args:
         client (MongoClient): The PyMongo client `pymongo.MongoClient` instance.
         collection (Collection): The collection `pymongo.collection.Collection` to load.
+        filter_ (Dict[str, Any]): The filter to apply to the collection.
+        projection (Optional[Union[Mapping[str, Any], Iterable[str]]]): The projection to select fields to create the Cursor.
+            when loading the collection. Supported inputs:
+                include (list) - ["year", "title"]
+                include (dict) - {"year": True, "title": True}
+                exclude (dict) - {"released": False, "runtime": False}
+            Note: Can't mix include and exclude statements '{"title": True, "released": False}`
+        pymongoarrow_schema (pymongoarrow.schema.Schema): The mapping of field types to convert BSON to Arrow.
         incremental (Optional[dlt.sources.incremental[Any]]): The incremental configuration.
         parallel (bool): Option to enable parallel loading for the collection. Default is False.
+        limit (Optional[int]): The maximum number of documents to load.
+        chunk_size (Optional[int]): The number of documents to load in each batch.
+        data_item_format (Optional[TDataItemFormat]): The data format to use for loading.
+            Supported formats:
+                object - Python objects (dicts, lists).
+                arrow - Apache Arrow tables.
 
     Returns:
         Iterable[DltResource]: A list of DLT resources for each collection to be loaded.
     """
-    LoaderClass = CollectionLoaderParallell if parallel else CollectionLoader
+    if data_item_format == "arrow" and not PYMONGOARROW_AVAILABLE:
+        dlt.common.logger.warn(
+            "'pymongoarrow' is not installed; falling back to standard MongoDB CollectionLoader."
+        )
+        data_item_format = "object"
+
+    if data_item_format != "arrow" and pymongoarrow_schema:
+        dlt.common.logger.warn(
+            "Received value for `pymongoarrow_schema`, but `data_item_format=='object'` "
+            "Use `data_item_format=='arrow'` to enforce schema."
+        )
+
+    if data_item_format == "arrow" and pymongoarrow_schema and projection:
+        dlt.common.logger.warn(
+            "Received values for both `pymongoarrow_schema` and `projection`. Since both "
+            "create a projection to select fields, `projection` will be ignored."
+        )
+
+    if parallel:
+        if data_item_format == "arrow":
+            LoaderClass = CollectionArrowLoaderParallel
+        else:
+            LoaderClass = CollectionLoaderParallel  # type: ignore
+    else:
+        if data_item_format == "arrow":
+            LoaderClass = CollectionArrowLoader  # type: ignore
+        else:
+            LoaderClass = CollectionLoader  # type: ignore
 
-    loader = LoaderClass(client, collection, incremental=incremental)
-    for data in loader.load_documents():
-        yield data
+    loader = LoaderClass(
+        client, collection, incremental=incremental, chunk_size=chunk_size
+    )
+    if isinstance(loader, (CollectionArrowLoader, CollectionArrowLoaderParallel)):
+        yield from loader.load_documents(
+            limit=limit,
+            filter_=filter_,
+            projection=projection,
+            pymongoarrow_schema=pymongoarrow_schema,
+        )
+    else:
+        yield from loader.load_documents(
+            limit=limit, filter_=filter_, projection=projection
+        )
 
 
 def convert_mongo_objs(value: Any) -> Any:
+    """MongoDB to dlt type conversion when using Python loaders.
+
+    Notes:
+        The method `ObjectId.__str__()` creates a hexstring using `binascii.hexlify(__id).decode()`
+
+    """
     if isinstance(value, (ObjectId, Decimal128)):
         return str(value)
     if isinstance(value, _datetime.datetime):
         return ensure_pendulum_datetime(value)
+    if isinstance(value, Regex):
+        return value.try_compile().pattern
+    if isinstance(value, Timestamp):
+        date = value.as_datetime()
+        return ensure_pendulum_datetime(date)
+
     return value
 
 
+def convert_arrow_columns(table: Any) -> Any:
+    """Convert the given table columns to Python types.
+
+    Notes:
+        Calling str() matches the `convert_mongo_obs()` used in non-arrow code.
+        Pymongoarrow converts ObjectId to `fixed_size_binary[12]`, which can't be
+        converted to a string as a vectorized operation because it contains ASCII characters.
+
+        Instead, you need to loop over values using:
+        ```python
+        pyarrow.array([v.as_buffer().hex() for v in object_id_array], type=pyarrow.string())
+        # pymongoarrow simplifies this by allowing this syntax
+        [str(v) for v in object_id_array]
+        ```
+
+    Args:
+        table (pyarrow.lib.Table): The table to convert.
+
+    Returns:
+        pyarrow.lib.Table: The table with the columns converted.
+    """
+    from dlt.common.libs.pyarrow import pyarrow
+    from pymongoarrow.types import (  # type: ignore
+        _is_binary,
+        _is_code,
+        _is_decimal128,
+        _is_objectid,
+    )
+
+    for i, field in enumerate(table.schema):
+        if _is_objectid(field.type) or _is_decimal128(field.type):
+            col_values = [str(value) for value in table[field.name]]
+            table = table.set_column(
+                i,
+                pyarrow.field(field.name, pyarrow.string()),
+                pyarrow.array(col_values, type=pyarrow.string()),
+            )
+        else:
+            type_ = None
+            if _is_binary(field.type):
+                type_ = pyarrow.binary()
+            elif _is_code(field.type):
+                type_ = pyarrow.string()
+
+            if type_:
+                col_values = [value.as_py() for value in table[field.name]]
+                table = table.set_column(
+                    i,
+                    pyarrow.field(field.name, type_),
+                    pyarrow.array(col_values, type=type_),
+                )
+    return table
+
+
 def client_from_credentials(connection_url: str) -> TMongoClient:
     client: TMongoClient = MongoClient(
         connection_url, uuidRepresentation="standard", tz_aware=True
@@ -148,6 +629,27 @@ def client_from_credentials(connection_url: str) -> TMongoClient:
     return client
 
 
+def _raise_if_intersection(filter1: Dict[str, Any], filter2: Dict[str, Any]) -> None:
+    """
+    Raise an exception, if the given filters'
+    fields are intersecting.
+
+    Args:
+        filter1 (Dict[str, Any]): The first filter.
+        filter2 (Dict[str, Any]): The second filter.
+    """
+    field_inter = filter1.keys() & filter2.keys()
+    for field in field_inter:
+        if filter1[field].keys() & filter2[field].keys():
+            str_repr = str({field: filter1[field]})
+            raise ValueError(
+                (
+                    f"Filtering operator {str_repr} is already used by the "
+                    "incremental and can't be used in the filter."
+                )
+            )
+
+
 @configspec
 class MongoDbCollectionConfiguration(BaseConfiguration):
     incremental: Optional[dlt.sources.incremental] = None  # type: ignore[type-arg]
@@ -155,12 +657,13 @@ class MongoDbCollectionConfiguration(BaseConfiguration):
 
 @configspec
 class MongoDbCollectionResourceConfiguration(BaseConfiguration):
-    connection_url: str = dlt.config.value
+    connection_url: dlt.TSecretValue = dlt.secrets.value
     database: Optional[str] = dlt.config.value
     collection: str = dlt.config.value
     incremental: Optional[dlt.sources.incremental] = None  # type: ignore[type-arg]
     write_disposition: Optional[str] = dlt.config.value
     parallel: Optional[bool] = False
+    projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = dlt.config.value
 
 
 __source_name__ = "mongodb"

ingestr/src/sources.py

@@ -21,6 +21,7 @@ from urllib.parse import ParseResult, parse_qs, quote, urlencode, urlparse
 import pendulum
 from dlt.common.time import ensure_pendulum_datetime
 from dlt.extract import Incremental
+from dlt.extract.exceptions import ResourcesNotFoundError
 from dlt.sources import incremental as dlt_incremental
 from dlt.sources.credentials import (
     ConnectionStringCredentials,
@@ -2392,3 +2393,29 @@ class ElasticsearchSource:
             verify_certs=verify_certs,
             incremental=incremental,
         ).with_resources(table)
+
+
+class AttioSource:
+    def handles_incrementality(self) -> bool:
+        return True
+
+    def dlt_source(self, uri: str, table: str, **kwargs):
+        parsed_uri = urlparse(uri)
+        query_params = parse_qs(parsed_uri.query)
+        api_key = query_params.get("api_key")
+
+        if api_key is None:
+            raise MissingValueError("api_key", "Attio")
+
+        parts = table.replace(" ", "").split(":")
+        table_name = parts[0]
+        params = parts[1:]
+
+        from ingestr.src.attio import attio_source
+
+        try:
+            return attio_source(api_key=api_key[0], params=params).with_resources(
+                table_name
+            )
+        except ResourcesNotFoundError:
+            raise UnsupportedResourceError(table_name, "Attio")

{ingestr-0.13.41.dist-info → ingestr-0.13.43.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ingestr
-Version: 0.13.41
+Version: 0.13.43
 Summary: ingestr is a command-line application that ingests data from various sources and stores them in any database.
 Project-URL: Homepage, https://github.com/bruin-data/ingestr
 Project-URL: Issues, https://github.com/bruin-data/ingestr/issues

{ingestr-0.13.41.dist-info → ingestr-0.13.43.dist-info}/RECORD

@@ -2,15 +2,16 @@ ingestr/conftest.py,sha256=Q03FIJIZpLBbpj55cfCHIKEjc1FCvWJhMF2cidUJKQU,1748
 ingestr/main.py,sha256=Pe_rzwcDRKIYa7baEVUAAPOHyqQbX29RUexMl0F_S1k,25273
 ingestr/src/.gitignore,sha256=8cX1AZTSI0TcdZFGTmS_oyBjpfCzhOEt0DdAo2dFIY8,203
 ingestr/src/blob.py,sha256=onMe5ZHxPXTdcB_s2oGNdMo-XQJ3ajwOsWE9eSTGFmc,1495
-ingestr/src/buildinfo.py,sha256=AK7sGUNx6CPDKJOXeMexFRun9bjpyv2c9t4DII73Pes,21
+ingestr/src/buildinfo.py,sha256=af3TiaPbJzAW9JyfpOIeOvjzAsiboFF4f1frweJenxc,21
 ingestr/src/destinations.py,sha256=MctbeJUyNr0DRB0XYt2xAbEKkHZ40-nXXEOYCs4KuoE,15420
 ingestr/src/errors.py,sha256=Ufs4_DfE77_E3vnA1fOQdi6cmuLVNm7_SbFLkL1XPGk,686
-ingestr/src/factory.py,sha256=KJKIL9q7kU4oAVXy5o0wDwLAU0nG9y0xC8D7HzksYak,5597
+ingestr/src/factory.py,sha256=c5WfqmRrXFj1PddnBOzTzzZUHJ-Fb42cvCvsBEqn6Yo,5682
 ingestr/src/filters.py,sha256=C-_TIVkF_cxZBgG-Run2Oyn0TAhJgA8IWXZ-OPY3uek,1136
+ingestr/src/http_client.py,sha256=dbAAf6Ptxf9pSn9RmNSHL4HEn5xhFimWjE2JZ67J00o,434
 ingestr/src/loader.py,sha256=9NaWAyfkXdqAZSS-N72Iwo36Lbx4PyqIfaaH1dNdkFs,1712
 ingestr/src/partition.py,sha256=BrIP6wFJvyR7Nus_3ElnfxknUXeCipK_E_bB8kZowfc,969
 ingestr/src/resource.py,sha256=XG-sbBapFVEM7OhHQFQRTdTLlh-mHB-N4V1t8F8Tsww,543
-ingestr/src/sources.py,sha256=SWZAa6bokLurQRPtH7rxi8K-GSVLp_p9Ig1ArGRsxCo,82635
+ingestr/src/sources.py,sha256=HZQGWPuDdf4sVq91KuaJ9p7MtSfjWue8d2vknZpIwg8,83456
 ingestr/src/table_definition.py,sha256=REbAbqdlmUMUuRh8nEQRreWjPVOQ5ZcfqGkScKdCrmk,390
 ingestr/src/time.py,sha256=H_Fk2J4ShXyUM-EMY7MqCLZQhlnZMZvO952bmZPc4yE,254
 ingestr/src/version.py,sha256=J_2xgZ0mKlvuHcjdKCx2nlioneLH0I47JiU_Slr_Nwc,189
@@ -30,6 +31,8 @@ ingestr/src/arrow/__init__.py,sha256=8fEntgHseKjFMiPQIzxYzw_raicNsEgnveLi1IzBca0
 ingestr/src/asana_source/__init__.py,sha256=QwQTCb5PXts8I4wLHG9UfRP-5ChfjSe88XAVfxMV5Ag,8183
 ingestr/src/asana_source/helpers.py,sha256=PukcdDQWIGqnGxuuobbLw4hUy4-t6gxXg_XywR7Lg9M,375
 ingestr/src/asana_source/settings.py,sha256=-2tpdkwh04RvLKFvwQodnFLYn9MaxOO1hsebGnDQMTU,2829
+ingestr/src/attio/__init__.py,sha256=Dvr9rSZUlw6HGsqF7iabUit-xRC17Nv6QcmA4cRF2wA,2864
+ingestr/src/attio/helpers.py,sha256=QvB-0BV_Z-cvMTBZDwOCuhxY1cB5PraPdrDkNyQ5TsM,1715
 ingestr/src/chess/__init__.py,sha256=y0Q8aKBigeKf3N7wuB_gadMQjVJzBPUT8Jhp1ObEWjk,6812
 ingestr/src/chess/helpers.py,sha256=v1HTImOMjAF7AzZUPDIuHu00e7ut0o5y1kWcVYo4QZw,549
 ingestr/src/chess/settings.py,sha256=p0RlCGgtXUacPDEvZmwzSWmzX0Apj1riwfz-nrMK89k,158
@@ -79,8 +82,8 @@ ingestr/src/klaviyo/helpers.py,sha256=_i-SHffhv25feLDcjy6Blj1UxYLISCwVCMgGtrlnYH
 ingestr/src/linkedin_ads/__init__.py,sha256=CAPWFyV24loziiphbLmODxZUXZJwm4JxlFkr56q0jfo,1855
 ingestr/src/linkedin_ads/dimension_time_enum.py,sha256=EmHRdkFyTAfo4chGjThrwqffWJxmAadZMbpTvf0xkQc,198
 ingestr/src/linkedin_ads/helpers.py,sha256=eUWudRVlXl4kqIhfXQ1eVsUpZwJn7UFqKSpnbLfxzds,4498
-ingestr/src/mongodb/__init__.py,sha256=aMr1PFIDUMRv--ne61lR17HudsN-fsrzMeyxe9PqK2s,4335
-ingestr/src/mongodb/helpers.py,sha256=y9rYKR8eyIqam_eNsZmwSYevgi8mghh7Zp8qhTHl65s,5652
+ingestr/src/mongodb/__init__.py,sha256=Ht5HGt9UJ8LeCtilgu7hZhrebo-knRLlPIlgGQojLgk,7221
+ingestr/src/mongodb/helpers.py,sha256=H0GpOK3bPBhFWBEhJZOjywUBdzih6MOpmyVO_cKSN14,24178
 ingestr/src/notion/__init__.py,sha256=36wUui8finbc85ObkRMq8boMraXMUehdABN_AMe_hzA,1834
 ingestr/src/notion/settings.py,sha256=MwQVZViJtnvOegfjXYc_pJ50oUYgSRPgwqu7TvpeMOA,82
 ingestr/src/notion/helpers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -128,8 +131,8 @@ ingestr/testdata/delete_insert_part2.csv,sha256=B_KUzpzbNdDY_n7wWop1mT2cz36TmayS
 ingestr/testdata/merge_expected.csv,sha256=DReHqWGnQMsf2PBv_Q2pfjsgvikYFnf1zYcQZ7ZqYN0,276
 ingestr/testdata/merge_part1.csv,sha256=Pw8Z9IDKcNU0qQHx1z6BUf4rF_-SxKGFOvymCt4OY9I,185
 ingestr/testdata/merge_part2.csv,sha256=T_GiWxA81SN63_tMOIuemcvboEFeAmbKc7xRXvL9esw,287
-ingestr-0.13.41.dist-info/METADATA,sha256=UzGBs9s0Kr6R1xji_ULG5Tuc383Klx2AIzfyZdXLBp4,13852
-ingestr-0.13.41.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ingestr-0.13.41.dist-info/entry_points.txt,sha256=oPJy0KBnPWYjDtP1k8qwAihcTLHSZokSQvRAw_wtfJM,46
-ingestr-0.13.41.dist-info/licenses/LICENSE.md,sha256=cW8wIhn8HFE-KLStDF9jHQ1O_ARWP3kTpk_-eOccL24,1075
-ingestr-0.13.41.dist-info/RECORD,,
+ingestr-0.13.43.dist-info/METADATA,sha256=n3hm9A9kfuALzOs9t2f4V2X4LH3hcIyFnj5y0DYP9Qo,13852
+ingestr-0.13.43.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ingestr-0.13.43.dist-info/entry_points.txt,sha256=oPJy0KBnPWYjDtP1k8qwAihcTLHSZokSQvRAw_wtfJM,46
+ingestr-0.13.43.dist-info/licenses/LICENSE.md,sha256=cW8wIhn8HFE-KLStDF9jHQ1O_ARWP3kTpk_-eOccL24,1075
+ingestr-0.13.43.dist-info/RECORD,,