ingestr 0.13.2__py3-none-any.whl → 0.14.104__py3-none-any.whl

ingestr/src/couchbase_source/__init__.py

@@ -0,0 +1,118 @@
+"""Source that loads data from Couchbase buckets, supports incremental loads."""
+
+from typing import Optional
+
+import dlt
+from dlt.sources import DltResource
+
+from .helpers import (
+    CouchbaseConfiguration,
+    client_from_credentials,
+    fetch_documents,
+)
+
+
+@dlt.source(max_table_nesting=0)
+def couchbase_source(
+    connection_string: str = dlt.secrets.value,
+    username: str = dlt.secrets.value,
+    password: str = dlt.secrets.value,
+    bucket: str = dlt.config.value,
+    scope: Optional[str] = dlt.config.value,
+    collection: Optional[str] = dlt.config.value,
+    incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
+    write_disposition: Optional[str] = dlt.config.value,
+    limit: Optional[int] = None,
+) -> DltResource:
+    """
+    A DLT source which loads data from a Couchbase bucket using Couchbase Python SDK.
+
+    Args:
+        connection_string (str): Couchbase connection string (e.g., 'couchbase://localhost')
+        username (str): Couchbase username
+        password (str): Couchbase password
+        bucket (str): Bucket name to load data from
+        scope (Optional[str]): Scope name (defaults to '_default')
+        collection (Optional[str]): Collection name (defaults to '_default')
+        incremental (Optional[dlt.sources.incremental]): Option to enable incremental loading.
+            E.g., `incremental=dlt.sources.incremental('updated_at', pendulum.parse('2022-01-01T00:00:00Z'))`
+        write_disposition (str): Write disposition of the resource.
+        limit (Optional[int]): The maximum number of documents to load.
+
+    Returns:
+        DltResource: A DLT resource for the Couchbase collection.
+    """
+    # Set up Couchbase client
+    cluster = client_from_credentials(connection_string, username, password)
+
+    resource_name = f"{bucket}_{scope}_{collection}"
+
+    return dlt.resource(  # type: ignore[call-overload, arg-type]
+        fetch_documents,
+        name=resource_name,
+        primary_key="id",
+        write_disposition=write_disposition or "replace",
+        spec=CouchbaseConfiguration,
+        max_table_nesting=0,
+    )(
+        cluster=cluster,
+        bucket_name=bucket,
+        scope_name=scope,
+        collection_name=collection,
+        incremental=incremental,
+        limit=limit,
+    )
+
+
+@dlt.resource(
+    name=lambda args: f"{args['bucket']}_{args['scope']}_{args['collection']}",
+    standalone=True,
+    spec=CouchbaseConfiguration,  # type: ignore[arg-type]
+)
+def couchbase_collection(
+    connection_string: str = dlt.secrets.value,
+    username: str = dlt.secrets.value,
+    password: str = dlt.secrets.value,
+    bucket: str = dlt.config.value,
+    scope: Optional[str] = dlt.config.value,
+    collection: Optional[str] = dlt.config.value,
+    incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
+    write_disposition: Optional[str] = dlt.config.value,
+    limit: Optional[int] = None,
+    chunk_size: Optional[int] = 1000,
+) -> DltResource:
+    """
+    A DLT resource which loads a collection from Couchbase.
+
+    Args:
+        connection_string (str): Couchbase connection string (e.g., 'couchbase://localhost')
+        username (str): Couchbase username
+        password (str): Couchbase password
+        bucket (str): Bucket name to load data from
+        scope (Optional[str]): Scope name (defaults to '_default')
+        collection (Optional[str]): Collection name (defaults to '_default')
+        incremental (Optional[dlt.sources.incremental]): Option to enable incremental loading.
+        write_disposition (str): Write disposition of the resource.
+        limit (Optional[int]): The maximum number of documents to load.
+        chunk_size (Optional[int]): The number of documents to load in each batch.
+
+    Returns:
+        DltResource: A DLT resource for the Couchbase collection.
+    """
+    # Set up Couchbase client
+    cluster = client_from_credentials(connection_string, username, password)
+
+    return dlt.resource(  # type: ignore[call-overload]
+        fetch_documents,
+        name=f"{bucket}_{scope}_{collection}",
+        primary_key="id",
+        write_disposition=write_disposition or "replace",
+    )(
+        cluster=cluster,
+        bucket_name=bucket,
+        scope_name=scope,
+        collection_name=collection,
+        incremental=incremental,
+        limit=limit,
+        chunk_size=chunk_size,
+    )

