omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/asana_source/settings.py

@@ -0,0 +1,158 @@
+# 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.
+
+"""Asana source settings and constants"""
+
+# Default start date for Asana API requests, only tasks started after this date will be collected
+DEFAULT_START_DATE = "2010-01-01T00:00:00.000Z"
+
+# Asana API request timeout
+REQUEST_TIMEOUT = 300
+
+# list of workspace fields to be retrieved from Asana API
+WORKSPACE_FIELDS = ("gid", "name", "is_organization", "resource_type", "email_domains")
+
+# List of project fields to be retrieved from Asana API
+PROJECT_FIELDS = (
+    "name",
+    "gid",
+    "owner",
+    "current_status",
+    "custom_fields",
+    "default_view",
+    "due_date",
+    "due_on",
+    "is_template",
+    "created_at",
+    "modified_at",
+    "start_on",
+    "archived",
+    "public",
+    "members",
+    "followers",
+    "color",
+    "notes",
+    "icon",
+    "permalink_url",
+    "workspace",
+    "team",
+    "resource_type",
+    "current_status_update",
+    "custom_field_settings",
+    "completed",
+    "completed_at",
+    "completed_by",
+    "created_from_template",
+    "project_brief",
+)
+
+# List of section fields to be retrieved from Asana API
+SECTION_FIELDS = (
+    "gid",
+    "resource_type",
+    "name",
+    "created_at",
+    "project",
+    "projects",
+)
+
+# List of tag fields to be retrieved from Asana API
+TAG_FIELDS = (
+    "gid",
+    "resource_type",
+    "created_at",
+    "followers",
+    "name",
+    "color",
+    "notes",
+    "permalink_url",
+    "workspace",
+)
+
+# List of task fields to be retrieved from Asana API
+TASK_FIELDS = (
+    "gid",
+    "resource_type",
+    "name",
+    "approval_status",
+    "assignee_status",
+    "created_at",
+    "assignee",
+    "start_on",
+    "start_at",
+    "due_on",
+    "due_at",
+    "completed",
+    "completed_at",
+    "completed_by",
+    "modified_at",
+    "dependencies",
+    "dependents",
+    "external",
+    "notes",
+    "num_subtasks",
+    "resource_subtype",
+    "followers",
+    "parent",
+    "permalink_url",
+    "tags",
+    "workspace",
+    "custom_fields",
+    "project",
+    "memberships",
+    "memberships.project.name",
+    "memberships.section.name",
+)
+
+# List of story fields to be retrieved from Asana API
+STORY_FIELDS = (
+    "gid",
+    "resource_type",
+    "created_at",
+    "created_by",
+    "resource_subtype",
+    "text",
+    "is_pinned",
+    "assignee",
+    "dependency",
+    "follower",
+    "new_section",
+    "old_section",
+    "new_text_value",
+    "old_text_value",
+    "preview",
+    "project",
+    "source",
+    "story",
+    "tag",
+    "target",
+    "task",
+    "sticker_name",
+    "custom_field",
+    "type",
+)
+
+# List of team fields to be retrieved from Asana API
+TEAMS_FIELD = (
+    "gid",
+    "resource_type",
+    "name",
+    "description",
+    "organization",
+    "permalink_url",
+    "visibility",
+)
+
+# List of user fields to be retrieved from Asana API
+USER_FIELDS = ("gid", "resource_type", "name", "email", "photo", "workspaces")

omniload/src/attio/__init__.py

