omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/factory.py

@@ -0,0 +1,305 @@
+from typing import Dict, Protocol, Type
+from urllib.parse import urlparse
+
+from dlt.common.destination import Destination
+
+from omniload.src.destinations import (
+    AthenaDestination,
+    BigQueryDestination,
+    ClickhouseDestination,
+    CrateDBDestination,
+    CsvDestination,
+    DatabricksDestination,
+    DuckDBDestination,
+    ElasticsearchDestination,
+    GCSDestination,
+    MongoDBDestination,
+    MotherduckDestination,
+    MsSQLDestination,
+    MySqlDestination,
+    PostgresDestination,
+    RedshiftDestination,
+    S3Destination,
+    SnowflakeDestination,
+    SqliteDestination,
+    SynapseDestination,
+    TrinoDestination,
+)
+from omniload.src.sources import (
+    AdjustSource,
+    AirtableSource,
+    AlliumSource,
+    AnthropicSource,
+    AppleAppStoreSource,
+    ApplovinMaxSource,
+    AppLovinSource,
+    AppsflyerSource,
+    ArrowMemoryMappedSource,
+    AsanaSource,
+    AttioSource,
+    BruinSource,
+    ChessSource,
+    ClickupSource,
+    CouchbaseSource,
+    CursorSource,
+    CustomerIoSource,
+    DoceboSource,
+    DuneSource,
+    DynamoDBSource,
+    ElasticsearchSource,
+    FacebookAdsSource,
+    FirefliesSource,
+    FluxxSource,
+    FrankfurterSource,
+    FreshdeskSource,
+    FundraiseupSource,
+    GCSSource,
+    GitHubSource,
+    GoogleAdsSource,
+    GoogleAnalyticsSource,
+    GoogleSheetsSource,
+    GorgiasSource,
+    HostawaySource,
+    HttpSource,
+    HubspotSource,
+    IndeedSource,
+    InfluxDBSource,
+    IntercomSource,
+    IsocPulseSource,
+    JiraSource,
+    KafkaSource,
+    KinesisSource,
+    KlaviyoSource,
+    LinearSource,
+    LinkedInAdsSource,
+    LocalCsvSource,
+    MailchimpSource,
+    MixpanelSource,
+    MondaySource,
+    MongoDbSource,
+    NotionSource,
+    PersonioSource,
+    PhantombusterSource,
+    PinterestSource,
+    PipedriveSource,
+    PlusVibeAISource,
+    PrimerSource,
+    QuickBooksSource,
+    RedditAdsSource,
+    RevenueCatSource,
+    S3Source,
+    SalesforceSource,
+    SFTPSource,
+    ShopifySource,
+    SlackSource,
+    SmartsheetSource,
+    SnapchatAdsSource,
+    SocrataSource,
+    SolidgateSource,
+    SqlSource,
+    StripeAnalyticsSource,
+    TikTokSource,
+    TrustpilotSource,
+    WiseSource,
+    ZendeskSource,
+    ZoomSource,
+)
+
+SQL_SOURCE_SCHEMES = [
+    "bigquery",
+    "crate",
+    "cratedb",
+    "duckdb",
+    "mssql",
+    "mssql+pyodbc",
+    "mysql",
+    "mysql+pymysql",
+    "mysql+mysqlconnector",
+    "md",
+    "motherduck",
+    "postgres",
+    "postgresql",
+    "postgresql+psycopg2",
+    "redshift",
+    "redshift+psycopg2",
+    "snowflake",
+    "sqlite",
+    "oracle",
+    "oracle+cx_oracle",
+    "oracle+oracledb",
+    "hana",
+    "clickhouse",
+    "databricks",
+    "db2",
+    "spanner",
+    "trino",
+]
+
+
+class SourceProtocol(Protocol):
+    def dlt_source(self, uri: str, table: str, **kwargs):
+        pass
+
+    def handles_incrementality(self) -> bool:
+        pass
+
+
+class DestinationProtocol(Protocol):
+    def dlt_dest(self, uri: str, **kwargs) -> Destination:
+        pass
+
+    def dlt_run_params(self, uri: str, table: str, **kwargs):
+        pass
+
+    def post_load(self) -> None:
+        pass
+
+
+def parse_scheme_from_uri(uri: str) -> str:
+    parsed = urlparse(uri)
+    if parsed.scheme != "":
+        return parsed.scheme
+
+    uri_parts = uri.split("://")
+    if len(uri_parts) > 1:
+        return uri_parts[0]
+
+    raise ValueError(f"Could not parse scheme from uri: {uri}")
+
+
+class SourceDestinationFactory:
+    source_scheme: str
+    destination_scheme: str
+    sources: Dict[str, Type[SourceProtocol]] = {
+        "allium": AlliumSource,
+        "anthropic": AnthropicSource,
+        "bruin": BruinSource,
+        "csv": LocalCsvSource,
+        "couchbase": CouchbaseSource,
+        "cursor": CursorSource,
+        "docebo": DoceboSource,
+        "dune": DuneSource,
+        "http": HttpSource,
+        "https": HttpSource,
+        "mongodb": MongoDbSource,
+        "mongodb+srv": MongoDbSource,
+        "notion": NotionSource,
+        "gsheets": GoogleSheetsSource,
+        "shopify": ShopifySource,
+        "gorgias": GorgiasSource,
+        "github": GitHubSource,
+        "chess": ChessSource,
+        "stripe": StripeAnalyticsSource,
+        "facebookads": FacebookAdsSource,
+        "fluxx": FluxxSource,
+        "slack": SlackSource,
+        "hostaway": HostawaySource,
+        "hubspot": HubspotSource,
+        "indeed": IndeedSource,
+        "intercom": IntercomSource,
+        "jira": JiraSource,
+        "airtable": AirtableSource,
+        "klaviyo": KlaviyoSource,
+        "mixpanel": MixpanelSource,
+        "appsflyer": AppsflyerSource,
+        "kafka": KafkaSource,
+        "adjust": AdjustSource,
+        "zendesk": ZendeskSource,
+        "mmap": ArrowMemoryMappedSource,
+        "s3": S3Source,
+        "dynamodb": DynamoDBSource,
+        "asana": AsanaSource,
+        "tiktok": TikTokSource,
+        "googleanalytics": GoogleAnalyticsSource,
+        "googleads": GoogleAdsSource,
+        "appstore": AppleAppStoreSource,
+        "gs": GCSSource,
+        "linkedinads": LinkedInAdsSource,
+        "linear": LinearSource,
+        "applovin": AppLovinSource,
+        "applovinmax": ApplovinMaxSource,
+        "salesforce": SalesforceSource,
+        "personio": PersonioSource,
+        "kinesis": KinesisSource,
+        "pipedrive": PipedriveSource,
+        "frankfurter": FrankfurterSource,
+        "freshdesk": FreshdeskSource,
+        "fundraiseup": FundraiseupSource,
+        "trustpilot": TrustpilotSource,
+        "phantombuster": PhantombusterSource,
+        "elasticsearch": ElasticsearchSource,
+        "attio": AttioSource,
+        "solidgate": SolidgateSource,
+        "quickbooks": QuickBooksSource,
+        "isoc-pulse": IsocPulseSource,
+        "smartsheet": SmartsheetSource,
+        "sftp": SFTPSource,
+        "pinterest": PinterestSource,
+        "redditads": RedditAdsSource,
+        "revenuecat": RevenueCatSource,
+        "socrata": SocrataSource,
+        "snapchatads": SnapchatAdsSource,
+        "zoom": ZoomSource,
+        "clickup": ClickupSource,
+        "influxdb": InfluxDBSource,
+        "wise": WiseSource,
+        "plusvibeai": PlusVibeAISource,
+        "monday": MondaySource,
+        "mailchimp": MailchimpSource,
+        "primer": PrimerSource,
+        "fireflies": FirefliesSource,
+        "customerio": CustomerIoSource,
+    }
+    destinations: Dict[str, Type[DestinationProtocol]] = {
+        "bigquery": BigQueryDestination,
+        "cratedb": CrateDBDestination,
+        "databricks": DatabricksDestination,
+        "duckdb": DuckDBDestination,
+        "motherduck": MotherduckDestination,
+        "md": MotherduckDestination,
+        "mssql": MsSQLDestination,
+        "postgres": PostgresDestination,
+        "postgresql": PostgresDestination,
+        "postgresql+psycopg2": PostgresDestination,
+        "redshift": RedshiftDestination,
+        "redshift+psycopg2": RedshiftDestination,
+        "redshift+redshift_connector": RedshiftDestination,
+        "snowflake": SnowflakeDestination,
+        "synapse": SynapseDestination,
+        "csv": CsvDestination,
+        "athena": AthenaDestination,
+        "clickhouse+native": ClickhouseDestination,
+        "clickhouse": ClickhouseDestination,
+        "elasticsearch": ElasticsearchDestination,
+        "mongodb": MongoDBDestination,
+        "mongodb+srv": MongoDBDestination,
+        "s3": S3Destination,
+        "gs": GCSDestination,
+        "sqlite": SqliteDestination,
+        "mysql": MySqlDestination,
+        "mysql+pymysql": MySqlDestination,
+        "trino": TrinoDestination,
+    }
+
+    def __init__(self, source_uri: str, destination_uri: str):
+        self.source_uri = source_uri
+        self.source_scheme = parse_scheme_from_uri(source_uri)
+
+        self.destination_uri = destination_uri
+        self.destination_scheme = parse_scheme_from_uri(destination_uri)
+
+    def get_source(self) -> SourceProtocol:
+        if self.source_scheme in SQL_SOURCE_SCHEMES:
+            return SqlSource()
+        elif self.source_scheme in self.sources:
+            return self.sources[self.source_scheme]()
+        else:
+            raise ValueError(f"Unsupported source scheme: {self.source_scheme}")
+
+    def get_destination(self) -> DestinationProtocol:
+        if self.destination_scheme in self.destinations:
+            return self.destinations[self.destination_scheme]()
+        else:
+            raise ValueError(
+                f"Unsupported destination scheme: {self.destination_scheme}"
+            )

