ingestr 0.13.75__py3-none-any.whl → 0.14.98__py3-none-any.whl

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,
+        )

ingestr/src/destinations.py

@@ -19,12 +19,26 @@ from dlt.destinations.impl.clickhouse.configuration import (
     ClickHouseCredentials,
 )
 
+from ingestr.src.elasticsearch.helpers import elasticsearch_insert
 from ingestr.src.errors import MissingValueError
 from ingestr.src.loader import load_dlt_file
+from ingestr.src.mongodb.helpers import mongodb_insert
 
 
 class GenericSqlDestination:
     def dlt_run_params(self, uri: str, table: str, **kwargs) -> dict:
+        if uri.startswith("databricks://"):
+            p = urlparse(uri)
+            q = parse_qs(p.query)
+            schema = q.get("schema", [None])[0]
+            if not schema:
+                raise ValueError("Databricks requires schema in the URI.")
+            res = {
+                "dataset_name": schema,
+                "table_name": table,
+            }
+            return res
+
         table_fields = table.split(".")
         if len(table_fields) != 2:
             raise ValueError("Table name must be in the format <schema>.<table>")
@@ -147,6 +161,24 @@ class DuckDBDestination(GenericSqlDestination):
         return dlt.destinations.duckdb(uri, **kwargs)
 
 
+class MotherduckDestination(GenericSqlDestination):
+    def dlt_dest(self, uri: str, **kwargs):
+        from urllib.parse import parse_qs, urlparse
+
+        parsed = urlparse(uri)
+        query = parse_qs(parsed.query)
+        token = query.get("token", [None])[0]
+        from dlt.destinations.impl.motherduck.configuration import MotherDuckCredentials
+
+        creds = {
+            "password": token,
+        }
+        if parsed.path.lstrip("/"):
+            creds["database"] = parsed.path.lstrip("/")
+
+        return dlt.destinations.motherduck(MotherDuckCredentials(creds), **kwargs)
+
+
 def handle_datetimeoffset(dto_value: bytes) -> datetime.datetime:
     # ref: https://github.com/mkleehammer/pyodbc/issues/134#issuecomment-281739794
     tup = struct.unpack(
@@ -252,7 +284,26 @@ class MsSQLDestination(GenericSqlDestination):
 
 class DatabricksDestination(GenericSqlDestination):
     def dlt_dest(self, uri: str, **kwargs):
-        return dlt.destinations.databricks(credentials=uri, **kwargs)
+        p = urlparse(uri)
+        q = parse_qs(p.query)
+        access_token = p.password
+        server_hostname = p.hostname
+        http_path = q.get("http_path", [None])[0]
+        catalog = q.get("catalog", [None])[0]
+        schema = q.get("schema", [None])[0]
+
+        creds = {
+            "access_token": access_token,
+            "server_hostname": server_hostname,
+            "http_path": http_path,
+            "catalog": catalog,
+            "schema": schema,
+        }
+
+        return dlt.destinations.databricks(
+            credentials=creds,
+            **kwargs,
+        )
 
 
 class SynapseDestination(GenericSqlDestination):
@@ -538,6 +589,76 @@ class MySqlDestination(GenericSqlDestination):
         }
 
 
+class TrinoTypeMapper:
+    """Custom type mapper for Trino to handle unsupported types."""
+
+    @staticmethod
+    def create_type_mapper():
+        """Create a custom type mapper for Trino."""
+        from dlt.destinations.impl.sqlalchemy.type_mapper import SqlalchemyTypeMapper
+        from sqlalchemy import BigInteger, Text
+        from sqlalchemy.sql import sqltypes
+
+        class CustomTrinoTypeMapper(SqlalchemyTypeMapper):
+            """Custom type mapper that converts unsupported Trino types."""
+
+            def to_destination_type(self, column, table=None):
+                # Handle special cases before calling parent
+                data_type = column.get("data_type", "")
+
+                # Convert JSON to VARCHAR for Trino's Iceberg catalog
+                if data_type == "json":
+                    # Use TEXT (unlimited VARCHAR) for JSON data
+                    return Text()
+
+                # Convert BINARY to VARCHAR
+                if data_type == "binary":
+                    return Text()
+
+                # Handle integer types - always use BIGINT for Trino
+                # Note: dlt uses "bigint" internally, not "integer"
+                if data_type in ["bigint", "integer", "int"]:
+                    return BigInteger()
+
+                # For other types, try parent mapper
+                try:
+                    type_ = super().to_destination_type(column, table)
+                except Exception:
+                    # If parent can't handle it, default to TEXT
+                    return Text()
+
+                # Convert any INTEGER type to BIGINT
+                if isinstance(type_, sqltypes.Integer) and not isinstance(
+                    type_, sqltypes.BigInteger
+                ):
+                    return BigInteger()
+
+                # Ensure VARCHAR types don't have constraints that Trino doesn't support
+                if isinstance(type_, sqltypes.String):
+                    # Return TEXT for unlimited string
+                    return Text()
+
+                return type_
+
+        return CustomTrinoTypeMapper
+
+
+class TrinoDestination(GenericSqlDestination):
+    def dlt_dest(self, uri: str, **kwargs):
+        # Import required modules
+        from dlt.destinations.impl.sqlalchemy.factory import (
+            sqlalchemy as sqlalchemy_factory,
+        )
+
+        # Create the destination with custom type mapper
+        # We need to use the factory to properly configure the type mapper
+        dest = sqlalchemy_factory(
+            credentials=uri, type_mapper=TrinoTypeMapper.create_type_mapper(), **kwargs
+        )
+
+        return dest
+
+
 class BlobStorageDestination(abc.ABC):
     @abc.abstractmethod
     def credentials(self, params: dict) -> FileSystemCredentials:
@@ -654,3 +775,50 @@ class GCSDestination(BlobStorageDestination):
             credentials = json.loads(base64.b64decode(credentials_base64[0]).decode())  # type: ignore
 
         return credentials
+
+
+class ElasticsearchDestination:
+    def dlt_dest(self, uri: str, **kwargs):
+        from urllib.parse import urlparse
+
+        parsed_uri = urlparse(uri)
+
+        # Extract connection details from URI
+        scheme = parsed_uri.scheme or "http"
+        host = parsed_uri.hostname or "localhost"
+        port = parsed_uri.port or 9200
+        username = parsed_uri.username
+        password = parsed_uri.password
+
+        # Build connection string
+        if username and password:
+            connection_string = f"{scheme}://{username}:{password}@{host}:{port}"
+        else:
+            connection_string = f"{scheme}://{host}:{port}"
+
+        # Add query parameters if any
+        if parsed_uri.query:
+            connection_string += f"?{parsed_uri.query}"
+
+        return elasticsearch_insert(connection_string=connection_string)
+
+    def dlt_run_params(self, uri: str, table: str, **kwargs) -> dict:
+        return {
+            "table_name": table,
+        }
+
+    def post_load(self):
+        pass
+
+
+class MongoDBDestination:
+    def dlt_dest(self, uri: str, **kwargs):
+        return mongodb_insert(uri)
+
+    def dlt_run_params(self, uri: str, table: str, **kwargs) -> dict:
+        return {
+            "table_name": table,
+        }
+
+    def post_load(self):
+        pass