@@ -0,0 +1,102 @@
+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},
+        },
+    )
+    # https://docs.attio.com/rest-api/endpoint-reference/objects/list-objects - does not support pagination
+    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_all(path, "get")
+
+    # https://docs.attio.com/rest-api/endpoint-reference/records/list-records
+    @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_paginated(path, "post")
+
+    # https://docs.attio.com/rest-api/endpoint-reference/lists/list-all-lists -- does not support pagination
+    @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_all(path, "get")
+
+    # https://docs.attio.com/rest-api/endpoint-reference/entries/list-entries
+    @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_paginated(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_all(path, "get"):
+            if params[0] in lst["parent_object"]:
+                path = f"lists/{lst['id']['list_id']}/entries/query"
+                yield from attio_client.fetch_paginated(path, "post")
+
+    return (
+        fetch_objects,
+        fetch_records,
+        fetch_lists,
+        fetch_list_entries,
+        fetch_all_list_entries,
+    )

omniload/src/attio/helpers.py

@@ -0,0 +1,65 @@
+from omniload.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_paginated(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 = {"limit": limit, "offset": offset, **params}
+            if method == "get":
+                response = self.client.get(
+                    url, headers=self.headers, params=query_params
+                )
+            else:
+                json_body = {**params, "limit": limit, "offset": offset}
+                response = self.client.post(url, headers=self.headers, json=json_body)
+
+            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:
+                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 fetch_all(self, path: str, method: str = "get", params=None):
+        url = f"{self.base_url}/{path}"
+        params = params or {}
+
+        if method == "get":
+            response = self.client.get(url, headers=self.headers, params=params)
+        else:
+            response = self.client.post(url, headers=self.headers, json=params)
+
+        response.raise_for_status()
+        data = response.json().get("data", [])
+        for item in data:
+            yield flatten_item(item)
+
+
+def flatten_item(item: dict) -> dict:
+    if "id" in item:
+        for key, value in item["id"].items():
+            item[key] = value
+    return item

omniload/src/blob.py

@@ -0,0 +1,95 @@
+import warnings
+from typing import Tuple, TypeAlias
+from urllib.parse import ParseResult, urlparse
+
+BucketName: TypeAlias = str
+FileGlob: TypeAlias = str
+
+
+class UnsupportedEndpointError(Exception):
+    pass
+
+
+def parse_uri(uri: ParseResult, table: str) -> Tuple[BucketName, FileGlob]:
+    """
+    parse the URI of a blob storage and
+    return the bucket name and the file glob.
+
+    Supports the following Forms:
+    - uri: "gs://"
+      table: "bucket-name/file-glob"
+    - uri: "gs://uri-bucket-name" (uri-bucket-name is preferred)
+      table: "gs://table-bucket-name/file-glob"
+    - uri: "gs://"
+      table: "gs://bucket-name/file-glob"
+    - uri: gs://bucket-name/file-glob
+      table: None
+    - uri: "gs://bucket-name"
+      table: "file-glob"
+
+    The first form is the prefered method. Other forms are supported but discouraged.
+    """
+
+    table = table.strip()
+    host = uri.netloc.strip()
+
+    if table == "" or uri.path.strip() != "":
+        warnings.warn(
+            f"Using the form '{uri.scheme}://bucket-name/file-glob' is deprecated and will be removed in future versions.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return host, uri.path.lstrip("/")
+
+    table_uri = urlparse(table)
+
+    if host != "":
+        return host, table_uri.path.lstrip("/")
+
+    if table_uri.hostname:
+        return table_uri.hostname, table_uri.path.lstrip("/")
+
+    parts = table_uri.path.lstrip("/").split("/", maxsplit=1)
+    if len(parts) != 2:
+        return "", parts[0]
+
+    return parts[0], parts[1]
+
+
+def parse_endpoint(path: str) -> str:
+    """
+    Parse the endpoint kind from the URI.
+
+    kind is a file format. one of [csv, jsonl, parquet]
+    """
+    file_extension = path.split(".")[-1]
+    if file_extension == "gz":
+        file_extension = path.split(".")[-2]
+    if file_extension == "csv":
+        endpoint = "read_csv"
+    elif file_extension == "jsonl":
+        endpoint = "read_jsonl"
+    elif file_extension == "parquet":
+        endpoint = "read_parquet"
+    else:
+        raise UnsupportedEndpointError(f"Unsupported file format: {file_extension}")
+    return endpoint
+
+
+def determine_endpoint(table: str, path: str) -> str:
+    """
+    determines the endpoint/method to use for reading data from a blob source
+    """
+
+    if "#" in table:
+        _, endpoint = table.split("#")
+        if endpoint not in ["csv", "csv_headless", "jsonl", "parquet"]:
+            raise UnsupportedEndpointError(f"Unsupported file format: {endpoint}")
+        endpoint = f"read_{endpoint}"
+    else:
+        try:
+            endpoint = parse_endpoint(path)
+        except Exception as e:
+            raise ValueError(f"Failed to parse endpoint from path: {path}") from e
+
+    return endpoint

omniload/src/bruin/__init__.py

@@ -0,0 +1,76 @@
+"""Bruin source for fetching pipeline and asset data from Bruin Cloud API"""
+
+from typing import Iterator
+
+import dlt
+from dlt.sources.helpers import requests
+
+BASE_URL = "https://cloud.getbruin.com/api/v1"
+
+
+def _fetch_pipelines(headers: dict) -> list:
+    """Fetch pipelines data from API."""
+    response = requests.get(f"{BASE_URL}/pipelines", headers=headers)
+    response.raise_for_status()
+    return response.json()
+
+
+@dlt.source(name="bruin", max_table_nesting=0)
+def bruin_source(api_token: str):
+    """
+    A dlt source for the Bruin Cloud API.
+
+    Args:
+        api_token (str): The API token for authentication.
+
+    Returns:
+        DltResource: Resources for pipelines and assets data.
+    """
+    headers = {"Authorization": f"Bearer {api_token}"}
+
+    @dlt.resource(write_disposition="replace")
+    def pipelines() -> Iterator[dict]:
+        """
+        Fetches all pipelines and yields pipeline_id and pipeline_name for each.
+        """
+        data = _fetch_pipelines(headers)
+
+        for pipeline in data:
+            yield {
+                "name": pipeline.get("name"),
+                "description": pipeline.get("description"),
+                "project": pipeline.get("project"),
+                "owner": pipeline.get("owner"),
+                "default_connections": pipeline.get("default_connections"),
+                "schedule": pipeline.get("schedule"),
+                "commit": pipeline.get("commit"),
+                "start_date": pipeline.get("start_date"),
+            }
+
+    @dlt.resource(write_disposition="replace")
+    def assets() -> Iterator[dict]:
+        """
+        Fetches all assets from all pipelines (same endpoint as pipelines).
+        """
+        data = _fetch_pipelines(headers)
+
+        for pipeline in data:
+            pipeline_assets = pipeline.get("assets", [])
+            for asset in pipeline_assets:
+                yield {
+                    "name": asset.get("name"),
+                    "type": asset.get("type"),
+                    "pipeline": asset.get("pipeline"),
+                    "project": asset.get("project"),
+                    "uri": asset.get("uri"),
+                    "description": asset.get("description"),
+                    "upstreams": asset.get("upstreams"),
+                    "downstream": asset.get("downstream"),
+                    "owner": asset.get("owner"),
+                    "content": asset.get("content"),
+                    "columns": asset.get("columns"),
+                    "materialization": asset.get("materialization"),
+                    "parameters": asset.get("parameters"),
+                }
+
+    return pipelines, assets

omniload/src/chess/__init__.py

@@ -0,0 +1,180 @@
+# 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.
+
+"""A source loading player profiles and games from chess.com api"""
+
+from typing import Any, Callable, Dict, Iterator, List, Sequence
+
+import dlt
+from dlt.common import pendulum
+from dlt.common.typing import TDataItem
+from dlt.sources import DltResource
+from dlt.sources.helpers import requests
+
+from .helpers import get_path_with_retry, get_url_with_retry, validate_month_string
+from .settings import UNOFFICIAL_CHESS_API_URL
+
+
+@dlt.source(name="chess", max_table_nesting=0)
+def source(
+    players: List[str], start_month: str = None, end_month: str = None
+) -> Sequence[DltResource]:
+    """
+    A dlt source for the chess.com api. It groups several resources (in this case chess.com API endpoints) containing
+    various types of data: user profiles or chess match results
+    Args:
+        players (List[str]): A list of the player usernames for which to get the data.
+        start_month (str, optional): Filters out all the matches happening before `start_month`. Defaults to None.
+        end_month (str, optional): Filters out all the matches happening after `end_month`. Defaults to None.
+    Returns:
+        Sequence[DltResource]: A sequence of resources that can be selected from including players_profiles,
+        players_archives, players_games, players_online_status
+    """
+    return (
+        players_profiles(players),
+        players_archives(players),
+        players_games(players, start_month=start_month, end_month=end_month),
+        players_online_status(players),
+    )
+
+
+@dlt.resource(
+    write_disposition="replace",
+    columns={
+        "last_online": {"data_type": "timestamp"},
+        "joined": {"data_type": "timestamp"},
+    },
+)
+def players_profiles(players: List[str]) -> Iterator[TDataItem]:
+    """
+    Yields player profiles for a list of player usernames.
+    Args:
+        players (List[str]): List of player usernames to retrieve profiles for.
+    Yields:
+        Iterator[TDataItem]: An iterator over player profiles data.
+    """
+
+    # get archives in parallel by decorating the http request with defer
+    @dlt.defer
+    def _get_profile(username: str) -> TDataItem:
+        return get_path_with_retry(f"player/{username}")
+
+    for username in players:
+        yield _get_profile(username)
+
+
+@dlt.resource(write_disposition="replace", selected=False)
+def players_archives(players: List[str]) -> Iterator[List[TDataItem]]:
+    """
+    Yields url to game archives for specified players.
+    Args:
+        players (List[str]): List of player usernames to retrieve archives for.
+    Yields:
+        Iterator[List[TDataItem]]: An iterator over list of player archive data.
+    """
+    for username in players:
+        data = get_path_with_retry(f"player/{username}/games/archives")
+        yield data.get("archives", [])
+
+
+@dlt.resource(
+    write_disposition="replace", columns={"end_time": {"data_type": "timestamp"}}
+)
+def players_games(
+    players: List[str], start_month: str = None, end_month: str = None
+) -> Iterator[Callable[[], List[TDataItem]]]:
+    """
+    Yields `players` games that happened between `start_month` and `end_month`.
+    Args:
+        players (List[str]): List of player usernames to retrieve games for.
+        start_month (str, optional): The starting month in the format "YYYY/MM". Defaults to None.
+        end_month (str, optional): The ending month in the format "YYYY/MM". Defaults to None.
+    Yields:
+        Iterator[Callable[[], List[TDataItem]]]: An iterator over callables that return a list of games for each player.
+    """  # do a simple validation to prevent common mistakes in month format
+    validate_month_string(start_month)
+    validate_month_string(end_month)
+
+    # get a list of already checked archives
+    # from your point of view, the state is python dictionary that will have the same content the next time this function is called
+    checked_archives = dlt.current.resource_state().setdefault("archives", [])
+    # get player archives, note that you can call the resource like any other function and just iterate it like a list
+    archives = players_archives(players)
+
+    # get archives in parallel by decorating the http request with defer
+    @dlt.defer
+    def _get_archive(url: str) -> List[TDataItem]:
+        try:
+            games = get_url_with_retry(url).get("games", [])
+            return games  # type: ignore
+        except requests.HTTPError as http_err:
+            # sometimes archives are not available and the error seems to be permanent
+            if http_err.response.status_code == 404:
+                return []
+            raise
+
+    # enumerate the archives
+    for url in archives:
+        # the `url` format is https://api.chess.com/pub/player/{username}/games/{YYYY}/{MM}
+        if start_month and url[-7:] < start_month:
+            continue
+        if end_month and url[-7:] > end_month:
+            continue
+        # do not download archive again
+        if url in checked_archives:
+            continue
+        checked_archives.append(url)
+        # get the filtered archive
+        yield _get_archive(url)
+
+
+@dlt.resource(write_disposition="append")
+def players_online_status(players: List[str]) -> Iterator[TDataItem]:
+    """
+    Returns current online status for a list of players.
+    Args:
+        players (List[str]): List of player usernames to check online status for.
+    Yields:
+        Iterator[TDataItem]: An iterator over the online status of each player.
+    """
+    # we'll use unofficial endpoint to get online status, the official seems to be removed
+    for player in players:
+        status = get_url_with_retry(f"{UNOFFICIAL_CHESS_API_URL}user/popup/{player}")
+        # return just relevant selection
+        yield {
+            "username": player,
+            "onlineStatus": status["onlineStatus"],
+            "lastLoginDate": status["lastLoginDate"],
+            "check_time": pendulum.now(),  # dlt can deal with native python dates
+        }
+
+
+@dlt.source
+def chess_dlt_config_example(
+    secret_str: str = dlt.secrets.value,
+    secret_dict: Dict[str, Any] = dlt.secrets.value,
+    config_int: int = dlt.config.value,
+) -> DltResource:
+    """
+    An example of a source that uses dlt to provide secrets and config values.
+    Args:
+        secret_str (str, optional): Secret string provided by dlt.secrets.value. Defaults to dlt.secrets.value.
+        secret_dict (Dict[str, Any], optional): Secret dictionary provided by dlt.secrets.value. Defaults to dlt.secrets.value.
+        config_int (int, optional): Config integer provided by dlt.config.value. Defaults to dlt.config.value.
+    Returns:
+        DltResource: Returns a resource yielding the configured values.
+    """
+
+    # returns a resource yielding the configured values - it is just a test
+    return dlt.resource([secret_str, secret_dict, config_int], name="config_values")