omniload/src/filesystem/__init__.py

@@ -0,0 +1,133 @@
+# 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.
+
+"""Reads files in s3, gs or azure buckets using fsspec and provides convenience resources for chunked reading of various file formats"""
+
+from typing import Iterator, List, Optional, Tuple, Union
+
+import dlt
+from dlt.sources import DltResource
+from dlt.sources.credentials import FileSystemCredentials
+from dlt.sources.filesystem import FileItem, FileItemDict, fsspec_filesystem, glob_files
+
+from .helpers import (
+    AbstractFileSystem,
+    FilesystemConfigurationResource,
+)
+from .readers import (
+    ReadersSource,
+    _read_csv,
+    _read_csv_duckdb,
+    _read_csv_headless,
+    _read_jsonl,
+    _read_parquet,
+)
+
+
+@dlt.source(_impl_cls=ReadersSource, spec=FilesystemConfigurationResource)
+def readers(
+    bucket_url: str,
+    credentials: Union[FileSystemCredentials, AbstractFileSystem],
+    file_glob: Optional[str] = "*",
+) -> Tuple[DltResource, ...]:
+    """This source provides a few resources that are chunked file readers. Readers can be further parametrized before use
+       read_csv(chunksize, **pandas_kwargs)
+       read_jsonl(chunksize)
+       read_parquet(chunksize)
+
+    Args:
+        bucket_url (str): The url to the bucket.
+        credentials (FileSystemCredentials | AbstractFilesystem): The credentials to the filesystem of fsspec `AbstractFilesystem` instance.
+        file_glob (str, optional): The filter to apply to the files in glob format. by default lists all files in bucket_url non-recursively
+    """
+    filesystem_resource = filesystem(bucket_url, credentials, file_glob=file_glob)
+
+    # NOTE: incremental support is disabled until we can figure out
+    #       how to support incremental loads per matching file, rather
+    #       than a blanket threshold.
+    #
+    # filesystem_resource.apply_hints(
+    #     incremental=dlt.sources.incremental("modification_date"),
+    # )
+    return (
+        filesystem_resource
+        | dlt.transformer(name="read_csv", max_table_nesting=0)(_read_csv),
+        filesystem_resource
+        | dlt.transformer(name="read_csv_headless", max_table_nesting=0)(
+            _read_csv_headless
+        ),
+        filesystem_resource
+        | dlt.transformer(name="read_jsonl", max_table_nesting=0)(_read_jsonl),
+        filesystem_resource
+        | dlt.transformer(name="read_parquet", max_table_nesting=0)(_read_parquet),
+        filesystem_resource
+        | dlt.transformer(name="read_csv_duckdb", max_table_nesting=0)(
+            _read_csv_duckdb
+        ),
+    )
+
+
+@dlt.resource(
+    primary_key="file_url", spec=FilesystemConfigurationResource, standalone=True
+)
+def filesystem(
+    bucket_url: str = dlt.secrets.value,
+    credentials: Union[FileSystemCredentials, AbstractFileSystem] = dlt.secrets.value,
+    file_glob: Optional[str] = "*",
+    files_per_page: int = 100,
+    extract_content: bool = True,
+) -> Iterator[List[FileItem]]:
+    """This resource lists files in `bucket_url` using `file_glob` pattern. The files are yielded as FileItem which also
+    provide methods to open and read file data. It should be combined with transformers that further process (ie. load files)
+
+    Args:
+        bucket_url (str): The url to the bucket.
+        credentials (FileSystemCredentials | AbstractFilesystem): The credentials to the filesystem of fsspec `AbstractFilesystem` instance.
+        file_glob (str, optional): The filter to apply to the files in glob format. by default lists all files in bucket_url non-recursively
+        files_per_page (int, optional): The number of files to process at once, defaults to 100.
+        extract_content (bool, optional): If true, the content of the file will be extracted if
+            false it will return a fsspec file, defaults to False.
+
+    Returns:
+        Iterator[List[FileItem]]: The list of files.
+    """
+
+    if isinstance(credentials, AbstractFileSystem):
+        fs_client = credentials
+    else:
+        fs_client = fsspec_filesystem(bucket_url, credentials)[0]
+
+    files_chunk: List[FileItem] = []
+    for file_model in glob_files(fs_client, bucket_url, file_glob):
+        file_dict = FileItemDict(file_model, credentials)
+        if extract_content:
+            file_dict["file_content"] = file_dict.read_bytes()
+        files_chunk.append(file_dict)  # type: ignore
+        # wait for the chunk to be full
+        if len(files_chunk) >= files_per_page:
+            yield files_chunk
+            files_chunk = []
+    if files_chunk:
+        yield files_chunk
+
+
+read_csv = dlt.transformer(standalone=True, max_table_nesting=0)(_read_csv)
+read_csv_headless = dlt.transformer(standalone=True, max_table_nesting=0)(
+    _read_csv_headless
+)
+read_jsonl = dlt.transformer(standalone=True, max_table_nesting=0)(_read_jsonl)
+read_parquet = dlt.transformer(standalone=True, max_table_nesting=0)(_read_parquet)
+read_csv_duckdb = dlt.transformer(standalone=True, max_table_nesting=0)(
+    _read_csv_duckdb
+)