ingestr/src/couchbase_source/helpers.py

@@ -0,0 +1,135 @@
+"""Helper functions for Couchbase source."""
+
+from datetime import datetime, timedelta
+from typing import Any, Dict, Iterator, Optional
+
+import dlt
+from couchbase.auth import PasswordAuthenticator  # type: ignore[import-untyped]
+from couchbase.cluster import Cluster  # type: ignore[import-untyped]
+from couchbase.options import (  # type: ignore[import-untyped]
+    ClusterOptions,
+    QueryOptions,
+)
+from dlt.common.configuration import configspec
+from dlt.common.time import ensure_pendulum_datetime
+
+
+@configspec
+class CouchbaseConfiguration:
+    """Configuration for Couchbase source."""
+
+    connection_string: str = dlt.secrets.value
+    username: str = dlt.secrets.value
+    password: str = dlt.secrets.value
+    bucket: str = dlt.config.value
+    scope: Optional[str] = dlt.config.value
+    collection: Optional[str] = dlt.config.value
+
+
+def client_from_credentials(
+    connection_string: str, username: str, password: str
+) -> Cluster:
+    """
+    Create a Couchbase cluster client from credentials.
+
+    Args:
+        connection_string: Couchbase connection string
+            - Local/self-hosted: 'couchbase://localhost'
+            - Capella (cloud): 'couchbases://your-instance.cloud.couchbase.com'
+        username: Couchbase username
+        password: Couchbase password
+
+    Returns:
+        Cluster: Connected Couchbase cluster instance
+    """
+    auth = PasswordAuthenticator(username, password)
+    options = ClusterOptions(auth)
+
+    # Apply wan_development profile for Capella (couchbases://) connections
+    # This helps avoid latency issues when accessing from different networks
+    if connection_string.startswith("couchbases://"):
+        options.apply_profile("wan_development")
+
+    cluster = Cluster(connection_string, options)
+    cluster.wait_until_ready(timedelta(seconds=30))
+
+    return cluster
+
+
+def fetch_documents(
+    cluster: Cluster,
+    bucket_name: str,
+    scope_name: str,
+    collection_name: str,
+    incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
+    limit: Optional[int] = None,
+    chunk_size: Optional[int] = 1000,
+) -> Iterator[Dict[str, Any]]:
+    """
+    Fetch documents from a Couchbase collection using N1QL queries.
+
+    Args:
+        cluster: Couchbase cluster instance
+        bucket_name: Name of the bucket
+        scope_name: Name of the scope
+        collection_name: Name of the collection
+        incremental: Incremental loading configuration
+        limit: Maximum number of documents to fetch
+        chunk_size: Number of documents to fetch per batch
+
+    Yields:
+        Dict[str, Any]: Document data
+    """
+    # Build N1QL query with full path
+    full_collection_path = f"`{bucket_name}`.`{scope_name}`.`{collection_name}`"
+    n1ql_query = f"SELECT META().id as id, c.* FROM {full_collection_path} c"
+
+    # Add incremental filter if provided
+    if incremental and incremental.cursor_path:
+        where_clause = f" WHERE {incremental.cursor_path} >= $start_value"
+        if incremental.end_value is not None:
+            where_clause += f" AND {incremental.cursor_path} < $end_value"
+        n1ql_query += where_clause
+
+    # Add limit if provided
+    if limit:
+        n1ql_query += f" LIMIT {limit}"
+
+    # Execute query
+    try:
+        query_options = QueryOptions()
+
+        # Add parameters if incremental
+        if incremental and incremental.cursor_path:
+            named_parameters = {"start_value": incremental.last_value}
+            if incremental.end_value is not None:
+                named_parameters["end_value"] = incremental.end_value
+            query_options = QueryOptions(named_parameters=named_parameters)
+
+        result = cluster.query(n1ql_query, query_options)
+
+        # Yield documents
+        count = 0
+        for row in result:
+            doc = dict(row)
+
+            # Convert datetime fields to proper format
+            if (
+                incremental
+                and incremental.cursor_path
+                and incremental.cursor_path in doc
+            ):
+                cursor_value = doc[incremental.cursor_path]
+                if isinstance(cursor_value, (str, datetime)):
+                    doc[incremental.cursor_path] = ensure_pendulum_datetime(
+                        cursor_value
+                    )
+
+            yield doc
+
+            count += 1
+            if limit and count >= limit:
+                break
+
+    except Exception as e:
+        raise Exception(f"Error executing Couchbase N1QL query: {str(e)}")

