omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/.gitignore

@@ -0,0 +1,10 @@
+# ignore secrets, virtual environments and typical python compilation artifacts
+secrets.toml
+# ignore basic python artifacts
+.env
+**/__pycache__/
+**/*.py[cod]
+**/*$py.class
+# ignore duckdb
+*.duckdb
+*.wal

omniload/src/adjust/__init__.py

@@ -0,0 +1,108 @@
+from typing import Optional, Sequence
+
+import dlt
+import pendulum
+from dlt.sources import DltResource
+
+from .adjust_helpers import DEFAULT_DIMENSIONS, DEFAULT_METRICS, AdjustAPI
+
+REQUIRED_CUSTOM_DIMENSIONS = [
+    "hour",
+    "day",
+    "week",
+    "month",
+    "quarter",
+    "year",
+]
+KNOWN_TYPE_HINTS = {
+    "hour": {"data_type": "timestamp"},
+    "day": {"data_type": "date"},
+    "week": {"data_type": "text"},
+    "month": {"data_type": "text"},
+    "quarter": {"data_type": "text"},
+    "year": {"data_type": "text"},
+    "campaign": {"data_type": "text"},
+    "adgroup": {"data_type": "text"},
+    "creative": {"data_type": "text"},
+    # metrics
+    "installs": {"data_type": "bigint"},
+    "clicks": {"data_type": "bigint"},
+    "cost": {"data_type": "decimal"},
+    "network_cost": {"data_type": "decimal"},
+    "impressions": {"data_type": "bigint"},
+    "ad_revenue": {"data_type": "decimal"},
+    "all_revenue": {"data_type": "decimal"},
+}
+
+
+@dlt.source(max_table_nesting=0)
+def adjust_source(
+    start_date: pendulum.DateTime,
+    end_date: pendulum.DateTime,
+    api_key: str,
+    dimensions: Optional[list[str]] = None,
+    metrics: Optional[list[str]] = None,
+    merge_key: Optional[str] = None,
+    filters: Optional[dict] = None,
+) -> Sequence[DltResource]:
+    @dlt.resource(write_disposition="merge", merge_key="day")
+    def campaigns() -> DltResource:
+        adjust_api = AdjustAPI(api_key=api_key)
+        yield from adjust_api.fetch_report_data(
+            start_date=start_date,
+            end_date=end_date,
+            dimensions=DEFAULT_DIMENSIONS,
+            metrics=DEFAULT_METRICS,
+            filters=filters,
+        )
+
+    @dlt.resource(write_disposition="replace", primary_key="id")
+    def events() -> DltResource:
+        adjust_api = AdjustAPI(api_key=api_key)
+        yield adjust_api.fetch_events()
+
+    @dlt.resource(write_disposition="merge", merge_key="day")
+    def creatives() -> DltResource:
+        adjust_api = AdjustAPI(api_key=api_key)
+        yield from adjust_api.fetch_report_data(
+            start_date=start_date,
+            end_date=end_date,
+            dimensions=DEFAULT_DIMENSIONS + ["adgroup", "creative"],
+            metrics=DEFAULT_METRICS,
+            filters=filters,
+        )
+
+    if not dimensions:
+        return campaigns, creatives, events
+
+    merge_key = merge_key
+    type_hints = {}
+    for dimension in REQUIRED_CUSTOM_DIMENSIONS:
+        if dimension in dimensions:
+            merge_key = dimension
+            break
+
+    for dimension in dimensions:
+        if dimension in KNOWN_TYPE_HINTS:
+            type_hints[dimension] = KNOWN_TYPE_HINTS[dimension]
+    for metric in metrics:
+        if metric in KNOWN_TYPE_HINTS:
+            type_hints[metric] = KNOWN_TYPE_HINTS[metric]
+
+    @dlt.resource(
+        write_disposition={"disposition": "merge", "strategy": "delete-insert"},
+        merge_key=merge_key,
+        primary_key=dimensions,
+        columns=type_hints,
+    )
+    def custom() -> DltResource:
+        adjust_api = AdjustAPI(api_key=api_key)
+        yield from adjust_api.fetch_report_data(
+            start_date=start_date,
+            end_date=end_date,
+            dimensions=dimensions,
+            metrics=metrics,
+            filters=filters,
+        )
+
+    return campaigns, creatives, custom, events

omniload/src/adjust/adjust_helpers.py