omniload/src/filesystem/helpers.py

@@ -0,0 +1,114 @@
+# 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.
+
+"""Helpers for the filesystem resource."""
+
+from typing import Any, Dict, Iterable, List, Optional, Type, Union
+
+import dlt
+from dlt.common.configuration import resolve_type
+from dlt.common.typing import TDataItem
+from dlt.sources import DltResource
+from dlt.sources.config import configspec, with_config
+from dlt.sources.credentials import (
+    CredentialsConfiguration,
+    FilesystemConfiguration,
+    FileSystemCredentials,
+)
+from dlt.sources.filesystem import fsspec_filesystem
+from fsspec import AbstractFileSystem  # type: ignore
+
+
+@configspec
+class FilesystemConfigurationResource(FilesystemConfiguration):
+    credentials: Union[FileSystemCredentials, AbstractFileSystem] = None
+    file_glob: Optional[str] = "*"
+    files_per_page: int = 100
+    extract_content: bool = False
+
+    @resolve_type("credentials")
+    def resolve_credentials_type(self) -> Type[CredentialsConfiguration]:
+        # use known credentials or empty credentials for unknown protocol
+        return Union[
+            self.PROTOCOL_CREDENTIALS.get(self.protocol)
+            or Optional[CredentialsConfiguration],
+            AbstractFileSystem,
+        ]  # type: ignore[return-value]
+
+
+def fsspec_from_resource(filesystem_instance: DltResource) -> AbstractFileSystem:
+    """Extract authorized fsspec client from a filesystem resource"""
+
+    @with_config(
+        spec=FilesystemConfiguration,
+        sections=("sources", filesystem_instance.section, filesystem_instance.name),
+    )
+    def _get_fsspec(
+        bucket_url: str, credentials: Optional[FileSystemCredentials]
+    ) -> AbstractFileSystem:
+        return fsspec_filesystem(bucket_url, credentials)[0]
+
+    return _get_fsspec(
+        filesystem_instance.explicit_args.get("bucket_url", dlt.config.value),
+        filesystem_instance.explicit_args.get("credentials", dlt.secrets.value),
+    )
+
+
+def add_columns(columns: List[str], rows: List[List[Any]]) -> List[Dict[str, Any]]:
+    """Adds column names to the given rows.
+
+    Args:
+        columns (List[str]): The column names.
+        rows (List[List[Any]]): The rows.
+
+    Returns:
+        List[Dict[str, Any]]: The rows with column names.
+    """
+    result = []
+    for row in rows:
+        result.append(dict(zip(columns, row)))
+
+    return result
+
+
+def fetch_arrow(file_data, chunk_size: int) -> Iterable[TDataItem]:  # type: ignore
+    """Fetches data from the given CSV file.
+
+    Args:
+        file_data (DuckDBPyRelation): The CSV file data.
+        chunk_size (int): The number of rows to read at once.
+
+    Yields:
+        Iterable[TDataItem]: Data items, read from the given CSV file.
+    """
+    batcher = file_data.fetch_arrow_reader(batch_size=chunk_size)
+    yield from batcher
+
+
+def fetch_json(file_data, chunk_size: int) -> List[Dict[str, Any]]:  # type: ignore
+    """Fetches data from the given CSV file.
+
+    Args:
+        file_data (DuckDBPyRelation): The CSV file data.
+        chunk_size (int): The number of rows to read at once.
+
+    Yields:
+        Iterable[TDataItem]: Data items, read from the given CSV file.
+    """
+    while True:
+        batch = file_data.fetchmany(chunk_size)
+        if not batch:
+            break
+
+        yield add_columns(file_data.columns, batch)