ingestr/src/cursor/__init__.py

@@ -0,0 +1,83 @@
+"""
+This source provides data extraction from Cursor via the REST API.
+
+It fetches team member information from the Cursor API.
+"""
+
+from typing import Any, Iterable, Optional
+
+import dlt
+from dlt.common.typing import TDataItem
+
+from .helpers import CursorClient
+
+
+@dlt.source
+def cursor_source() -> Any:
+    """
+    The main function that fetches data from Cursor API.
+
+    Returns:
+        Sequence[DltResource]: A sequence of DltResource objects containing the fetched data.
+    """
+    return [
+        team_members,
+        daily_usage_data,
+        team_spend,
+        filtered_usage_events,
+    ]
+
+
+@dlt.resource(
+    write_disposition="replace",
+    max_table_nesting=0,
+)
+def team_members(
+    api_key: str = dlt.secrets.value,
+) -> Iterable[TDataItem]:
+    client = CursorClient(api_key=api_key)
+
+    members = client.get_team_members()
+    yield from members
+
+
+@dlt.resource(
+    write_disposition="replace",
+    max_table_nesting=0,
+)
+def daily_usage_data(
+    api_key: str = dlt.secrets.value,
+    start_date: Optional[int] = dlt.config.value,
+    end_date: Optional[int] = dlt.config.value,
+) -> Iterable[TDataItem]:
+    client = CursorClient(api_key=api_key)
+
+    yield from client.get_daily_usage_data(start_date=start_date, end_date=end_date)
+
+
+@dlt.resource(
+    write_disposition="replace",
+    max_table_nesting=0,
+)
+def team_spend(
+    api_key: str = dlt.secrets.value,
+) -> Iterable[TDataItem]:
+    client = CursorClient(api_key=api_key)
+
+    yield from client.get_team_spend()
+
+
+@dlt.resource(
+    write_disposition="replace",
+    max_table_nesting=0,
+)
+def filtered_usage_events(
+    api_key: str = dlt.secrets.value,
+    start_date: Optional[int] = dlt.config.value,
+    end_date: Optional[int] = dlt.config.value,
+) -> Iterable[TDataItem]:
+    client = CursorClient(api_key=api_key)
+
+    yield from client.get_filtered_usage_events(
+        start_date=start_date, end_date=end_date
+    )

ingestr/src/cursor/helpers.py