@@ -0,0 +1,122 @@
+from typing import Optional
+
+import pendulum
+import requests
+from dlt.sources.helpers.requests import Client
+from requests.exceptions import HTTPError
+
+DEFAULT_DIMENSIONS = ["campaign", "day", "app", "store_type", "channel", "country"]
+
+DEFAULT_METRICS = [
+    "installs",
+    "network_cost",
+    "all_revenue_total_d0",
+    "ad_revenue_total_d0",
+    "revenue_total_d0",
+    "all_revenue_total_d1",
+    "ad_revenue_total_d1",
+    "revenue_total_d1",
+    "all_revenue_total_d3",
+    "ad_revenue_total_d3",
+    "revenue_total_d3",
+    "all_revenue_total_d7",
+    "ad_revenue_total_d7",
+    "revenue_total_d7",
+    "all_revenue_total_d14",
+    "ad_revenue_total_d14",
+    "revenue_total_d14",
+    "all_revenue_total_d21",
+]
+
+
+def retry_on_limit(response: requests.Response, exception: BaseException) -> bool:
+    return response.status_code == 429
+
+
+class AdjustAPI:
+    def __init__(self, api_key):
+        self.api_key = api_key
+        self.request_client = Client(
+            request_timeout=1000,  # Adjust support recommends 1000 seconds of read timeout.
+            raise_for_status=False,
+            retry_condition=retry_on_limit,
+            request_max_attempts=12,
+            request_backoff_factor=2,
+        ).session
+
+    def fetch_report_data(
+        self,
+        start_date: pendulum.DateTime,
+        end_date: pendulum.DateTime,
+        dimensions=DEFAULT_DIMENSIONS,
+        metrics=DEFAULT_METRICS,
+        filters: Optional[dict] = None,
+    ):
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        params = {}
+
+        if filters:
+            for key, value in filters.items():
+                if isinstance(value, list):
+                    params[key] = ",".join(value)
+                else:
+                    params[key] = value
+
+        params["date_period"] = (
+            f"{start_date.format('YYYY-MM-DD')}:{end_date.format('YYYY-MM-DD')}"
+        )
+        params["dimensions"] = ",".join(dimensions)
+        params["metrics"] = ",".join(metrics)
+
+        if start_date > end_date:
+            raise ValueError(
+                f"Invalid date range: Start date ({start_date}) must be earlier than end date ({end_date})."
+            )
+
+        response = self.request_client.get(
+            "https://automate.adjust.com/reports-service/report",
+            headers=headers,
+            params=params,
+        )
+        if response.status_code == 200:
+            result = response.json()
+            items = result.get("rows", [])
+            yield items
+        else:
+            raise HTTPError(
+                f"Request failed with status code: {response.status_code}, {response.text}."
+            )
+
+    def fetch_events(self):
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        response = self.request_client.get(
+            "https://automate.adjust.com/reports-service/events", headers=headers
+        )
+        if response.status_code == 200:
+            result = response.json()
+            yield result
+        else:
+            raise HTTPError(
+                f"Request failed with status code: {response.status_code}, {response.text}."
+            )
+
+
+def parse_filters(filters_raw: str) -> dict:
+    # Parse filter string like "key1=value1,key2=value2,value3,value4"
+    filters = {}
+    current_key = None
+
+    for item in filters_raw.split(","):
+        if "=" in item:
+            # Start of a new key-value pair
+            key, value = item.split("=")
+            filters[key] = [value]  # Always start with a list
+            current_key = key
+        elif current_key is not None:
+            # Additional value for the current key
+            filters[current_key].append(item)
+
+    # Convert single-item lists to simple values
+    filters = {k: v[0] if len(v) == 1 else v for k, v in filters.items()}
+
+    return filters

omniload/src/airtable/__init__.py

