From 19d4c4526713597b162f2bcbd958984fff6f659d Mon Sep 17 00:00:00 2001 From: TiM Date: Tue, 4 Aug 2026 07:46:17 +1200 Subject: [PATCH] feat(sql): load remote SQLite and DuckDB databases Load a SQLite or DuckDB database that lives on a remote filesystem from the URI its storage service already names it with, `s3://analytics/snapshots/ events.duckdb`. A filesystem-scheme source URI whose object carries a database extension routes to the SQL source instead of a format reader, so `--source-table` keeps selecting a table inside the database and SQL reflection, chunking, and type mapping stay on the path. Stage each object for the complete ingestion lifetime with size checks, byte verification, cleanup on success and failure, and no remote writeback. Reuse the filesystem credential parsing so the storage options are the same query parameters the matching file source takes, and keep credentials out of logs, errors, and object representations. Resolve the engine from the staged file's header (SQLite's marker at offset 0, DuckDB's `DUCK` behind its checksum), because `.db` names both engines and an object's name is not evidence of its contents. An empty object, a database created but not yet written to, falls back to an unambiguous extension. Cover the schemes whose credential parsing is already shared: s3, r2, gs, az, adls, and abfss. Reject the split form (`--source-uri s3://` plus the object path on `--source-table`) by naming the carrier to use, because the table selects a table inside the database. Validate a dry run against the SQL source the real run uses, without downloading the object. Document the remote URI contract and cover parsing, lifecycle, compatibility, and emulator-backed ingestion. --- docs/changelog.md | 3 + .../ADR-001-stage-remote-file-databases.md | 85 +++++ docs/index.md | 11 +- docs/supported-sources/duckdb.md | 29 +- docs/supported-sources/filesystem.md | 46 +++ docs/supported-sources/sqlite.md | 32 +- pyproject.toml | 2 +- src/dlt_filesystem/source/impl/remote.py | 78 +--- src/dlt_filesystem/staging.py | 193 ++++++++++ src/dlt_filesystem/util/auth.py | 99 ++++- src/omniload/api.py | 46 ++- src/omniload/model.py | 2 + src/omniload/source/sql_database/remote.py | 151 ++++++++ tests/dlt_filesystem/test_staging.py | 161 ++++++++ .../filesystem/test_remote_integration.py | 91 ++++- tests/main/filesystem/test_remote_read.py | 12 + tests/main/test_remote_database.py | 358 ++++++++++++++++++ tests/main/test_sources.py | 31 ++ 18 files changed, 1338 insertions(+), 92 deletions(-) create mode 100644 docs/decisions/ADR-001-stage-remote-file-databases.md create mode 100644 src/dlt_filesystem/staging.py create mode 100644 src/omniload/source/sql_database/remote.py create mode 100644 tests/dlt_filesystem/test_staging.py create mode 100644 tests/main/test_remote_database.py diff --git a/docs/changelog.md b/docs/changelog.md index bb99051d0..f027debc2 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,9 @@ ## in progress +- Database: Load SQLite and DuckDB databases that live on a remote filesystem, + addressed by their plain object URI, for example + `s3://analytics/snapshots/events.duckdb`. Thanks, @hampsterx. - Filesystem: Treat XLSX and ODS workbooks as multi-table sources by default, while preserving one-table loads through explicit worksheet selectors. Thanks, @hampsterx. diff --git a/docs/decisions/ADR-001-stage-remote-file-databases.md b/docs/decisions/ADR-001-stage-remote-file-databases.md new file mode 100644 index 000000000..b9cbdee82 --- /dev/null +++ b/docs/decisions/ADR-001-stage-remote-file-databases.md @@ -0,0 +1,85 @@ +# ADR-001: Stage remote file databases before SQL ingestion + +**Status**: Proposed +**Date**: 2026-08-03 + +## Context + +SQLite and DuckDB require random access to a database file through their +existing SQLAlchemy integrations. Object-storage streams do not provide a +shared database interface that both engines can open directly. + +Remote database access needs the same URI and lifecycle semantics across +supported storage transports. Storage credentials must remain separate from +logged object locations, and the database file must remain available while dlt +constructs and consumes its lazy SQL resource. + +## Decision + +We materialize remote SQLite and DuckDB source files into a run-scoped local +temporary directory before invoking the existing SQL source path. + +A remote database is addressed by its plain object URI, the same carrier the +filesystem sources read files from, with storage options as its query +parameters. A filesystem source URI whose object name carries a database +extension routes to the SQL source rather than to a format reader, so +`--source-table` keeps selecting a table inside the database. Transport support +covers the storage schemes that share the repository's existing filesystem +credential parsing and emulator coverage: `s3://`, `r2://`, `gs://`, `az://`, +`adls://`, and `abfss://`. + +The engine is resolved from the staged file's header, with an unambiguous +extension as the fallback, because `.db` names both engines and an object's name +is not evidence of its contents. + +The staging operation checks the remote object's reported size and available +disk space, streams the object while counting bytes, and keeps the local copy +alive through the complete ingestion run. It removes the copy on normal +completion and exception exits. The staged database is source-only, and local +changes are never written back to object storage. + +## Alternatives considered + +- **Carry the object URI in a `location` query parameter of the SQL URI** + (`duckdb:///?location=&`): rejected + because it makes the caller encode a URI inside a URI for a location the + storage scheme already names unambiguously. +- **Stack storage schemes inside the SQL URI**: rejected because nested URI + authorities, query strings, and credential ownership are ambiguous to + standard URL parsers. +- **Decode database files through a filesystem reader**: rejected because the + SQL source already owns reflection, chunking, type mapping, and table + selection. A reader would reimplement them one file format at a time. +- **Use DuckDB `httpfs` or `ATTACH` as the primary path**: rejected because it + does not provide the same mechanism for SQLite. Engine-specific fast paths + can be added later without changing the public URI grammar. +- **Expose object storage through an engine-specific VFS**: rejected because + there is no shared, maintained random-access interface for both engines and + all three launch transports. +- **Keep a persistent local cache**: rejected because cache invalidation, + credential boundaries, and stale-object behavior require a separate policy. + Run-scoped staging has explicit ownership and cleanup. +- **Write staged database changes back remotely**: rejected because concurrent + writers and atomic replacement require consistency guarantees that a source + ingestion does not need. + +## Consequences + +- Remote file databases reuse the established SQL reflection, extraction, and + loading path. +- Each run downloads the whole object and requires enough local disk for the + database plus any engine sidecar files. +- Normal success and failure paths remove staged data. A process terminated + without cleanup may leave a temporary directory when a persistent staging + parent is configured. +- Credentials remain storage options and are excluded from safe object + locations and object representations. +- A database extension on a storage URI is reserved: an object named `.db`, + `.ddb`, `.duckdb`, `.sqlite`, or `.sqlite3` is never offered to a format + reader, and one that holds neither engine reports that instead of a format + error. +- Additional transports need a credential adapter and integration coverage. A + transport that builds its filesystem client inside its own source class needs + that construction lifted out before it can stage objects. +- A future engine-specific fast path must preserve the same source-only behavior + and public URI contract. diff --git a/docs/index.md b/docs/index.md index f383870e7..6adc79ff5 100644 --- a/docs/index.md +++ b/docs/index.md @@ -16,13 +16,14 @@ inherited by [dlt]: `append`, `merge`, and `delete+insert`. We recommend using [uv](https://github.com/astral-sh/uv) to run `omniload`. -``` +```bash pip install uv uvx omniload ``` Alternatively, if you'd like to install it globally: -``` + +```bash uv pip install --system omniload ``` @@ -35,7 +36,8 @@ Check out the {ref}`Quickstart` guide to get started with omniload. ### License The project is licensed under the MIT License, see the [LICENSE] file for details. -Some components are licensed under the Apache 2.0 license, see the [NOTICE] file for details. +Some components are licensed under the Apache 2.0 license, see the [NOTICE] +file for details. ### Acknowledgements @@ -43,7 +45,6 @@ This project would not have been possible without the amazing work by the authors and contributors to [SQLAlchemy], [dlt], and [ingestr], turtles all the way down. Kudos. - ```{toctree} :caption: Commands and adapters :maxdepth: 1 @@ -79,13 +80,13 @@ tutorials/* :maxdepth: 1 :hidden: :glob: +decisions/* sandbox changelog contributors backlog ``` - [dlt]: https://github.com/dlt-hub/dlt [ingestr]: https://bruin-data.github.io/ingestr/ [LICENSE]: https://github.com/panodata/omniload/blob/main/LICENSE diff --git a/docs/supported-sources/duckdb.md b/docs/supported-sources/duckdb.md index ba8ecbf8e..4402d7631 100644 --- a/docs/supported-sources/duckdb.md +++ b/docs/supported-sources/duckdb.md @@ -1,9 +1,11 @@ # DuckDB + DuckDB is an in-memory database designed to be fast and reliable. omniload supports DuckDB as both a source and destination. ## URI format + The URI format for DuckDB is as follows: ```text @@ -11,6 +13,31 @@ duckdb:/// ``` URI parameters: + - `database-file`: the path to the DuckDB database file -The same URI structure can be used both for sources and destinations. You can read more about SQLAlchemy's DuckDB dialect [here](https://github.com/Mause/duckdb_engine). +The same URI structure can be used both for sources and destinations. See the +[DuckDB SQLAlchemy dialect][duckdb-sqlalchemy] for details. + +## Remote source files + +A DuckDB source file that lives in [Amazon S3](s3.md), [Cloudflare R2](r2.md), +[Azure Blob Storage](azure-storage.md), or +[Google Cloud Storage](google-cloud-storage.md) is addressed by its object URI, +with the storage credentials as query parameters: + +```bash +omniload ingest \ + --source-uri 's3://analytics/snapshots/events.duckdb?access_key_id=ACCESS&secret_access_key=SECRET' \ + --source-table 'main.events' \ + --dest-uri 'duckdb:///local.duckdb' \ + --dest-table 'raw.events' +``` + +Percent-encode every query value that contains reserved characters such as `+`, +`/`, `=`, `&`, or `?`. + +The object is staged locally for the duration of the run, so remote databases +are sources only. See {ref}`database-files` for the whole contract. + +[duckdb-sqlalchemy]: https://github.com/Mause/duckdb_engine diff --git a/docs/supported-sources/filesystem.md b/docs/supported-sources/filesystem.md index 2d8dc11f6..97fb68834 100644 --- a/docs/supported-sources/filesystem.md +++ b/docs/supported-sources/filesystem.md @@ -215,6 +215,52 @@ cloud blob destinations currently address one table path. They reject a plural workbook load. Select one worksheet or use a dataset-capable destination. ::: +(database-files)= + +## Database files + +An object whose name carries a database extension is a database, not a file to +decode. omniload hands it to the SQL source instead of a reader, so +`--source-table` selects a table inside the database and the usual SQL +reflection, chunking, and type mapping apply: + +```sh +omniload ingest \ + --source-uri 's3://analytics/snapshots/events.duckdb?access_key_id=ACCESS&secret_access_key=SECRET' \ + --source-table 'main.events' \ + --dest-uri 'duckdb:///local.duckdb' \ + --dest-table 'raw.events' +``` + +| Extension | Engine | +|:-----------------------|:-----------------------------------| +| .db | SQLite or DuckDB, read from header | +| .ddb, .duckdb | DuckDB | +| .sqlite, .sqlite3 | SQLite | + +The engine is read from the file header, so a mislabeled object still loads. An +empty file falls back to its extension, and a `.db` file that is neither engine +reports both candidates. + +Databases are read from `s3://`, `r2://`, `gs://`, `az://`, `adls://`, and +`abfss://`. Credentials and storage options are the same query parameters the +matching {ref}`filesystem source ` takes. The URI names one +object: globs, `#` fragments, and credentials in the authority are rejected +rather than silently reinterpreted. The object path also has to ride +`--source-uri`, because `--source-table` names a table inside the database +rather than the object, unlike the split form the file sources accept. + +Local databases need no filesystem source. Address them with +[DuckDB](duckdb.md) or [SQLite](sqlite.md) directly, as +`duckdb:///path/to/events.duckdb`. + +:::{note} +The object is downloaded whole into a run-scoped temporary directory before it +is opened, so the database must fit on local disk. The copy is removed after +success and after failure, and changes to it are never written back. Remote +databases are sources only. +::: + (file-format-routing)= ## File format routing diff --git a/docs/supported-sources/sqlite.md b/docs/supported-sources/sqlite.md index b949165be..1b7909437 100644 --- a/docs/supported-sources/sqlite.md +++ b/docs/supported-sources/sqlite.md @@ -1,9 +1,12 @@ # SQLite -SQLite is a C-language library that implements a small, fast, self-contained, high-reliability, full-featured, SQL database engine. + +SQLite is a C-language library that implements a small, fast, self-contained, +high-reliability, full-featured SQL database engine. omniload supports SQLite as a source and a destination. ## URI format + The URI format for SQLite is as follows: ```text @@ -11,6 +14,31 @@ sqlite:/// ``` URI parameters: + - `database-file`: the path to the SQLite database file. -The same URI structure can be used both for sources and destinations. You can read more about SQLAlchemy's SQLite dialect [here](https://docs.sqlalchemy.org/en/20/core/engines.html#sqlite). +The same URI structure can be used both for sources and destinations. See the +[SQLite SQLAlchemy dialect][sqlite-sqlalchemy] for details. + +## Remote source files + +An SQLite source file that lives in [Amazon S3](s3.md), [Cloudflare R2](r2.md), +[Azure Blob Storage](azure-storage.md), or +[Google Cloud Storage](google-cloud-storage.md) is addressed by its object URI, +with the storage credentials as query parameters: + +```bash +omniload ingest \ + --source-uri 's3://analytics/snapshots/events.sqlite?access_key_id=ACCESS&secret_access_key=SECRET' \ + --source-table 'main.events' \ + --dest-uri 'duckdb:///local.duckdb' \ + --dest-table 'raw.events' +``` + +Percent-encode every query value that contains reserved characters such as `+`, +`/`, `=`, `&`, or `?`. + +The object is staged locally for the duration of the run, so remote databases +are sources only. See {ref}`database-files` for the whole contract. + +[sqlite-sqlalchemy]: https://docs.sqlalchemy.org/en/20/dialects/sqlite.html diff --git a/pyproject.toml b/pyproject.toml index 1dc48df8a..bb616fb4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,7 @@ dependencies = [ "google-cloud-bigquery-storage<3", "google-cloud-spanner<4", "influxdb-client<2", - "mq-bridge-py>=0.3.2,<0.4", + "mq-bridge-py>=0.3.2,<0.4,!=0.3.9", "mysql-connector-python<27", "omniload[filesystem]", "oracledb<5", diff --git a/src/dlt_filesystem/source/impl/remote.py b/src/dlt_filesystem/source/impl/remote.py index 99750e87d..a80b8ae9b 100644 --- a/src/dlt_filesystem/source/impl/remote.py +++ b/src/dlt_filesystem/source/impl/remote.py @@ -1,5 +1,3 @@ -import base64 -import json from abc import abstractmethod from typing import TYPE_CHECKING, Any, Dict, Type from urllib.parse import parse_qs, urlparse @@ -16,7 +14,12 @@ parse_uri, source_selects_single_file, ) -from dlt_filesystem.util.auth import AzureBlobAuth, parse_azure_blob_auth +from dlt_filesystem.util.auth import ( + azure_blob_filesystem_kwargs, + gcs_filesystem_kwargs, + parse_azure_blob_auth, + s3_filesystem_kwargs, +) if TYPE_CHECKING: from fsspec import AbstractFileSystem @@ -52,25 +55,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): bucket_url = f"gs://{bucket_name}" - credentials_path = params.pop("credentials_path", [None])[0] - credentials_base64 = params.pop("credentials_base64", [None])[0] - - # Merge params into fs kwargs, without overriding kwargs already - # supplied by the caller (e.g. filesystem_incremental, column_types). - for key, value in params.items(): - kwargs.setdefault(key, value[0]) - - if "token" not in kwargs: - credentials = None - if credentials_path: - credentials = credentials_path - elif credentials_base64: - credentials = json.loads(base64.b64decode(credentials_base64).decode()) - else: - credentials = "anon" - kwargs["token"] = credentials - - fs = self.fs_class(**kwargs) + fs = self.fs_class(**gcs_filesystem_kwargs(params, kwargs)) try: endpoint: str = determine_endpoint(table, path_to_file) @@ -133,14 +118,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): parsed_uri = urlparse(uri) source_fields = parse_qs(parsed_uri.query) - access_key_id = source_fields.get("access_key_id") - if not access_key_id: - raise MissingConnectorOption("access_key_id", self.fs_name) - - secret_access_key = source_fields.get("secret_access_key") - if not secret_access_key: - raise MissingConnectorOption("secret_access_key", self.fs_name) - + fs_kwargs = s3_filesystem_kwargs(source_fields, self.fs_name) bucket_name, path_to_file = parse_uri(parsed_uri, table) if not bucket_name or not path_to_file: raise InvalidBlobTableError(self.fs_name) @@ -148,15 +126,6 @@ def dlt_source(self, uri: str, table: str, **kwargs): bucket_url = f"{self.fs_protocol}://{bucket_name}/" endpoint_url = source_fields.get("endpoint_url") - fs_kwargs: dict = { - "key": access_key_id[0], - "secret": secret_access_key[0], - # S3FileSystem caches directory listings by default. Disable the cache so - # long-lived processes see objects created between incremental runs. - "use_listings_cache": False, - } - if endpoint_url: - fs_kwargs["endpoint_url"] = endpoint_url[0] fs = self.fs_class(**fs_kwargs) @@ -193,35 +162,6 @@ def fs_name(self) -> str: return "S3" -def _azure_kwargs(auth: AzureBlobAuth): - """Return AzureBlobAuth information as dictionary. - - The ingestr-style short names already match adlfs kwargs, so they pass - straight through; only the supplied ones are forwarded. ``adlfs`` is - imported lazily so the CLI ``--help`` and every non-Azure path never load - the Azure SDK (matching the s3fs/gcsfs deferred-import convention). - """ - - kwargs = {"account_name": auth.account_name} - if auth.account_key is not None: - kwargs["account_key"] = auth.account_key - if auth.sas_token is not None: - kwargs["sas_token"] = auth.sas_token - if auth.tenant_id is not None: - kwargs["tenant_id"] = auth.tenant_id - if auth.client_id is not None: - kwargs["client_id"] = auth.client_id - if auth.client_secret is not None: - kwargs["client_secret"] = auth.client_secret - if auth.account_host is not None: - kwargs["account_host"] = auth.account_host - if auth.connection_string is not None: - kwargs["connection_string"] = auth.connection_string - if auth.api_version is not None: - kwargs["api_version"] = auth.api_version - return kwargs - - class AzureSource(FilesystemSource): """Azure Blob Storage / ADLS Gen2 source (``az://``, ``adls://``, ``abfss://``). @@ -259,7 +199,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): bucket_url = f"az://{bucket_name}" - kwargs.update(_azure_kwargs(auth)) + kwargs.update(azure_blob_filesystem_kwargs(auth)) fs = self.fs_class(**kwargs) try: diff --git a/src/dlt_filesystem/staging.py b/src/dlt_filesystem/staging.py new file mode 100644 index 000000000..3c8a35f2e --- /dev/null +++ b/src/dlt_filesystem/staging.py @@ -0,0 +1,193 @@ +"""Run-scoped materialization of remote objects for random-access consumers.""" + +from __future__ import annotations + +import logging +import shutil +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, Iterator +from urllib.parse import parse_qs, unquote, urlsplit, urlunsplit + +from dlt_filesystem.util.auth import ( + azure_blob_filesystem_kwargs, + gcs_filesystem_kwargs, + parse_azure_blob_auth, + s3_filesystem_kwargs, +) + +if TYPE_CHECKING: + from fsspec import AbstractFileSystem + +logger = logging.getLogger(__name__) + +# Source schemes that can be staged, mapped onto the backend that serves them. +# The alias schemes are the same ones the matching sources accept, so a URI that +# reads a file also stages an object. +STAGING_BACKENDS: dict[str, str] = { + "abfss": "az", + "adls": "az", + "az": "az", + "gs": "gs", + "r2": "s3", + "s3": "s3", +} +COPY_CHUNK_SIZE = 8 * 1024 * 1024 + + +class RemoteObjectNotFoundError(FileNotFoundError): + """A remote object selected for staging is missing or is not a file.""" + + def __init__(self, location: str): + super().__init__(f"Remote object not found: {location}") + + +class RemoteObjectSizeError(OSError): + """A staged copy does not match the size reported before download.""" + + +@dataclass(frozen=True) +class RemoteObject: + """One object on a remote filesystem, with the storage options that reach it.""" + + backend: str + path: str + safe_location: str + storage_options: dict[str, list[str]] = field(repr=False) + + @classmethod + def from_uri(cls, uri: str) -> RemoteObject: + """Parse a filesystem URI that names one object. + + The URI is the same one the matching source reads files from: the path + addresses the object and the query carries the storage options, so + ``s3://bucket/path/object?access_key_id=...`` needs no second encoding + layer. + """ + parsed = urlsplit(uri) + scheme = parsed.scheme.lower() + backend = STAGING_BACKENDS.get(scheme) + if backend is None: + supported = ", ".join(sorted(STAGING_BACKENDS)) + raise ValueError( + f"Unsupported storage scheme {scheme!r} for object staging; " + f"expected one of: {supported}." + ) + if parsed.username is not None or parsed.password is not None: + raise ValueError( + "Storage credentials must be supplied as URI query parameters, " + "not embedded in the object URI." + ) + if parsed.fragment: + raise ValueError( + "The object URI must not carry a '#' fragment; it names one " + "object rather than selecting a reader." + ) + if not parsed.netloc or not parsed.path.lstrip("/"): + raise ValueError( + "The object URI must identify one object as " + "':///'." + ) + + safe_location = urlunsplit((scheme, parsed.netloc, parsed.path, "", "")) + path = f"{unquote(parsed.netloc)}{unquote(parsed.path)}" + return cls( + backend=backend, + path=path, + safe_location=safe_location, + storage_options=parse_qs(parsed.query), + ) + + +def _filesystem_for(remote: RemoteObject) -> AbstractFileSystem: + """Construct the same authorized fsspec client as the matching source.""" + if remote.backend == "s3": + from s3fs import S3FileSystem + + return S3FileSystem(**s3_filesystem_kwargs(remote.storage_options)) + if remote.backend == "gs": + from gcsfs import GCSFileSystem + + return GCSFileSystem(**gcs_filesystem_kwargs(remote.storage_options)) + if remote.backend == "az": + from adlfs import AzureBlobFileSystem + + auth = parse_azure_blob_auth(remote.storage_options) + return AzureBlobFileSystem(**azure_blob_filesystem_kwargs(auth)) + raise AssertionError(f"Unhandled staging backend: {remote.backend}") + + +def _expected_size(info: dict[str, Any]) -> int | None: + """Return fsspec's normalized object size when it is available.""" + value = info.get("size") + if isinstance(value, int) and value >= 0: + return value + return None + + +@contextmanager +def materialize_remote_object( + remote: RemoteObject, + *, + filename: str, + staging_root: str | Path | None = None, +) -> Iterator[Path]: + """Download one remote object and remove its run-scoped copy on every exit.""" + root = Path(staging_root) if staging_root is not None else None + if root is not None: + root.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory( + prefix="omniload-staged-object-", dir=root + ) as staging_dir: + local_path = Path(staging_dir) / filename + filesystem = _filesystem_for(remote) + + try: + info = filesystem.info(remote.path) + except FileNotFoundError as exc: + raise RemoteObjectNotFoundError(remote.safe_location) from exc + if info.get("type") == "directory": + raise RemoteObjectNotFoundError(remote.safe_location) + + expected_size = _expected_size(info) + if expected_size is not None: + available = shutil.disk_usage(staging_dir).free + if expected_size > available: + raise OSError( + f"Not enough local disk space to stage {remote.safe_location}: " + f"need {expected_size} bytes, have {available} bytes." + ) + + logger.info( + "Staging remote object %s (%s bytes)", + remote.safe_location, + expected_size if expected_size is not None else "unknown", + ) + + downloaded = 0 + try: + with ( + filesystem.open(remote.path, "rb") as source, + local_path.open("wb") as destination, + ): + while chunk := source.read(COPY_CHUNK_SIZE): + downloaded += len(chunk) + destination.write(chunk) + except FileNotFoundError as exc: + raise RemoteObjectNotFoundError(remote.safe_location) from exc + + if expected_size is not None and downloaded != expected_size: + raise RemoteObjectSizeError( + f"Remote object size changed while staging {remote.safe_location}: " + f"expected {expected_size} bytes, downloaded {downloaded} bytes." + ) + + logger.info( + "Staged remote object %s (%d bytes)", + remote.safe_location, + downloaded, + ) + yield local_path diff --git a/src/dlt_filesystem/util/auth.py b/src/dlt_filesystem/util/auth.py index 497681ef5..16d912419 100644 --- a/src/dlt_filesystem/util/auth.py +++ b/src/dlt_filesystem/util/auth.py @@ -1,11 +1,67 @@ +import base64 +import json from dataclasses import dataclass -from typing import Optional +from typing import Any, Optional from dlt_filesystem.error import MissingConnectorOption AZURE_SERVICE_PRINCIPAL_FIELDS = ("tenant_id", "client_id", "client_secret") +def _first(params: dict[str, list[str]], key: str) -> Optional[str]: + """Return the first query-parameter value, matching existing URI semantics.""" + return params.get(key, [None])[0] + + +def s3_filesystem_kwargs( + params: dict[str, list[str]], connector: str = "S3" +) -> dict[str, Any]: + """Translate omniload S3 URI parameters into ``s3fs`` arguments.""" + access_key_id = _first(params, "access_key_id") + if not access_key_id: + raise MissingConnectorOption("access_key_id", connector) + + secret_access_key = _first(params, "secret_access_key") + if not secret_access_key: + raise MissingConnectorOption("secret_access_key", connector) + + kwargs: dict[str, Any] = { + "key": access_key_id, + "secret": secret_access_key, + # S3FileSystem caches directory listings by default. Disable the cache so + # long-lived processes see objects created between incremental runs. + "use_listings_cache": False, + } + endpoint_url = _first(params, "endpoint_url") + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + return kwargs + + +def gcs_filesystem_kwargs( + params: dict[str, list[str]], inherited: dict[str, Any] | None = None +) -> dict[str, Any]: + """Translate omniload GCS URI parameters into ``gcsfs`` arguments.""" + kwargs = dict(inherited or {}) + remaining = {key: list(values) for key, values in params.items()} + credentials_path = _first(remaining, "credentials_path") + credentials_base64 = _first(remaining, "credentials_base64") + remaining.pop("credentials_path", None) + remaining.pop("credentials_base64", None) + + for key, values in remaining.items(): + kwargs.setdefault(key, values[0]) + + if "token" not in kwargs: + if credentials_path: + kwargs["token"] = credentials_path + elif credentials_base64: + kwargs["token"] = json.loads(base64.b64decode(credentials_base64).decode()) + else: + kwargs["token"] = "anon" # noqa: S105 - gcsfs anonymous-access sentinel + return kwargs + + @dataclass class AzureBlobAuth: """Resolved Azure blob-storage credentials parsed from URI query params. @@ -66,16 +122,15 @@ def parse_azure_blob_auth(params: dict) -> AzureBlobAuth: rather than silently picking one. """ - def one(key: str) -> Optional[str]: - return params.get(key, [None])[0] - - connection_string = one("connection_string") - api_version = one("api_version") - account_name = one("account_name") - account_key = one("account_key") - sas_token = one("sas_token") - sp_values = {field: one(field) for field in AZURE_SERVICE_PRINCIPAL_FIELDS} - account_host = one("account_host") + connection_string = _first(params, "connection_string") + api_version = _first(params, "api_version") + account_name = _first(params, "account_name") + account_key = _first(params, "account_key") + sas_token = _first(params, "sas_token") + sp_values = { + field: _first(params, field) for field in AZURE_SERVICE_PRINCIPAL_FIELDS + } + account_host = _first(params, "account_host") if connection_string is not None: conflicting_fields = [ @@ -139,3 +194,25 @@ def one(key: str) -> Optional[str]: api_version=api_version, **sp_values, ) + + +def azure_blob_filesystem_kwargs(auth: AzureBlobAuth) -> dict[str, Any]: + """Translate parsed Azure credentials into ``adlfs`` arguments.""" + kwargs: dict[str, Any] = {"account_name": auth.account_name} + if auth.account_key is not None: + kwargs["account_key"] = auth.account_key + if auth.sas_token is not None: + kwargs["sas_token"] = auth.sas_token + if auth.tenant_id is not None: + kwargs["tenant_id"] = auth.tenant_id + if auth.client_id is not None: + kwargs["client_id"] = auth.client_id + if auth.client_secret is not None: + kwargs["client_secret"] = auth.client_secret + if auth.account_host is not None: + kwargs["account_host"] = auth.account_host + if auth.connection_string is not None: + kwargs["connection_string"] = auth.connection_string + if auth.api_version is not None: + kwargs["api_version"] = auth.api_version + return kwargs diff --git a/src/omniload/api.py b/src/omniload/api.py index 8ea82456e..47b58966e 100644 --- a/src/omniload/api.py +++ b/src/omniload/api.py @@ -61,10 +61,52 @@ def run_ingest(**kwargs) -> LoadInfo | None: jr = LoadRequest(**kwargs) if jr.pipelines_dir is not None: - return _run_ingest(jr, jr.pipelines_dir, is_pipelines_dir_temp=False) + return _run_ingest_with_remote_database( + jr, jr.pipelines_dir, is_pipelines_dir_temp=False + ) with tempfile.TemporaryDirectory() as pipelines_dir: - return _run_ingest(jr, pipelines_dir, is_pipelines_dir_temp=True) + return _run_ingest_with_remote_database( + jr, pipelines_dir, is_pipelines_dir_temp=True + ) + + +def _run_ingest_with_remote_database( + jr: LoadRequest, pipelines_dir: str, *, is_pipelines_dir_temp: bool +) -> LoadInfo | None: + """Stage a remote SQLite/DuckDB source for the complete pipeline lifetime.""" + import dataclasses + + from omniload.source.sql_database.remote import ( + dry_run_database_uri, + parse_remote_database_uri, + stage_remote_database, + ) + + remote_database = parse_remote_database_uri(jr.source_uri, jr.source_table or "") + if remote_database is None: + return _run_ingest( + jr, pipelines_dir, is_pipelines_dir_temp=is_pipelines_dir_temp + ) + + if jr.dry_run: + # Validate the SQL source the real run uses, without downloading anything. + return _run_ingest( + dataclasses.replace(jr, source_uri=dry_run_database_uri(remote_database)), + pipelines_dir, + is_pipelines_dir_temp=is_pipelines_dir_temp, + ) + + staging_root = jr.remote_database_staging_root or pipelines_dir + with stage_remote_database( + remote_database, staging_root=staging_root + ) as staged_source_uri: + staged_request = dataclasses.replace(jr, source_uri=staged_source_uri) + return _run_ingest( + staged_request, + pipelines_dir, + is_pipelines_dir_temp=is_pipelines_dir_temp, + ) def _run_ingest( diff --git a/src/omniload/model.py b/src/omniload/model.py index 8b65befd1..992b70461 100644 --- a/src/omniload/model.py +++ b/src/omniload/model.py @@ -93,6 +93,8 @@ class LoadRequest: loader_file_size: int = 100000 schema_naming: SchemaNaming | str = SchemaNaming.default pipelines_dir: str | None = None + # Parent directory for run-scoped remote database staging directories. + remote_database_staging_root: str | None = None extract_parallelism: int = 5 sql_reflection_level: SqlReflectionLevel | str = SqlReflectionLevel.full sql_limit: int | None = None diff --git a/src/omniload/source/sql_database/remote.py b/src/omniload/source/sql_database/remote.py new file mode 100644 index 000000000..a2015694e --- /dev/null +++ b/src/omniload/source/sql_database/remote.py @@ -0,0 +1,151 @@ +"""Route filesystem URIs that name a SQLite or DuckDB file into the SQL source.""" + +from __future__ import annotations + +from contextlib import contextmanager +from pathlib import Path, PurePosixPath +from typing import Iterator, Optional +from urllib.parse import unquote, urlsplit + +from dlt_filesystem.source.impl.util import has_glob_magic +from dlt_filesystem.staging import ( + STAGING_BACKENDS, + RemoteObject, + materialize_remote_object, +) + +# File extensions that name a file-based database, mapped onto the engine they +# imply. `.db` is used by both engines, so it carries no implication and is +# resolved from the file header alone. +DATABASE_EXTENSIONS: dict[str, Optional[str]] = { + ".db": None, + ".ddb": "duckdb", + ".duckdb": "duckdb", + ".sqlite": "sqlite", + ".sqlite3": "sqlite", +} + +# Both engines start their files with a fixed marker: SQLite at offset 0, DuckDB +# behind an 8-byte checksum. Reading them identifies the engine of a staged file +# whatever its extension says. +SQLITE_MAGIC = b"SQLite format 3\x00" +DUCKDB_MAGIC = b"DUCK" +DUCKDB_MAGIC_OFFSET = 8 +DATABASE_HEADER_SIZE = 16 + + +def _names_database(carrier: str) -> bool: + """Return whether a URI or table carrier ends in a database extension.""" + path = unquote(urlsplit(carrier).path) + return PurePosixPath(path).suffix.lower() in DATABASE_EXTENSIONS + + +def parse_remote_database_uri(uri: str, table: str = "") -> Optional[RemoteObject]: + """Return the remote object for a filesystem URI naming a database file. + + ``s3://analytics/snapshots/events.duckdb`` reads as the object URI it is: + the same carrier the filesystem sources take, with credentials in its query. + Every other URI returns ``None`` and keeps its own source, including a + filesystem URI that selects a regular file. + """ + parsed = urlsplit(uri) + if parsed.scheme.lower() not in STAGING_BACKENDS: + return None + + if not _names_database(parsed.path): + # The filesystem sources also accept the object path on `--source-table`. + # A database cannot use that form, because the table names a table inside + # the database, so say which carrier takes the object. + if _names_database(table): + raise ValueError( + "Name a remote database on --source-uri, as " + "'s3://bucket/path/events.duckdb'. --source-table selects a table " + "inside the database, so it cannot also carry the object path." + ) + return None + # Glob syntax is written literally, so the still-encoded path is what decides: + # a percent-encoded `?` is part of an object name, not a wildcard. + if has_glob_magic(parsed.path): + raise ValueError( + "A database source names one object; wildcards select a file set " + "that no single database connection can open." + ) + + return RemoteObject.from_uri(uri) + + +def resolve_database_engine(local_path: Path, location: str) -> str: + """Identify the engine of a staged database file. + + The file header decides, so a mislabeled object still loads. Only an empty + file has no header to read, and it falls back to an unambiguous extension so + a freshly created, still-empty database stays loadable. + """ + with local_path.open("rb") as database: + header = database.read(DATABASE_HEADER_SIZE) + + if header.startswith(SQLITE_MAGIC): + return "sqlite" + if header[DUCKDB_MAGIC_OFFSET : DUCKDB_MAGIC_OFFSET + len(DUCKDB_MAGIC)] == ( + DUCKDB_MAGIC + ): + return "duckdb" + if header: + raise ValueError( + f"Cannot identify the database engine of {location}: its header " + f"matches neither SQLite nor DuckDB." + ) + + engine = _engine_from_extension(location) + if engine is None: + raise ValueError( + f"Cannot identify the database engine of {location}: the object is " + f"empty, and its extension names both SQLite and DuckDB." + ) + return engine + + +def _engine_from_extension(location: str) -> Optional[str]: + """Return the engine a location's extension implies, if it implies one.""" + suffix = PurePosixPath(unquote(urlsplit(location).path)).suffix.lower() + return DATABASE_EXTENSIONS.get(suffix) + + +def dry_run_database_uri(remote: RemoteObject) -> str: + """Return the SQL URI a dry run validates in place of the object URI. + + A dry run stops before extraction, so the object is never downloaded and its + engine cannot be read. Only the scheme matters here: it routes validation + through the SQL source the real run uses, rather than the filesystem source + the URI came from. Nothing connects, so an extension that names both engines + resolves to either one. + """ + engine = _engine_from_extension(remote.safe_location) or "sqlite" + return f"{engine}:///{remote.path}" + + +def _database_uri(engine: str, path: Path) -> str: + """Build an absolute SQLAlchemy file URI, including on Windows.""" + normalized = path.resolve().as_posix() + if "?" in normalized: + raise ValueError( + "Remote database staging paths must not contain '?', which SQLAlchemy " + "treats as the start of connection query parameters." + ) + return f"{engine}:///{normalized}" + + +@contextmanager +def stage_remote_database( + remote: RemoteObject, + *, + staging_root: str | Path | None = None, +) -> Iterator[str]: + """Yield a local SQL URI backed by a run-scoped copy of the remote database.""" + with materialize_remote_object( + remote, + filename="database", + staging_root=staging_root, + ) as local_path: + engine = resolve_database_engine(local_path, remote.safe_location) + yield _database_uri(engine, local_path) diff --git a/tests/dlt_filesystem/test_staging.py b/tests/dlt_filesystem/test_staging.py new file mode 100644 index 000000000..1370eef72 --- /dev/null +++ b/tests/dlt_filesystem/test_staging.py @@ -0,0 +1,161 @@ +import logging +from typing import Any, ClassVar + +import pytest +from fsspec.implementations.memory import MemoryFileSystem + +from dlt_filesystem.staging import ( + RemoteObject, + RemoteObjectNotFoundError, + RemoteObjectSizeError, + _filesystem_for, + materialize_remote_object, +) + + +class _IsolatedMemoryFileSystem(MemoryFileSystem): + """Keep staging-test objects out of fsspec's process-global memory store.""" + + store: ClassVar[dict[str, Any]] = {} + + def __init__(self) -> None: + super().__init__() + self.pseudo_dirs = [""] + + +@pytest.fixture +def memory_filesystem(): + filesystem = _IsolatedMemoryFileSystem() + filesystem.store.clear() + filesystem.pipe_file("bucket/path/database.duckdb", b"database bytes") + yield filesystem + filesystem.store.clear() + + +def _remote() -> RemoteObject: + return RemoteObject.from_uri( + "s3://bucket/path/database.duckdb" + "?access_key_id=access&secret_access_key=do-not-log" + ) + + +def test_remote_object_rejects_a_scheme_it_cannot_stage(): + with pytest.raises(ValueError, match="Unsupported storage scheme 'sftp'"): + RemoteObject.from_uri("sftp://host/path/database.duckdb") + + +def test_filesystem_for_serves_an_s3_compatible_alias_through_s3fs(): + """R2 is S3-compatible, so its objects are staged with the S3 client.""" + from s3fs import S3FileSystem + + remote = RemoteObject.from_uri( + "r2://bucket/path/database.duckdb" + "?access_key_id=access&secret_access_key=secret" + "&endpoint_url=https://account.r2.cloudflarestorage.com" + ) + + assert remote.backend == "s3" + assert isinstance(_filesystem_for(remote), S3FileSystem) + + +def test_materialize_remote_object_lives_for_context_and_is_removed( + tmp_path, memory_filesystem, mocker, caplog +): + mocker.patch( + "dlt_filesystem.staging._filesystem_for", return_value=memory_filesystem + ) + + with caplog.at_level(logging.INFO, logger="dlt_filesystem.staging"): + with materialize_remote_object( + _remote(), filename="database", staging_root=tmp_path + ) as local_path: + assert local_path.exists() + assert local_path.read_bytes() == b"database bytes" + assert local_path.parent.parent == tmp_path + + assert not local_path.exists() + assert list(tmp_path.iterdir()) == [] + assert "s3://bucket/path/database.duckdb" in caplog.text + assert "do-not-log" not in caplog.text + + +def test_materialize_remote_object_cleans_up_when_consumer_fails( + tmp_path, memory_filesystem, mocker +): + mocker.patch( + "dlt_filesystem.staging._filesystem_for", return_value=memory_filesystem + ) + + with pytest.raises(RuntimeError, match="extract failed"): + with materialize_remote_object( + _remote(), filename="database", staging_root=tmp_path + ) as local_path: + assert local_path.exists() + raise RuntimeError("extract failed") + + assert list(tmp_path.iterdir()) == [] + + +def test_materialize_remote_object_reports_missing_object_without_credentials( + tmp_path, memory_filesystem, mocker +): + mocker.patch( + "dlt_filesystem.staging._filesystem_for", return_value=memory_filesystem + ) + remote = RemoteObject.from_uri( + "s3://bucket/path/missing.duckdb" + "?access_key_id=access&secret_access_key=do-not-report" + ) + + with pytest.raises(RemoteObjectNotFoundError) as excinfo: + with materialize_remote_object( + remote, filename="database", staging_root=tmp_path + ): + pass + + assert "s3://bucket/path/missing.duckdb" in str(excinfo.value) + assert "do-not-report" not in str(excinfo.value) + assert list(tmp_path.iterdir()) == [] + + +def test_materialize_remote_object_rejects_short_read( + tmp_path, memory_filesystem, mocker +): + mocker.patch( + "dlt_filesystem.staging._filesystem_for", return_value=memory_filesystem + ) + mocker.patch.object( + memory_filesystem, + "info", + return_value={"name": "database.duckdb", "size": 99, "type": "file"}, + ) + + with pytest.raises(RemoteObjectSizeError, match="expected 99 bytes"): + with materialize_remote_object( + _remote(), filename="database", staging_root=tmp_path + ): + pass + + assert list(tmp_path.iterdir()) == [] + + +def test_materialize_remote_object_checks_available_disk_before_download( + tmp_path, memory_filesystem, mocker +): + mocker.patch( + "dlt_filesystem.staging._filesystem_for", return_value=memory_filesystem + ) + mocker.patch( + "dlt_filesystem.staging.shutil.disk_usage", + return_value=mocker.Mock(free=1), + ) + open_spy = mocker.spy(memory_filesystem, "open") + + with pytest.raises(OSError, match="Not enough local disk space"): + with materialize_remote_object( + _remote(), filename="database", staging_root=tmp_path + ): + pass + + open_spy.assert_not_called() + assert list(tmp_path.iterdir()) == [] diff --git a/tests/main/filesystem/test_remote_integration.py b/tests/main/filesystem/test_remote_integration.py index 46eee6b99..094157059 100644 --- a/tests/main/filesystem/test_remote_integration.py +++ b/tests/main/filesystem/test_remote_integration.py @@ -1,11 +1,13 @@ import gzip import io import json +import sqlite3 from collections.abc import Callable +from contextlib import closing from dataclasses import dataclass from pathlib import Path from time import sleep -from urllib.parse import quote +from urllib.parse import parse_qsl, quote, urlencode from uuid import uuid4 import duckdb @@ -16,6 +18,7 @@ from dlt.sources.filesystem import FileItemDict from dlt_filesystem.source.error import NoFilesFoundError +from dlt_filesystem.staging import RemoteObjectNotFoundError from omniload import run_ingest from tests.util import invoke_ingest_command from tests.util.common import has_exception @@ -32,8 +35,14 @@ for backend in FAST_SOURCE_BACKENDS for file_format in ["csv", "jsonl", "parquet", "csv.gz"] ] + [pytest.param("gcs", "csv", id="gcs-csv")] +DATABASE_MATRIX = [ + pytest.param("s3", "duckdb", id="s3-duckdb"), + pytest.param("azure", "sqlite", id="azure-sqlite"), + pytest.param("gcs", "duckdb", id="gcs-duckdb"), +] CSV_ROWS = b"name,value\nAlice,1\nBob,2\n" +WIDGET_ROWS = [(1, "alpha"), (2, "beta"), (3, "gamma")] @dataclass(frozen=True) @@ -50,6 +59,11 @@ class RemoteFilesystemEmulator: def source_table(self, path: str) -> str: return f"{self.namespace}/{path}" + def object_uri(self, path: str) -> str: + """Address one object by the URI its storage service already names it with.""" + scheme, _, query = self.source_uri.partition("?") + return f"{scheme}{self.namespace}/{path}?{query}" + def destination_table(self, prefix: str, table: str = "data") -> str: return f"{self.namespace}/{prefix}/{table}" @@ -203,6 +217,21 @@ def format_payload(file_format: str) -> tuple[str, bytes]: raise ValueError(f"Unknown test format: {file_format}") +def database_payload(engine: str, path: Path) -> bytes: + """Build a small single-table database and return it as uploadable bytes.""" + if engine == "duckdb": + with duckdb.connect(path) as db: + db.execute("CREATE TABLE widgets (id INTEGER, name VARCHAR)") + db.executemany("INSERT INTO widgets VALUES (?, ?)", WIDGET_ROWS) + elif engine == "sqlite": + with closing(sqlite3.connect(path)) as db, db: + db.execute("CREATE TABLE widgets (id INTEGER, name TEXT)") + db.executemany("INSERT INTO widgets VALUES (?, ?)", WIDGET_ROWS) + else: + raise ValueError(f"Unknown test database engine: {engine}") + return path.read_bytes() + + def parquet_rows(objects: dict[str, bytes]) -> list[dict]: """Read all destination parquet data objects returned by an emulator SDK.""" parquet_objects = { @@ -303,6 +332,66 @@ def test_s3_missing_concrete_source_fails(s3_emulator, tmp_path, missing_bucket) assert "path/to/missing.csv" in str(result.exception) +@pytest.mark.parametrize( + "remote_filesystem,engine", DATABASE_MATRIX, indirect=["remote_filesystem"] +) +def test_remote_database_source_stages_and_cleans_up( + remote_filesystem, engine, tmp_path +): + """A remote database object loads through SQL extraction and leaves nothing staged. + + The object keeps the URI its storage service names it with, so the extension + routes it to the SQL source and ``--source-table`` selects a table inside it. + """ + key = f"databases/source.{engine}" + remote_filesystem.upload( + key, database_payload(engine, tmp_path / f"source.{engine}") + ) + + staging_root = tmp_path / "staging" + # The destination file name doubles as the DuckDB catalog name, so keep it + # distinct from the `remote` dataset the rows land in. + db_path = tmp_path / f"{remote_filesystem.backend}-remote.duckdb" + result = run_ingest( + source_uri=remote_filesystem.object_uri(key), + dest_uri=f"duckdb:///{db_path}", + source_table="main.widgets", + dest_table="remote.widgets", + remote_database_staging_root=str(staging_root), + progress="log", + ) + + assert result is not None + assert duckdb_table_cardinality(db_path, "remote.widgets") == 3 + assert list(staging_root.iterdir()) == [] + + +def test_missing_remote_database_fails_and_removes_staging(s3_emulator, tmp_path): + """A missing object reports its safe location and leaves no staged files.""" + staging_root = tmp_path / "staging" + options = dict(parse_qsl(s3_emulator.source_uri.partition("?")[2])) + options["secret_access_key"] = "do-not-report" # noqa: S105 - leak canary + source_uri = ( + f"s3://{s3_emulator.namespace}/databases/missing.duckdb?{urlencode(options)}" + ) + + with pytest.raises(RemoteObjectNotFoundError) as excinfo: + run_ingest( + source_uri=source_uri, + dest_uri=f"duckdb:///{tmp_path / 'missing.duckdb'}", + source_table="main.widgets", + dest_table="remote.widgets", + remote_database_staging_root=str(staging_root), + progress="log", + ) + + assert f"s3://{s3_emulator.namespace}/databases/missing.duckdb" in str( + excinfo.value + ) + assert "do-not-report" not in str(excinfo.value) + assert list(staging_root.iterdir()) == [] + + @pytest.mark.parametrize("remote_filesystem", DESTINATION_BACKENDS, indirect=True) def test_remote_destination_writes_native_parquet(remote_filesystem): """S3 and Azure destinations write data readable through their native SDKs.""" diff --git a/tests/main/filesystem/test_remote_read.py b/tests/main/filesystem/test_remote_read.py index a45414538..b13d6d5e4 100644 --- a/tests/main/filesystem/test_remote_read.py +++ b/tests/main/filesystem/test_remote_read.py @@ -326,6 +326,18 @@ def test_touch_s3_filesystem_without_secret_access_key(): assert exc_info.match("secret_access_key is required") +def test_touch_s3_filesystem_without_bucket_or_credentials_reports_credentials(): + """Credential validation keeps precedence for a doubly invalid S3 source.""" + source_uri = "s3://" + factory = SourceDestinationFactory(source_uri, "file://") + source = factory.get_source() + + with pytest.raises(MissingConnectorOption) as exc_info: + source.dlt_source(uri=source_uri, table="") + + assert exc_info.match("access_key_id is required") + + @pytest.mark.parametrize( "source_uri", SCHEMES_WITH_HOST, ids=[str(item) for item in SCHEMES_WITH_HOST] ) diff --git a/tests/main/test_remote_database.py b/tests/main/test_remote_database.py new file mode 100644 index 000000000..ecc0d1b6e --- /dev/null +++ b/tests/main/test_remote_database.py @@ -0,0 +1,358 @@ +import sqlite3 +from pathlib import Path +from typing import Any, ClassVar +from urllib.parse import urlencode + +import duckdb +import pytest +from fsspec.implementations.memory import MemoryFileSystem +from sqlalchemy.engine import make_url + +from dlt_filesystem.staging import RemoteObject +from omniload import run_ingest +from omniload.source.sql_database.remote import ( + _database_uri, + dry_run_database_uri, + parse_remote_database_uri, + resolve_database_engine, +) + + +class _IsolatedMemoryFileSystem(MemoryFileSystem): + """Keep remote-database fixtures out of fsspec's shared memory store.""" + + store: ClassVar[dict[str, Any]] = {} + + def __init__(self) -> None: + super().__init__() + self.pseudo_dirs = [""] + + +def _remote_uri(object_name: str = "database.duckdb") -> str: + query = urlencode( + { + "access_key_id": "access", + "secret_access_key": "secret", + } + ) + return f"s3://bucket/path/{object_name}?{query}" + + +def test_parse_remote_database_uri_reads_the_object_uri_as_written(): + query = urlencode( + { + "access_key_id": "+/=&?", + "secret_access_key": "secret+/=&?", + } + ) + + remote = parse_remote_database_uri( + f"s3://bucket/reports/data%3Fpart%3D1.duckdb?{query}" + ) + + assert remote is not None + assert remote.backend == "s3" + assert remote.path == "bucket/reports/data?part=1.duckdb" + assert remote.safe_location == "s3://bucket/reports/data%3Fpart%3D1.duckdb" + assert remote.storage_options == { + "access_key_id": ["+/=&?"], + "secret_access_key": ["secret+/=&?"], + } + assert "+/=&?" not in repr(remote) + assert "secret+/=&?" not in repr(remote) + + +@pytest.mark.parametrize( + ("uri", "backend"), + [ + ("s3://bucket/events.duckdb", "s3"), + ("r2://bucket/events.ddb", "s3"), + ("gs://bucket/events.sqlite", "gs"), + ("az://container/events.sqlite3", "az"), + ("adls://container/events.db", "az"), + ("abfss://container/EVENTS.DuckDB", "az"), + ], +) +def test_parse_remote_database_uri_covers_every_staging_scheme(uri, backend): + remote = parse_remote_database_uri(uri) + + assert remote is not None + assert remote.backend == backend + + +@pytest.mark.parametrize( + "uri", + [ + # Database schemes address their own files. + "sqlite:///tmp/source.sqlite", + "duckdb:///tmp/source.duckdb", + "md://analytics?token=secret", + "motherduck://analytics?token=secret", + # Local files keep the local source. + "file://path/to/source.sqlite", + # Regular objects keep their filesystem source. + "s3://bucket/path/events.csv", + "gs://bucket/path/events.parquet", + "s3://bucket/path/*.csv", + ], +) +def test_parse_remote_database_uri_leaves_other_sources_alone(uri): + assert parse_remote_database_uri(uri) is None + + +@pytest.mark.parametrize( + ("uri", "message"), + [ + ("s3://bucket/path/*.duckdb", "names one object"), + ("s3://bucket/path/data[12].duckdb", "names one object"), + ("s3://bucket/path/data.duckdb#csv", "must not carry a '#' fragment"), + ("s3://access:secret@bucket/data.duckdb", "not embedded in the object URI"), + ("s3:///data.duckdb", "must identify one object"), + ], +) +def test_parse_remote_database_uri_rejects_ambiguous_selections(uri, message): + with pytest.raises(ValueError, match=message): + parse_remote_database_uri(uri) + + +def test_parse_remote_database_uri_keeps_an_encoded_wildcard_literal(): + remote = parse_remote_database_uri("s3://bucket/path/data%2A.duckdb") + + assert remote is not None + assert remote.path == "bucket/path/data*.duckdb" + + +@pytest.mark.parametrize( + "table", + ["bucket/path/events.duckdb", "s3://bucket/path/events.sqlite"], +) +def test_parse_remote_database_uri_rejects_the_split_form(table): + """The object path belongs on the URI: the table names a table inside it.""" + with pytest.raises(ValueError, match="Name a remote database on --source-uri"): + parse_remote_database_uri("s3://?access_key_id=access", table) + + +def test_parse_remote_database_uri_leaves_the_split_form_for_files_alone(): + assert parse_remote_database_uri("s3://", "bucket/path/events.csv") is None + + +def test_resolve_database_engine_reads_the_file_header(tmp_path): + sqlite_path = tmp_path / "opaque" + with sqlite3.connect(sqlite_path) as db: + db.execute("CREATE TABLE widgets (id INTEGER)") + duckdb_path = tmp_path / "opaque.duckdb" + with duckdb.connect(duckdb_path) as db: + db.execute("CREATE TABLE widgets (id INTEGER)") + + # The extension of the remote object never overrides the header. + assert resolve_database_engine(sqlite_path, "s3://bucket/a.duckdb") == "sqlite" + assert resolve_database_engine(duckdb_path, "s3://bucket/a.sqlite") == "duckdb" + + +def test_resolve_database_engine_falls_back_only_for_an_empty_file(tmp_path): + """A fresh, still-empty database has no header to read.""" + empty = tmp_path / "empty" + empty.touch() + + assert resolve_database_engine(empty, "s3://bucket/fresh.sqlite") == "sqlite" + assert resolve_database_engine(empty, "s3://bucket/fresh.ddb") == "duckdb" + + with pytest.raises(ValueError, match="the object is empty"): + resolve_database_engine(empty, "s3://bucket/fresh.db") + + +@pytest.mark.parametrize( + "location", + ["s3://bucket/export.db", "s3://bucket/export.sqlite"], +) +def test_resolve_database_engine_reports_an_unidentifiable_file(tmp_path, location): + """A non-empty file is judged by its header, never by its extension.""" + not_a_database = tmp_path / "opaque" + not_a_database.write_bytes(b"id,name\n1,alpha\n") + + with pytest.raises(ValueError, match="matches neither SQLite nor DuckDB"): + resolve_database_engine(not_a_database, location) + + +def test_database_uri_rejects_question_mark_in_staged_path(tmp_path): + with pytest.raises(ValueError, match="must not contain '\\?'"): + _database_uri("sqlite", tmp_path / "staging ? root" / "database") + + +def test_run_ingest_keeps_staged_database_through_pipeline_and_cleans_it_up( + tmp_path, mocker +): + filesystem = _IsolatedMemoryFileSystem() + filesystem.store.clear() + filesystem.pipe_file("bucket/path/database.duckdb", b"\x00" * 8 + b"DUCK" + b"rest") + mocker.patch("dlt_filesystem.staging._filesystem_for", return_value=filesystem) + staged_paths = [] + + def fake_run(jr, pipelines_dir, *, is_pipelines_dir_temp): + staged_url = make_url(jr.source_uri) + assert staged_url.drivername == "duckdb" + staged_path = staged_url.database + assert staged_path is not None + staged_paths.append(staged_path) + with open(staged_path, "rb") as staged: + assert staged.read() == b"\x00" * 8 + b"DUCK" + b"rest" + return "loaded" + + mocker.patch("omniload.api._run_ingest", side_effect=fake_run) + + result = run_ingest( + source_uri=_remote_uri(), + dest_uri="duckdb:///destination.duckdb", + source_table="main.widgets", + dest_table="out.widgets", + remote_database_staging_root=str(tmp_path), + ) + + assert result == "loaded" + assert len(staged_paths) == 1 + assert not Path(staged_paths[0]).exists() + assert list(tmp_path.iterdir()) == [] + filesystem.store.clear() + + +def test_run_ingest_remote_sqlite_to_duckdb_end_to_end(tmp_path, mocker): + source_path = tmp_path / "source.sqlite" + with sqlite3.connect(source_path) as db: + db.execute("CREATE TABLE widgets (id INTEGER, name TEXT)") + db.executemany( + "INSERT INTO widgets VALUES (?, ?)", + [(1, "alpha"), (2, "beta"), (3, "gamma")], + ) + + filesystem = _IsolatedMemoryFileSystem() + filesystem.store.clear() + filesystem.pipe_file("bucket/path/source.sqlite", source_path.read_bytes()) + mocker.patch("dlt_filesystem.staging._filesystem_for", return_value=filesystem) + staging_root = tmp_path / "staging #% café" + destination = tmp_path / "destination.duckdb" + + result = run_ingest( + source_uri=_remote_uri("source.sqlite"), + dest_uri=f"duckdb:///{destination}", + source_table="main.widgets", + dest_table="out.widgets", + remote_database_staging_root=str(staging_root), + progress="log", + ) + + assert result is not None + with duckdb.connect(destination) as db: + rows = db.sql("select id, name from out.widgets order by id").fetchall() + assert rows == [(1, "alpha"), (2, "beta"), (3, "gamma")] + assert list(staging_root.iterdir()) == [] + filesystem.store.clear() + + +def test_run_ingest_cleans_staged_database_up_when_pipeline_fails(tmp_path, mocker): + filesystem = _IsolatedMemoryFileSystem() + filesystem.store.clear() + filesystem.pipe_file("bucket/path/database.duckdb", b"\x00" * 8 + b"DUCK") + mocker.patch("dlt_filesystem.staging._filesystem_for", return_value=filesystem) + + def fail_run(jr, pipelines_dir, *, is_pipelines_dir_temp): + staged_path = make_url(jr.source_uri).database + assert staged_path is not None + assert Path(staged_path).exists() + raise RuntimeError("pipeline failed") + + mocker.patch("omniload.api._run_ingest", side_effect=fail_run) + + with pytest.raises(RuntimeError, match="pipeline failed"): + run_ingest( + source_uri=_remote_uri(), + dest_uri="duckdb:///destination.duckdb", + source_table="main.widgets", + dest_table="out.widgets", + remote_database_staging_root=str(tmp_path), + ) + + assert list(tmp_path.iterdir()) == [] + filesystem.store.clear() + + +@pytest.mark.parametrize( + "source_uri", + [ + "sqlite:///local.sqlite", + "duckdb:///local.duckdb", + "md://analytics?token=secret", + "s3://bucket/path/events.csv?access_key_id=access&secret_access_key=secret", + ], +) +def test_run_ingest_other_sources_do_not_construct_stager(source_uri, mocker): + stage = mocker.patch("omniload.source.sql_database.remote.stage_remote_database") + mocker.patch("omniload.api._run_ingest", return_value="loaded") + + assert ( + run_ingest( + source_uri=source_uri, + dest_uri="duckdb:///destination.duckdb", + source_table="main.widgets", + dest_table="out.widgets", + ) + == "loaded" + ) + stage.assert_not_called() + + +def test_run_ingest_dry_run_does_not_download_remote_database(tmp_path, mocker): + filesystem = mocker.patch("dlt_filesystem.staging._filesystem_for") + + assert ( + run_ingest( + source_uri=_remote_uri(), + dest_uri=f"duckdb:///{tmp_path / 'destination.duckdb'}", + source_table="main.widgets", + dest_table="out.widgets", + dry_run=True, + progress="log", + ) + is None + ) + filesystem.assert_not_called() + + +def test_run_ingest_dry_run_validates_the_sql_source(tmp_path, mocker): + """A dry run must validate the source the real run uses, not the storage one.""" + captured = [] + mocker.patch( + "omniload.api._run_ingest", + side_effect=lambda jr, *args, **kwargs: captured.append(jr.source_uri), + ) + + run_ingest( + source_uri=_remote_uri(), + dest_uri=f"duckdb:///{tmp_path / 'destination.duckdb'}", + source_table="main.widgets", + dest_table="out.widgets", + dry_run=True, + ) + + assert captured == ["duckdb:///bucket/path/database.duckdb"] + + +@pytest.mark.parametrize( + ("object_name", "expected"), + [ + ("events.duckdb", "duckdb:///bucket/path/events.duckdb"), + ("events.sqlite3", "sqlite:///bucket/path/events.sqlite3"), + # Nothing connects during a dry run, so an ambiguous name may resolve + # to either engine. + ("events.db", "sqlite:///bucket/path/events.db"), + ], +) +def test_dry_run_database_uri_names_a_sql_source(object_name, expected): + remote = parse_remote_database_uri(_remote_uri(object_name)) + + assert remote is not None + assert dry_run_database_uri(remote) == expected + + +def test_remote_object_rejects_embedded_credentials(): + with pytest.raises(ValueError, match="not embedded in the object URI"): + RemoteObject.from_uri("s3://access:secret@bucket/database.duckdb") diff --git a/tests/main/test_sources.py b/tests/main/test_sources.py index 450062ce7..b48ca502e 100644 --- a/tests/main/test_sources.py +++ b/tests/main/test_sources.py @@ -141,6 +141,37 @@ def sql_table( self.assertIn("token:***", captured_uri) self.assertIn("hostname", captured_uri) + def test_motherduck_uri_rewrite_is_unchanged(self): + captured_uri = "" + + def sql_table(credentials: ConnectionStringCredentials, **kwargs): + nonlocal captured_uri + captured_uri = credentials.to_native_representation() + return dlt.resource() + + source = SqlSourceRouter(table_builder=sql_table) + source.dlt_source("md://analytics?token=secret", "main.widgets") + + self.assertEqual( + captured_uri, + "duckdb:///md:analytics?motherduck_token=secret", + ) + + def test_cratedb_uri_rewrite_is_unchanged(self): + captured_uri = "" + + def sql_table(credentials: ConnectionStringCredentials, **kwargs): + nonlocal captured_uri + captured_uri = credentials.to_native_representation() + return dlt.resource() + + source = SqlSourceRouter(table_builder=sql_table) + source.dlt_source( + "cratedb://crate@localhost:4200/?sslmode=require", "doc.widgets" + ) + + self.assertEqual(captured_uri, "crate://crate@localhost:4200/?ssl=true") + class MongoDbSourceTest(unittest.TestCase): def test_sql_source_requires_two_fields_in_table(self):