@@ -0,0 +1,188 @@
+"""Cursor source helpers"""
+
+from typing import Any, Callable, Dict, Iterator, List, Optional
+
+import requests
+
+REQUEST_TIMEOUT = 30
+
+
+class CursorClient:
+    """Cursor REST API client with API key authentication."""
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.cursor.com",
+        timeout: int = REQUEST_TIMEOUT,
+    ):
+        """
+        Initialize Cursor client with API key authentication.
+
+        Args:
+            api_key: API key for authentication
+            base_url: Cursor API base URL
+            timeout: Request timeout in seconds
+        """
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.api_key = api_key
+
+    def _make_request(
+        self,
+        endpoint: str,
+        method: str = "POST",
+        json_data: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Make HTTP request to Cursor API.
+
+        Args:
+            endpoint: API endpoint path
+            method: HTTP method (default: POST)
+            json_data: JSON data for request body
+
+        Returns:
+            JSON response data
+        """
+        url = f"{self.base_url}/{endpoint.lstrip('/')}"
+
+        if json_data is not None:
+            response = requests.request(
+                method=method,
+                url=url,
+                auth=(self.api_key, ""),
+                timeout=self.timeout,
+                headers={"Content-Type": "application/json"},
+                json=json_data,
+            )
+        else:
+            response = requests.request(
+                method=method,
+                url=url,
+                auth=(self.api_key, ""),
+                timeout=self.timeout,
+                headers={"Content-Type": "application/json"},
+                json={},
+            )
+
+        response.raise_for_status()
+        return response.json()
+
+    def _paginate(
+        self,
+        endpoint: str,
+        data_key: str,
+        base_payload: Optional[Dict[str, Any]] = None,
+        page_size: Optional[int] = 100,
+        has_next_page_check: Optional[Callable[[Dict[str, Any]], bool]] = None,
+    ) -> Iterator[Dict[str, Any]]:
+        """
+        Generic pagination helper for API endpoints.
+
+        Args:
+            endpoint: API endpoint to call
+            data_key: Key in response containing the data array
+            base_payload: Base payload to include in each request
+            page_size: Number of results per page (default: 100)
+            has_next_page_check: Optional function to check if there's a next page from response
+
+        Yields:
+            Individual records from the paginated response
+        """
+        page = 1
+        base_payload = base_payload or {}
+
+        while True:
+            payload = base_payload.copy()
+
+            if page_size:
+                payload["pageSize"] = page_size
+                payload["page"] = page
+
+            result = self._make_request(endpoint, json_data=payload)
+            data = result.get(data_key, [])
+
+            if not data:
+                break
+
+            for record in data:
+                yield record
+
+            # If page_size is not set, we get all data in one request
+            if not page_size:
+                break
+
+            # Custom check for next page if provided
+            if has_next_page_check:
+                if not has_next_page_check(result):
+                    break
+            # Default: if we got less data than page_size, we've reached the end
+            elif len(data) < page_size:
+                break
+
+            page += 1
+
+    def get_team_members(self) -> List[Dict[str, Any]]:
+        response = self._make_request("teams/members", method="GET")
+        return response.get("teamMembers", [])
+
+    def get_daily_usage_data(
+        self,
+        start_date: Optional[int] = None,
+        end_date: Optional[int] = None,
+        page_size: Optional[int] = 100,
+    ) -> Iterator[Dict[str, Any]]:
+        base_payload = {}
+        if start_date is not None:
+            base_payload["startDate"] = start_date
+        if end_date is not None:
+            base_payload["endDate"] = end_date
+
+        yield from self._paginate(
+            endpoint="teams/daily-usage-data",
+            data_key="data",
+            base_payload=base_payload,
+            page_size=page_size,
+        )
+
+    def get_team_spend(
+        self,
+        page_size: Optional[int] = 100,
+    ) -> Iterator[Dict[str, Any]]:
+        def check_has_next_page(response: Dict[str, Any]) -> bool:
+            current_page = response.get("currentPage", 1)
+            total_pages = response.get("totalPages", 1)
+            return current_page < total_pages
+
+        yield from self._paginate(
+            endpoint="teams/spend",
+            data_key="teamMemberSpend",
+            page_size=page_size,
+            has_next_page_check=check_has_next_page,
+        )
+
+    def get_filtered_usage_events(
+        self,
+        start_date: Optional[int] = None,
+        end_date: Optional[int] = None,
+        page_size: Optional[int] = 100,
+    ) -> Iterator[Dict[str, Any]]:
+        base_payload = {}
+        if start_date is not None:
+            base_payload["startDate"] = start_date
+        if end_date is not None:
+            base_payload["endDate"] = end_date
+
+        # Custom check for hasNextPage
+        def check_has_next_page(response: Dict[str, Any]) -> bool:
+            pagination = response.get("pagination", {})
+            return pagination.get("hasNextPage", False)
+
+        yield from self._paginate(
+            endpoint="teams/filtered-usage-events",
+            data_key="usageEvents",
+            base_payload=base_payload,
+            page_size=page_size,
+            has_next_page_check=check_has_next_page,
+        )