@@ -0,0 +1,84 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Source that loads tables form Airtable.
+Supports whitelisting of tables or loading of all tables from a specified base.
+"""
+
+from typing import Any, Dict, Iterable, Iterator, List, Optional
+
+import dlt
+import pyairtable
+from dlt.sources import DltResource
+
+
+@dlt.source(max_table_nesting=1)
+def airtable_source(
+    base_id: str = dlt.config.value,
+    table_names: Optional[List[str]] = dlt.config.value,
+    access_token: str = dlt.secrets.value,
+) -> Iterable[DltResource]:
+    """
+    Represents tables for a single Airtable base.
+    Args:
+        base_id (str): The id of the base. Obtain it e.g. from the URL in your webbrowser.
+            It starts with "app". See https://support.airtable.com/docs/finding-airtable-ids
+        table_names (Optional[List[str]]): A list of table IDs or table names to load.
+            Unless specified otherwise, all tables in the schema are loaded.
+            Names are freely user-defined. IDs start with "tbl". See https://support.airtable.com/docs/finding-airtable-ids
+        access_token (str): The personal access token.
+            See https://support.airtable.com/docs/creating-and-using-api-keys-and-access-tokens#personal-access-tokens-basic-actions
+    """
+    api = pyairtable.Api(access_token)
+    all_tables_url = api.build_url(f"meta/bases/{base_id}/tables")
+    tables = api.request(method="GET", url=all_tables_url).get("tables")
+    for t in tables:
+        if table_names:
+            if t.get("id") in table_names or t.get("name") in table_names:
+                yield airtable_resource(api, base_id, t)
+        else:
+            yield airtable_resource(api, base_id, t)
+
+
+def airtable_resource(
+    api: pyairtable.Api,
+    base_id: str,
+    table: Dict[str, Any],
+) -> DltResource:
+    """
+    Represents a single airtable.
+    Args:
+        api (pyairtable.Api): The API connection object
+        base_id (str): The id of the base. Obtain it e.g. from the URL in your webbrowser.
+            It starts with "app". See https://support.airtable.com/docs/finding-airtable-ids
+        table (Dict[str, Any]): Metadata about an airtable, does not contain the actual records
+    """
+
+    primary_key_id = table["primaryFieldId"]
+    primary_key_field = [
+        field for field in table["fields"] if field["id"] == primary_key_id
+    ][0]
+    table_name: str = table["name"]
+    primary_key: List[str] = [f"fields__{primary_key_field['name']}".lower()]
+    air_table = api.table(base_id, table["id"])
+
+    # Table.iterate() supports rich customization options, such as chunk size, fields, cell format, timezone, locale, and view
+    air_table_generator: Iterator[List[Any]] = air_table.iterate()
+
+    return dlt.resource(
+        air_table_generator,
+        name=table_name,
+        primary_key=primary_key,
+        write_disposition="replace",
+    )

omniload/src/allium/__init__.py

@@ -0,0 +1,128 @@
+"""
+Allium source for data extraction via REST API.
+
+This source provides access to Allium blockchain data via asynchronous query execution.
+"""
+
+import time
+from typing import Any, Iterator
+
+import dlt
+
+from omniload.src.http_client import create_client
+
+
+@dlt.source(max_table_nesting=0, name="allium_source")
+def allium_source(
+    api_key: str,
+    query_id: str,
+    parameters: dict[str, Any] | None = None,
+    limit: int | None = None,
+    compute_profile: str | None = None,
+) -> Any:
+    """
+    Allium data source for blockchain data extraction.
+
+    This source connects to Allium API, runs async queries, and fetches results.
+
+    Args:
+        api_key: Allium API key for authentication
+        query_id: The query ID to execute (e.g., 'abc123')
+        parameters: Optional parameters for the query (e.g., {'start_date': '2025-02-01', 'end_date': '2025-02-02'})
+        limit: Limit the number of rows in the result (max 250,000)
+        compute_profile: Compute profile identifier
+
+    Yields:
+        DltResource: Data resources for Allium query results
+    """
+    base_url = "https://api.allium.so/api/v1/explorer"
+    session = create_client()
+    headers = {"X-API-Key": api_key}
+
+    @dlt.resource(
+        name="query_results",
+        write_disposition="replace",
+    )
+    def fetch_query_results() -> Iterator[dict[str, Any]]:
+        """
+        Fetch query results from Allium.
+
+        This function:
+        1. Starts an async query execution
+        2. Polls for completion status
+        3. Fetches and yields the results
+        """
+        # Step 1: Start async query execution
+        run_config: dict[str, Any] = {}
+        if limit is not None:
+            run_config["limit"] = limit
+        if compute_profile is not None:
+            run_config["compute_profile"] = compute_profile
+
+        run_payload = {"parameters": parameters or {}, "run_config": run_config}
+
+        run_response = session.post(
+            f"{base_url}/queries/{query_id}/run-async",
+            json=run_payload,
+            headers=headers,
+        )
+
+        run_data = run_response.json()
+
+        if "run_id" not in run_data:
+            raise ValueError(f"Failed to start query execution: {run_data}")
+
+        run_id = run_data["run_id"]
+
+        # Step 2: Poll for completion
+        max_retries = 8640  # Max 12 hours with 5-second intervals
+        retry_count = 0
+        poll_interval = 5  # seconds
+
+        while retry_count < max_retries:
+            status_response = session.get(
+                f"{base_url}/query-runs/{run_id}/status",
+                headers=headers,
+            )
+            status_response.raise_for_status()
+            status_data = status_response.json()
+
+            # Handle both string and dict responses
+            if isinstance(status_data, str):
+                status = status_data
+            else:
+                status = status_data.get("status")
+
+            if status == "success":
+                break
+            elif status == "failed":
+                error_msg = (
+                    status_data.get("error", "Unknown error")
+                    if isinstance(status_data, dict)
+                    else "Unknown error"
+                )
+                raise ValueError(f"Query execution failed: {error_msg}")
+            elif status in ["pending", "running", "queued"]:
+                time.sleep(poll_interval)
+                retry_count += 1
+            else:
+                raise ValueError(f"Unknown status: {status}")
+
+        if retry_count >= max_retries:
+            raise TimeoutError(
+                f"Query execution timed out after {max_retries * poll_interval} seconds"
+            )
+
+        # Step 3: Fetch results
+        results_response = session.get(
+            f"{base_url}/query-runs/{run_id}/results",
+            headers=headers,
+            params={"f": "json"},
+        )
+        results_response.raise_for_status()
+        query_output = results_response.json()
+
+        # Extract and yield all data
+        yield query_output.get("data", [])
+
+    return (fetch_query_results,)