omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/dune/helpers.py

@@ -0,0 +1,108 @@
+import time
+from typing import Any, Iterator
+
+BASE_URL = "https://api.dune.com/api/v1"
+
+
+def poll_execution(session, headers: dict, execution_id: str) -> None:
+    max_retries = 8640  # Max 12 hours with 5-second intervals
+    retry_count = 0
+    poll_interval = 5
+
+    while retry_count < max_retries:
+        status_response = session.get(
+            f"{BASE_URL}/execution/{execution_id}/status",
+            headers=headers,
+        )
+        status_response.raise_for_status()
+        status_data = status_response.json()
+        state = status_data.get("state")
+
+        if state == "QUERY_STATE_COMPLETED":
+            return
+        elif state == "QUERY_STATE_FAILED":
+            error = status_data.get("error", {})
+            error_msg = (
+                error.get("message", "Unknown error")
+                if isinstance(error, dict)
+                else str(error)
+            )
+            raise ValueError(f"Query execution failed: {error_msg}")
+        elif state in ("QUERY_STATE_PENDING", "QUERY_STATE_EXECUTING"):
+            time.sleep(poll_interval)
+            retry_count += 1
+        elif state == "QUERY_STATE_CANCELLED":
+            raise ValueError("Query execution was cancelled")
+        elif state == "QUERY_STATE_EXPIRED":
+            raise ValueError("Query execution expired")
+        else:
+            raise ValueError(f"Unknown query state: {state}")
+
+    raise TimeoutError(
+        f"Query execution timed out after {max_retries * poll_interval} seconds"
+    )
+
+
+def fetch_results(
+    session, headers: dict, execution_id: str
+) -> Iterator[dict[str, Any]]:
+    offset = 0
+    page_limit = 1000
+
+    while True:
+        params: dict[str, Any] = {
+            "limit": page_limit,
+            "offset": offset,
+        }
+
+        results_response = session.get(
+            f"{BASE_URL}/execution/{execution_id}/results",
+            headers=headers,
+            params=params,
+        )
+        results_response.raise_for_status()
+        results_data = results_response.json()
+
+        result = results_data.get("result", {})
+        rows = result.get("rows", [])
+
+        if not rows:
+            break
+
+        yield from rows
+
+        next_offset = results_data.get("next_offset")
+        if not next_offset:
+            break
+
+        offset = next_offset
+
+
+def fetch_queries(session, headers: dict) -> Iterator[dict[str, Any]]:
+    offset = 0
+    page_limit = 100
+
+    while True:
+        params: dict[str, Any] = {
+            "limit": page_limit,
+            "offset": offset,
+        }
+
+        response = session.get(
+            f"{BASE_URL}/queries",
+            headers=headers,
+            params=params,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        rows = data.get("queries", [])
+        if not rows:
+            break
+
+        yield from rows
+
+        total = data.get("total", 0)
+        offset += len(rows)
+        if offset >= total:
+            break

omniload/src/dynamodb/__init__.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import boto3
+import dlt
+from boto3.dynamodb.conditions import Attr
+from dlt.common.configuration.specs import AwsCredentials
+
+PAGINATION_KEY = "LastEvaluatedKey"
+FILTER_KEY = "FilterExpression"
+DATA_KEY = "Items"
+
+
+@dataclass
+class TableSchema:
+    primary_key: Optional[str]
+    sort_key: Optional[str]
+
+
+def parseSchema(table) -> TableSchema:
+    schema = TableSchema(None, None)
+    for key in table.key_schema:
+        match key["KeyType"]:
+            case "HASH":
+                schema.primary_key = key["AttributeName"]
+            case "RANGE":
+                schema.sort_key = key["AttributeName"]
+
+    if schema.primary_key is None:
+        raise ValueError(f"Table {table.name} has no primary key!")
+
+    return schema
+
+
+@dlt.source
+def dynamodb(
+    table_name: str,
+    credentials: AwsCredentials,
+    incremental: Optional[dlt.sources.incremental] = None,
+):
+    sesh = boto3.Session(
+        aws_access_key_id=credentials.aws_access_key_id,
+        aws_secret_access_key=credentials.aws_secret_access_key,
+        region_name=credentials.region_name,
+    )
+    db = sesh.resource("dynamodb", endpoint_url=credentials.endpoint_url)
+    table = db.Table(table_name)
+    schema = parseSchema(table)
+    resource = dlt.resource(
+        dynamodb_table,
+        primary_key=schema.primary_key,
+    )
+
+    yield resource(table, incremental)
+
+
+def dynamodb_table(
+    table,
+    incremental: Optional[dlt.sources.incremental] = None,
+):
+    args = build_scan_args(incremental)
+    scan = table.scan(**args)
+    while True:
+        yield from scan[DATA_KEY]
+        if PAGINATION_KEY not in scan:
+            break
+        scan = table.scan(ExclusiveStartKey=scan[PAGINATION_KEY], **args)
+
+
+def build_scan_args(
+    incremental: Optional[dlt.sources.incremental] = None,
+):
+    scan_args = {}
+
+    if incremental is None:
+        return scan_args
+
+    if incremental.last_value:
+        criteria = Attr(incremental.cursor_path).gte(incremental.last_value)
+        if incremental.end_value:
+            criteria = Attr(incremental.cursor_path).between(
+                incremental.last_value, incremental.end_value
+            )
+        scan_args[FILTER_KEY] = criteria
+
+    return scan_args

omniload/src/elasticsearch/__init__.py

@@ -0,0 +1,80 @@
+from datetime import date, datetime
+from typing import Any, Optional
+
+import dlt
+import pendulum
+from dlt.common.time import ensure_pendulum_datetime
+from pendulum import parse
+
+from elasticsearch import Elasticsearch
+
+
+@dlt.source
+def elasticsearch_source(
+    connection_url: str,
+    index: str,
+    verify_certs: bool,
+    incremental: Optional[dlt.sources.incremental] = None,
+):
+    client = Elasticsearch(connection_url, verify_certs=verify_certs)
+
+    @dlt.resource(
+        name=index, primary_key="id", write_disposition="merge", incremental=incremental
+    )
+    def get_documents(incremental=incremental):
+        body = {"query": {"match_all": {}}}
+
+        if incremental:
+            start_value = incremental.last_value
+            range_filter = {"gte": start_value}
+            if incremental.end_value is not None:
+                range_filter["lt"] = incremental.end_value
+            body = {"query": {"range": {incremental.cursor_path: range_filter}}}
+
+        page = client.search(index=index, scroll="5m", size=1000, body=body)
+
+        sid = page["_scroll_id"]
+        hits = page["hits"]["hits"]
+
+        if not hits:
+            return
+
+        # fetching first page (via .search)
+        for doc in hits:
+            doc_data = {"id": doc["_id"], **doc["_source"]}
+            if incremental:
+                doc_data[incremental.cursor_path] = convert_elasticsearch_objs(
+                    doc_data[incremental.cursor_path]
+                )
+            yield doc_data
+
+        while True:
+            # fetching page 2 and other pages (via .scroll)
+            page = client.scroll(scroll_id=sid, scroll="5m")
+            sid = page["_scroll_id"]
+            hits = page["hits"]["hits"]
+            if not hits:
+                break
+            for doc in hits:
+                doc_data = {"id": doc["_id"], **doc["_source"]}
+                if incremental:
+                    doc_data[incremental.cursor_path] = convert_elasticsearch_objs(
+                        doc_data[incremental.cursor_path]
+                    )
+                yield doc_data
+
+        client.clear_scroll(scroll_id=sid)
+
+    return get_documents
+
+
+def convert_elasticsearch_objs(value: Any) -> Any:
+    if isinstance(value, str):
+        parsed_date = parse(value, strict=False)
+        if parsed_date is not None:
+            if isinstance(
+                parsed_date,
+                (pendulum.DateTime, pendulum.Date, datetime, date, str, float, int),
+            ):
+                return ensure_pendulum_datetime(parsed_date)
+    return value

omniload/src/elasticsearch/helpers.py

@@ -0,0 +1,141 @@
+"""Elasticsearch destination helpers"""
+
+import json
+import logging
+from typing import Any, Dict, Iterator, Set
+from urllib.parse import urlparse
+
+import dlt
+
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import bulk
+
+# Suppress Elasticsearch transport logging
+logging.getLogger("elasticsearch.transport").setLevel(logging.WARNING)
+logging.getLogger("elastic_transport.transport").setLevel(logging.WARNING)
+
+_cleared_indices: Set[str] = set()
+
+
+def process_file_items(file_path: str) -> Iterator[Dict[str, Any]]:
+    """Process items from a file path (JSONL format)."""
+    with open(file_path, "r") as f:
+        for line in f:
+            if line.strip():
+                doc = json.loads(line.strip())
+                # Clean DLT metadata
+                cleaned_doc = {
+                    k: v for k, v in doc.items() if not k.startswith("_dlt_")
+                }
+                yield cleaned_doc
+
+
+def process_iterable_items(items: Any) -> Iterator[Dict[str, Any]]:
+    """Process items from an iterable."""
+    for item in items:
+        if isinstance(item, dict):
+            # Clean DLT metadata
+            cleaned_item = {k: v for k, v in item.items() if not k.startswith("_dlt_")}
+            yield cleaned_item
+
+
+@dlt.destination(
+    name="elasticsearch",
+    loader_file_format="typed-jsonl",
+    batch_size=1000,
+    naming_convention="snake_case",
+)
+def elasticsearch_insert(
+    items, table, connection_string: str = dlt.secrets.value
+) -> None:
+    """Insert data into Elasticsearch index.
+
+    Args:
+        items: Data items (file path or iterable)
+        table: Table metadata containing name and schema info
+        connection_string: Elasticsearch connection string
+    """
+    # Parse connection string
+    parsed = urlparse(connection_string)
+
+    # Build Elasticsearch client configuration
+    actual_url = connection_string
+    secure = True  # Default to HTTPS (secure by default)
+
+    if connection_string.startswith("elasticsearch://"):
+        actual_url = connection_string.replace("elasticsearch://", "")
+
+        # Parse to check for query parameters
+        temp_parsed = urlparse("http://" + actual_url)
+        from urllib.parse import parse_qs
+
+        query_params = parse_qs(temp_parsed.query)
+
+        # Check ?secure parameter (defaults to true)
+        if "secure" in query_params:
+            secure = query_params["secure"][0].lower() in ["true", "1", "yes"]
+
+        # Remove query params from URL for ES client
+        actual_url = actual_url.split("?")[0]
+
+        # Add scheme
+        scheme = "https" if secure else "http"
+        actual_url = f"{scheme}://{actual_url}"
+
+        parsed = urlparse(actual_url)
+
+    es_config: Dict[str, Any] = {
+        "hosts": [actual_url],
+        "verify_certs": secure,
+        "ssl_show_warn": False,
+    }
+
+    # Add authentication if present
+    if parsed.username and parsed.password:
+        es_config["http_auth"] = (parsed.username, parsed.password)
+
+    # Get index name from table metadata
+    index_name = table["name"]
+
+    # Connect to Elasticsearch
+    client = Elasticsearch(**es_config)
+
+    if index_name not in _cleared_indices:
+        if client.indices.exists(index=index_name):
+            client.indices.delete(index=index_name)
+        _cleared_indices.add(index_name)
+
+    # Process and insert documents
+    if isinstance(items, str):
+        documents = process_file_items(items)
+    else:
+        documents = process_iterable_items(items)
+
+    # Prepare documents for bulk insert as generator
+    def doc_generator():
+        for doc in documents:
+            es_doc: Dict[str, Any] = {"_index": index_name, "_source": doc.copy()}
+
+            # Use _id if present, otherwise let ES generate one
+            if "_id" in doc:
+                es_doc["_id"] = str(doc["_id"])
+                # Remove _id from source since it's metadata
+                if "_id" in es_doc["_source"]:
+                    del es_doc["_source"]["_id"]
+            elif "id" in doc:
+                es_doc["_id"] = str(doc["id"])
+
+            yield es_doc
+
+    # Bulk insert
+    try:
+        _, failed_items = bulk(client, doc_generator(), request_timeout=60)
+        if failed_items:
+            failed_count = (
+                len(failed_items) if isinstance(failed_items, list) else failed_items
+            )
+            raise Exception(
+                f"Failed to insert {failed_count} documents: {failed_items}"
+            )
+    except Exception as e:
+        raise Exception(f"Elasticsearch bulk insert failed: {str(e)}")

omniload/src/errors.py

@@ -0,0 +1,26 @@
+import requests
+
+
+class MissingValueError(Exception):
+    def __init__(self, value, source):
+        super().__init__(f"{value} is required to connect to {source}")
+
+
+class UnsupportedResourceError(Exception):
+    def __init__(self, resource, source):
+        super().__init__(
+            f"Resource '{resource}' is not supported for {source} source yet, if you are interested in it please create a GitHub issue at https://github.com/panodata/omniload"
+        )
+
+
+class InvalidBlobTableError(Exception):
+    def __init__(self, source):
+        super().__init__(
+            f"Invalid source table for {source} "
+            "Ensure that the table is in the format {bucket-name}/{file glob}"
+        )
+
+
+class HTTPError(Exception):
+    def __init__(self, source: requests.HTTPError):
+        super().__init__(f"HTTP {source.response.status_code}: {source.response.text}")