diff --git a/docs/changelog.rst b/docs/changelog.rst index bc8cfea8b..ababb3879 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -23,6 +23,27 @@ Unreleased * Removed the undocumented ``sqlspec.exceptions.wrap_exceptions`` helper, superseded by the typed per-adapter exception handlers. +**Added:** + +* Storage pipelines expose ``resolve_destination()``, returning a + ``ResolvedStorageTarget(uri, protocol)`` without opening a database session. + Direct remote URIs retain their address, alias paths resolve relative to the + configured backend, and local paths become absolute. Backend options come + only from the method's explicit ``storage_options`` argument, not pipeline + writer defaults. + +**Breaking changes:** + +* Removed the unimplemented driver methods ``stage_artifact()``, + ``flush_staging_artifacts()``, and ``get_storage_job()``, and the exported + ``StorageLoadRequest`` and ``StagedArtifact`` types. Retain the + ``StorageBridgeJob`` returned by a storage operation instead of looking it up. +* Removed pipeline ``allocate_staging_artifacts()`` and + ``cleanup_staging_artifacts()``, the ``requires_staging_for_load`` and + ``staging_protocols`` capability settings, and the unused + ``storage_bridge.partitions_created`` diagnostic counter. Working storage + import, export, and per-operation partition telemetry remain available. + v0.62.2 - Litestar config lookup diagnostics --------------------------------------------- diff --git a/docs/reference/storage.rst b/docs/reference/storage.rst index 1058ec8be..a687e37e1 100644 --- a/docs/reference/storage.rst +++ b/docs/reference/storage.rst @@ -64,10 +64,12 @@ Pipelines .. autoclass:: sqlspec.storage.SyncStoragePipeline :members: + :inherited-members: :show-inheritance: .. autoclass:: sqlspec.storage.AsyncStoragePipeline :members: + :inherited-members: :show-inheritance: Registry @@ -88,19 +90,15 @@ Configuration Types :members: :show-inheritance: -.. autoclass:: sqlspec.storage.StorageLoadRequest - :members: - :show-inheritance: - -.. autoclass:: sqlspec.storage.StagedArtifact +.. autoclass:: sqlspec.storage.StorageTelemetry :members: :show-inheritance: -.. autoclass:: sqlspec.storage.StorageTelemetry +.. autoclass:: sqlspec.storage.StorageBridgeJob :members: :show-inheritance: -.. autoclass:: sqlspec.storage.StorageBridgeJob +.. autoclass:: sqlspec.storage.ResolvedStorageTarget :members: :show-inheritance: diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py index e52509992..6278975b0 100644 --- a/sqlspec/adapters/bigquery/config.py +++ b/sqlspec/adapters/bigquery/config.py @@ -170,8 +170,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]): supports_arrow_streaming: ClassVar[bool] = True supports_native_row_streaming: ClassVar[bool] = True supports_native_parquet_export: ClassVar[bool] = True - requires_staging_for_load: ClassVar[bool] = True - staging_protocols: "ClassVar[tuple[str, ...]]" = ("gs://",) _connection_context_class: "ClassVar[type[BigQueryConnectionContext]]" = BigQueryConnectionContext _session_factory_class: "ClassVar[type[_BigQuerySessionConnectionHandler]]" = _BigQuerySessionConnectionHandler _session_context_class: "ClassVar[type[BigQuerySessionContext]]" = BigQuerySessionContext diff --git a/sqlspec/adapters/spanner/config.py b/sqlspec/adapters/spanner/config.py index fc42029c4..9631345f9 100644 --- a/sqlspec/adapters/spanner/config.py +++ b/sqlspec/adapters/spanner/config.py @@ -255,7 +255,6 @@ class SpannerSyncConfig(SyncDatabaseConfig["SpannerConnection", "AbstractSession supports_native_arrow_import: ClassVar[bool] = True supports_native_parquet_export: ClassVar[bool] = False supports_native_parquet_import: ClassVar[bool] = False - requires_staging_for_load: ClassVar[bool] = False _connection_context_class: "ClassVar[type[SpannerConnectionContext]]" = SpannerConnectionContext _session_factory_class: "ClassVar[type[_SpannerSessionConnectionHandler]]" = _SpannerSessionConnectionHandler _session_context_class: "ClassVar[type[SpannerSessionContext]]" = SpannerSessionContext diff --git a/sqlspec/config.py b/sqlspec/config.py index c862cc628..9b6a1b415 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -943,8 +943,6 @@ class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]): supports_migration_schemas: "ClassVar[bool]" = False supports_native_parquet_import: "ClassVar[bool]" = False supports_native_parquet_export: "ClassVar[bool]" = False - requires_staging_for_load: "ClassVar[bool]" = False - staging_protocols: "ClassVar[tuple[str, ...]]" = () default_storage_profile: "ClassVar[str | None]" = None storage_partition_strategies: "ClassVar[tuple[str, ...]]" = ("fixed",) bind_key: "str | None" @@ -1349,8 +1347,6 @@ def _build_storage_capabilities(self) -> "StorageCapabilities": "arrow_import_enabled": bool(self.supports_native_arrow_import and arrow_dependency_ready), "parquet_export_enabled": bool(self.supports_native_parquet_export and parquet_dependency_ready), "parquet_import_enabled": bool(self.supports_native_parquet_import and parquet_dependency_ready), - "requires_staging_for_load": self.requires_staging_for_load, - "staging_protocols": list(self.staging_protocols), "partition_strategies": list(self.storage_partition_strategies), } if self.default_storage_profile is not None: diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 2aedca99e..12a89ba79 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -1648,39 +1648,6 @@ async def load_from_records( arrow_table = self._records_to_arrow_table(prepared_records, columns) return await self.load_from_arrow(table, arrow_table, overwrite=overwrite) - def stage_artifact(self, request: "dict[str, Any]") -> "dict[str, Any]": - """Provision staging metadata for adapters that require remote URIs. - - Args: - request: Staging request configuration. - - Returns: - Staging metadata dict. - """ - self._raise_storage_not_implemented("stage_artifact") - raise NotImplementedError - - def flush_staging_artifacts(self, artifacts: "list[dict[str, Any]]", *, error: Exception | None = None) -> None: - """Clean up staged artifacts after a job completes. - - Args: - artifacts: List of staging artifacts to clean up. - error: Optional error that triggered cleanup. - """ - if artifacts: - self._raise_storage_not_implemented("flush_staging_artifacts") - - def get_storage_job(self, job_id: str) -> "StorageBridgeJob | None": - """Fetch a previously created job handle. - - Args: - job_id: Job identifier. - - Returns: - StorageBridgeJob if found, None otherwise. - """ - return None - # ───────────────────────────────────────────────────────────────────────────── # UTILITY METHODS # ───────────────────────────────────────────────────────────────────────────── diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index 4f8a537e5..de6c1c096 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -1636,39 +1636,6 @@ def load_from_records( arrow_table = self._records_to_arrow_table(prepared_records, columns) return self.load_from_arrow(table, arrow_table, overwrite=overwrite) - def stage_artifact(self, request: "dict[str, Any]") -> "dict[str, Any]": - """Provision staging metadata for adapters that require remote URIs. - - Args: - request: Staging request configuration. - - Returns: - Staging metadata dict. - """ - self._raise_storage_not_implemented("stage_artifact") - raise NotImplementedError - - def flush_staging_artifacts(self, artifacts: "list[dict[str, Any]]", *, error: Exception | None = None) -> None: - """Clean up staged artifacts after a job completes. - - Args: - artifacts: List of staging artifacts to clean up. - error: Optional error that triggered cleanup. - """ - if artifacts: - self._raise_storage_not_implemented("flush_staging_artifacts") - - def get_storage_job(self, job_id: str) -> "StorageBridgeJob | None": - """Fetch a previously created job handle. - - Args: - job_id: Job identifier. - - Returns: - StorageBridgeJob if found, None otherwise. - """ - return None - # ───────────────────────────────────────────────────────────────────────────── # UTILITY METHODS # ───────────────────────────────────────────────────────────────────────────── diff --git a/sqlspec/storage/__init__.py b/sqlspec/storage/__init__.py index b21833a41..5285e7fdb 100644 --- a/sqlspec/storage/__init__.py +++ b/sqlspec/storage/__init__.py @@ -12,12 +12,11 @@ from sqlspec.storage.pipeline import ( AsyncStoragePipeline, PartitionStrategyConfig, - StagedArtifact, + ResolvedStorageTarget, StorageBridgeJob, StorageCapabilities, StorageDestination, StorageFormat, - StorageLoadRequest, StorageTelemetry, SyncStoragePipeline, create_storage_bridge_job, @@ -30,12 +29,11 @@ __all__ = ( "AsyncStoragePipeline", "PartitionStrategyConfig", - "StagedArtifact", + "ResolvedStorageTarget", "StorageBridgeJob", "StorageCapabilities", "StorageDestination", "StorageFormat", - "StorageLoadRequest", "StorageRegistry", "StorageTelemetry", "SyncStoragePipeline", diff --git a/sqlspec/storage/pipeline.py b/sqlspec/storage/pipeline.py index 3bdd7bd38..03b4ba5b8 100644 --- a/sqlspec/storage/pipeline.py +++ b/sqlspec/storage/pipeline.py @@ -3,19 +3,21 @@ from collections import deque from functools import partial from pathlib import Path -from time import perf_counter, time +from time import perf_counter from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast +from urllib.parse import unquote, urlparse from mypy_extensions import mypyc_attr from typing_extensions import NotRequired, TypedDict from sqlspec.exceptions import ImproperConfigurationError, StorageCapabilityError from sqlspec.storage._arrow_payload import StorageFormat, decode_arrow_payload, encode_arrow_payload +from sqlspec.storage._paths import FILE_PROTOCOL, FILE_SCHEME_PREFIX, strip_windows_drive_prefix from sqlspec.storage.errors import execute_async_storage_operation, execute_sync_storage_operation from sqlspec.storage.registry import StorageRegistry, storage_registry from sqlspec.utils.serializers import get_serializer_metrics, serialize_collection, to_json from sqlspec.utils.sync_tools import async_ -from sqlspec.utils.type_guards import supports_async_delete, supports_async_read_bytes, supports_async_write_bytes +from sqlspec.utils.type_guards import supports_async_read_bytes, supports_async_write_bytes from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: @@ -28,13 +30,12 @@ __all__ = ( "AsyncStoragePipeline", "PartitionStrategyConfig", - "StagedArtifact", + "ResolvedStorageTarget", "StorageBridgeJob", "StorageCapabilities", "StorageDestination", "StorageDiagnostics", "StorageFormat", - "StorageLoadRequest", "StorageTelemetry", "SyncStoragePipeline", "create_storage_bridge_job", @@ -57,8 +58,6 @@ class StorageCapabilities(TypedDict): arrow_import_enabled: bool parquet_export_enabled: bool parquet_import_enabled: bool - requires_staging_for_load: bool - staging_protocols: "list[str]" partition_strategies: "list[str]" default_storage_profile: NotRequired[str | None] @@ -72,27 +71,6 @@ class PartitionStrategyConfig(TypedDict, total=False): manifest_path: str -class StorageLoadRequest(TypedDict): - """Request describing a staging allocation.""" - - partition_id: str - destination_uri: str - ttl_seconds: int - correlation_id: str - source_uri: NotRequired[str] - - -class StagedArtifact(TypedDict): - """Metadata describing a staged artifact managed by the pipeline.""" - - partition_id: str - uri: str - cleanup_token: str - ttl_seconds: int - expires_at: float - correlation_id: str - - class StorageTelemetry(TypedDict, total=False): """Telemetry payload for storage bridge operations.""" @@ -109,6 +87,13 @@ class StorageTelemetry(TypedDict, total=False): bind_key: str +class ResolvedStorageTarget(NamedTuple): + """A storage destination resolved to an address and its backend protocol.""" + + uri: str + protocol: str + + class StorageBridgeJob(NamedTuple): """Handle representing a storage bridge operation.""" @@ -118,27 +103,19 @@ class StorageBridgeJob(NamedTuple): class _StorageBridgeMetrics: - __slots__ = ("bytes_written", "partitions_created") + __slots__ = ("bytes_written",) def __init__(self) -> None: self.bytes_written = 0 - self.partitions_created = 0 def record_bytes(self, count: int) -> None: self.bytes_written += max(count, 0) - def record_partitions(self, count: int) -> None: - self.partitions_created += max(count, 0) - def snapshot(self) -> "dict[str, int]": - return { - "storage_bridge.bytes_written": self.bytes_written, - "storage_bridge.partitions_created": self.partitions_created, - } + return {"storage_bridge.bytes_written": self.bytes_written} def reset(self) -> None: self.bytes_written = 0 - self.partitions_created = 0 _METRICS = _StorageBridgeMetrics() @@ -269,12 +246,6 @@ def _encode_arrow_payload( return encode_arrow_payload(table, format_choice, compression=compression, write_options=write_options) -def _delete_backend_sync(backend: "ObjectStoreProtocol", path: str, *, backend_name: str) -> None: - execute_sync_storage_operation( - partial(backend.delete_sync, path), backend=backend_name, operation="delete", path=path - ) - - def _write_backend_sync(backend: "ObjectStoreProtocol", path: str, payload: bytes, *, backend_name: str) -> None: execute_sync_storage_operation( partial(backend.write_bytes_sync, path, payload), backend=backend_name, operation="write_bytes", path=path @@ -370,6 +341,39 @@ def _backend( self._resolved_backend_cache[cache_key] = resolved return resolved + def resolve_destination( + self, destination: StorageDestination, storage_options: "dict[str, Any] | None" = None + ) -> ResolvedStorageTarget: + """Resolve a destination without opening a database session or reading an object. + + Direct remote URIs retain their address. Alias paths resolve relative to + the registered backend. Local paths resolve to absolute filesystem paths + through the backend's path checks. + + Args: + destination: Remote URI, local path, or ``alias://name/path``. + storage_options: Explicit backend options. Pipeline writer defaults + are not inherited by this method. + + Returns: + The resolved address and backend protocol. + + Raises: + ImproperConfigurationError: If the destination or alias is invalid. + StoragePathTraversalError: If the backend rejects the local path. + """ + backend, path, _backend_name = self._backend(destination, storage_options) + destination_str = str(destination) + if destination_str.startswith("alias://"): + uri = backend.resolve_uri(path) + elif backend.protocol == FILE_PROTOCOL: + if destination_str.startswith(FILE_SCHEME_PREFIX): + path = strip_windows_drive_prefix(unquote(urlparse(destination_str).path)) + uri = backend.resolve_uri(Path(path).expanduser().resolve()) + else: + uri = destination_str + return ResolvedStorageTarget(uri, backend.protocol) + @mypyc_attr(allow_interpreted_subclasses=True) class SyncStoragePipeline(_StoragePipelineBase): @@ -449,38 +453,6 @@ def stream_read( backend, path, _backend_name = self._backend(source, storage_options) return backend.stream_read_sync(path, chunk_size=chunk_size) - def allocate_staging_artifacts(self, requests: "list[StorageLoadRequest]") -> "list[StagedArtifact]": - """Allocate staging metadata for upcoming loads.""" - - artifacts: list[StagedArtifact] = [] - now = time() - - for request in requests: - ttl = max(request["ttl_seconds"], 0) - cleanup_token = f"{request['correlation_id']}::{request['partition_id']}" - artifacts.append({ - "partition_id": request["partition_id"], - "uri": request["destination_uri"], - "cleanup_token": cleanup_token, - "ttl_seconds": ttl, - "expires_at": now + ttl if ttl else now, - "correlation_id": request["correlation_id"], - }) - if artifacts: - _METRICS.record_partitions(len(artifacts)) - return artifacts - - def cleanup_staging_artifacts(self, artifacts: "list[StagedArtifact]", *, ignore_errors: bool = True) -> None: - """Delete staged artifacts best-effort.""" - - for artifact in artifacts: - backend, path, backend_name = self._backend(artifact["uri"], None) - try: - _delete_backend_sync(backend, path, backend_name=backend_name) - except Exception: - if not ignore_errors: - raise - def _write_bytes( self, payload: bytes, @@ -552,25 +524,6 @@ async def write_arrow( payload, destination, rows=int(table.num_rows), format_label=format_choice, storage_options=resolved_options ) - async def cleanup_staging_artifacts(self, artifacts: "list[StagedArtifact]", *, ignore_errors: bool = True) -> None: - for artifact in artifacts: - backend, path, backend_name = self._backend(artifact["uri"], None) - if supports_async_delete(backend): - try: - await execute_async_storage_operation( - partial(backend.delete_async, path), backend=backend_name, operation="delete", path=path - ) - except Exception: - if not ignore_errors: - raise - continue - - try: - await async_(_delete_backend_sync)(backend=backend, path=path, backend_name=backend_name) - except Exception: - if not ignore_errors: - raise - async def _write_bytes_async( self, payload: bytes, diff --git a/tests/unit/adapters/test_adbc/test_arrow_streaming.py b/tests/unit/adapters/test_adbc/test_arrow_streaming.py index 31cf31343..ce51d16fe 100644 --- a/tests/unit/adapters/test_adbc/test_arrow_streaming.py +++ b/tests/unit/adapters/test_adbc/test_arrow_streaming.py @@ -15,8 +15,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py index 6e02d4687..6b94fce26 100644 --- a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py @@ -14,8 +14,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_aiosqlite/test_load_from_arrow_transaction.py b/tests/unit/adapters/test_aiosqlite/test_load_from_arrow_transaction.py index 1044c3328..b837a6bab 100644 --- a/tests/unit/adapters/test_aiosqlite/test_load_from_arrow_transaction.py +++ b/tests/unit/adapters/test_aiosqlite/test_load_from_arrow_transaction.py @@ -15,8 +15,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_arrow_odbc/test_load_from_arrow.py b/tests/unit/adapters/test_arrow_odbc/test_load_from_arrow.py index 1d9255a9a..9cb4a0a8c 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_load_from_arrow.py +++ b/tests/unit/adapters/test_arrow_odbc/test_load_from_arrow.py @@ -11,8 +11,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py index 0ea841720..f6ae77385 100644 --- a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py +++ b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py @@ -13,8 +13,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_bigquery/test_job_controls.py b/tests/unit/adapters/test_bigquery/test_job_controls.py index 272d2c776..3a00afb58 100644 --- a/tests/unit/adapters/test_bigquery/test_job_controls.py +++ b/tests/unit/adapters/test_bigquery/test_job_controls.py @@ -19,8 +19,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": ["fixed"], } diff --git a/tests/unit/adapters/test_bigquery/test_storage_write_api.py b/tests/unit/adapters/test_bigquery/test_storage_write_api.py index f7c4a64f5..8c4f13f31 100644 --- a/tests/unit/adapters/test_bigquery/test_storage_write_api.py +++ b/tests/unit/adapters/test_bigquery/test_storage_write_api.py @@ -15,8 +15,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": ["fixed"], } diff --git a/tests/unit/adapters/test_duckdb/test_type_converter.py b/tests/unit/adapters/test_duckdb/test_type_converter.py index 8ed4601e6..3e181fe7c 100644 --- a/tests/unit/adapters/test_duckdb/test_type_converter.py +++ b/tests/unit/adapters/test_duckdb/test_type_converter.py @@ -10,8 +10,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py index 86fafa048..c309205ae 100644 --- a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py @@ -11,8 +11,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py b/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py index 446496214..cb2c2b07c 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py +++ b/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py @@ -14,8 +14,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_oracledb/test_direct_path_load.py b/tests/unit/adapters/test_oracledb/test_direct_path_load.py index 72bbdf9b1..96fa68061 100644 --- a/tests/unit/adapters/test_oracledb/test_direct_path_load.py +++ b/tests/unit/adapters/test_oracledb/test_direct_path_load.py @@ -11,8 +11,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py index dbc37afba..a5b80014c 100644 --- a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py @@ -13,8 +13,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": False, "parquet_import_enabled": False, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/adapters/test_spanner/test_batch_write_api.py b/tests/unit/adapters/test_spanner/test_batch_write_api.py index 5f88e638e..8a002e02a 100644 --- a/tests/unit/adapters/test_spanner/test_batch_write_api.py +++ b/tests/unit/adapters/test_spanner/test_batch_write_api.py @@ -14,8 +14,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": ["fixed"], } diff --git a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py index b9871b654..f459d5848 100644 --- a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py +++ b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py @@ -13,8 +13,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": ["fixed"], } diff --git a/tests/unit/adapters/test_sqlite/test_load_from_arrow_transaction.py b/tests/unit/adapters/test_sqlite/test_load_from_arrow_transaction.py index 5987ef423..9d1f5e175 100644 --- a/tests/unit/adapters/test_sqlite/test_load_from_arrow_transaction.py +++ b/tests/unit/adapters/test_sqlite/test_load_from_arrow_transaction.py @@ -15,8 +15,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": [], } diff --git a/tests/unit/config/test_storage_capabilities.py b/tests/unit/config/test_storage_capabilities.py index b6553134b..0f8501237 100644 --- a/tests/unit/config/test_storage_capabilities.py +++ b/tests/unit/config/test_storage_capabilities.py @@ -126,8 +126,6 @@ class _CapabilityConfig(_NoPoolSyncConfigBase): supports_native_arrow_import = True supports_native_parquet_export = False supports_native_parquet_import = False - requires_staging_for_load = True - staging_protocols = ("s3://",) storage_partition_strategies = ("fixed", "rows_per_chunk") default_storage_profile = "local-temp" @@ -148,8 +146,6 @@ class _AsyncCapabilityConfig(_NoPoolAsyncConfigBase): connection_type = object supports_native_arrow_export = True supports_native_arrow_import = True - requires_staging_for_load = True - staging_protocols = ("s3://",) storage_partition_strategies = ("fixed", "rows_per_chunk") async def create_connection(self) -> object: @@ -216,7 +212,6 @@ def test_storage_capabilities_snapshot(monkeypatch): assert capabilities["arrow_export_enabled"] is True assert capabilities["arrow_import_enabled"] is True assert capabilities["parquet_export_enabled"] is False - assert capabilities["requires_staging_for_load"] is True assert capabilities["partition_strategies"] == ["fixed", "rows_per_chunk"] assert capabilities["default_storage_profile"] == "local-temp" diff --git a/tests/unit/storage/test_bridge.py b/tests/unit/storage/test_bridge.py index dcce37e5d..2edd63ae1 100644 --- a/tests/unit/storage/test_bridge.py +++ b/tests/unit/storage/test_bridge.py @@ -29,15 +29,11 @@ from sqlspec.adapters.pymysql import default_statement_config as pymysql_statement_config from sqlspec.adapters.sqlite import SqliteDriver from sqlspec.adapters.sqlite import default_statement_config as sqlite_statement_config +from sqlspec.exceptions import StoragePathTraversalError +from sqlspec.protocols import ObjectStoreProtocol from sqlspec.storage import SyncStoragePipeline, get_storage_bridge_diagnostics, reset_storage_bridge_metrics -from sqlspec.storage.pipeline import ( - AsyncStoragePipeline, - StagedArtifact, - StorageDestination, - _encode_row_payload, - _StoragePipelineBase, -) -from sqlspec.storage.registry import storage_registry +from sqlspec.storage.pipeline import AsyncStoragePipeline, StorageDestination, _encode_row_payload, _StoragePipelineBase +from sqlspec.storage.registry import StorageRegistry, storage_registry from sqlspec.utils.serializers import reset_serializer_cache, serialize_collection CAPABILITIES = { @@ -45,8 +41,6 @@ "arrow_import_enabled": True, "parquet_export_enabled": True, "parquet_import_enabled": True, - "requires_staging_for_load": False, - "staging_protocols": [], "partition_strategies": ["fixed"], } @@ -169,13 +163,10 @@ class _CountingStorageBackend: backend_type = "counting" def __init__(self) -> None: - self.deleted_paths: list[str] = [] - - def delete_sync(self, path: str) -> None: - self.deleted_paths.append(path) + self.written_paths: list[str] = [] - async def delete_async(self, path: str) -> None: - self.deleted_paths.append(path) + async def write_bytes_async(self, path: str, payload: bytes) -> None: + self.written_paths.append(path) class _CountingStorageRegistry: @@ -235,36 +226,20 @@ def test_sync_pipeline_bypasses_resolution_cache_for_storage_options() -> None: ] -async def test_async_pipeline_cleanup_reuses_cached_backend_resolution() -> None: +async def test_async_pipeline_write_reuses_cached_backend_resolution() -> None: registry = _CountingStorageRegistry() pipeline = AsyncStoragePipeline(registry=cast(Any, registry)) - artifacts: list[StagedArtifact] = [ - { - "partition_id": "0", - "uri": "file://tmp/payload.jsonl", - "cleanup_token": "cleanup::0", - "ttl_seconds": 0, - "expires_at": 0.0, - "correlation_id": "cleanup", - }, - { - "partition_id": "1", - "uri": "file://tmp/payload.jsonl", - "cleanup_token": "cleanup::1", - "ttl_seconds": 0, - "expires_at": 0.0, - "correlation_id": "cleanup", - }, - ] + table = pa.table({"id": [1]}) - await pipeline.cleanup_staging_artifacts(artifacts) + await pipeline.write_arrow(table, "file://tmp/payload.parquet") + await pipeline.write_arrow(table, "file://tmp/payload.parquet") - assert registry.calls == [("file://tmp/payload.jsonl", {})] - assert registry.backend.deleted_paths == ["tmp/payload.jsonl", "tmp/payload.jsonl"] + assert registry.calls == [("file://tmp/payload.parquet", {})] + assert registry.backend.written_paths == ["tmp/payload.parquet", "tmp/payload.parquet"] pipeline.clear_cache() - await pipeline.cleanup_staging_artifacts(artifacts[:1]) - assert registry.calls == [("file://tmp/payload.jsonl", {}), ("file://tmp/payload.jsonl", {})] + await pipeline.write_arrow(table, "file://tmp/payload.parquet") + assert registry.calls == [("file://tmp/payload.parquet", {}), ("file://tmp/payload.parquet", {})] async def test_asyncpg_load_from_storage(monkeypatch: pytest.MonkeyPatch) -> None: @@ -706,6 +681,127 @@ def test_storage_bridge_diagnostics_include_serializer_metrics() -> None: assert "serializer.size" in diagnostics +class _ResolutionCountingRegistry(StorageRegistry): + def __init__(self) -> None: + super().__init__() + self.calls: list[tuple[str, str | None, dict[str, Any]]] = [] + + def get(self, uri_or_alias: str | Path, *, backend: str | None = None, **kwargs: Any) -> ObjectStoreProtocol: + self.calls.append((str(uri_or_alias), backend, kwargs)) + return super().get(uri_or_alias, backend=backend, **kwargs) + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +@pytest.mark.parametrize( + ("destination", "options", "protocol"), + [ + ( + "s3://example-bucket/prefix/file.parquet", + {"backend": "obstore", "skip_signature": True, "region": "us-east-1"}, + "s3", + ), + ("gs://example-bucket/prefix/file.parquet", {"backend": "obstore", "skip_signature": True}, "gs"), + ("s3://example-bucket/prefix/file.parquet", {"backend": "fsspec", "anon": True}, "s3"), + ], +) +def test_resolve_destination_remote_uri( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], + destination: str, + options: dict[str, Any], + protocol: str, +) -> None: + pipeline = pipeline_type(registry=StorageRegistry()) + + target = pipeline.resolve_destination(destination, options) + + assert target.uri == destination + assert target.protocol == protocol + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +@pytest.mark.parametrize("backend", ["obstore", "fsspec"]) +def test_resolve_destination_remote_alias_prefix_and_options( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], backend: str +) -> None: + registry = StorageRegistry() + options = {"skip_signature": True, "region": "us-east-1"} if backend == "obstore" else {"anon": True} + registry.register_alias("assets", "s3://example-bucket/prefix", backend=backend, base_path="default", **options) + pipeline = pipeline_type(registry=registry) + + target = pipeline.resolve_destination("alias://assets/sub/file.parquet", {"base_path": "override"}) + + assert target.uri == "s3://example-bucket/prefix/override/sub/file.parquet" + assert target.protocol == "s3" + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +@pytest.mark.parametrize("path_kind", ["absolute", "relative", "uri"]) +def test_resolve_destination_local_path( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], + path_kind: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + path = tmp_path / "folder" / "data file.parquet" + destination: str | Path = path + if path_kind == "relative": + destination = path.relative_to(tmp_path) + elif path_kind == "uri": + destination = path.as_uri() + pipeline = pipeline_type(registry=StorageRegistry()) + + target = pipeline.resolve_destination(destination, {"backend": "local"}) + + assert target.uri == str(path) + assert target.protocol == "file" + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +@pytest.mark.parametrize("backend", ["local", "obstore", "fsspec"]) +def test_resolve_destination_local_alias_prefix( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], backend: str, tmp_path: Path +) -> None: + registry = StorageRegistry() + registry.register_alias("assets", tmp_path.as_uri(), backend=backend, base_path="prefix") + pipeline = pipeline_type(registry=registry) + + target = pipeline.resolve_destination("alias://assets/sub/file.parquet") + + assert target.uri == str(tmp_path / "prefix" / "sub" / "file.parquet") + assert target.protocol == "file" + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +@pytest.mark.parametrize("backend", ["local", "obstore"]) +def test_resolve_destination_local_options_preserve_root_guard( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], backend: str, tmp_path: Path +) -> None: + pipeline = pipeline_type(registry=StorageRegistry()) + + with pytest.raises(StoragePathTraversalError): + pipeline.resolve_destination(tmp_path / "file.parquet", {"backend": backend, "base_path": "nested"}) + + +@pytest.mark.parametrize("pipeline_type", [SyncStoragePipeline, AsyncStoragePipeline]) +def test_resolve_destination_cache_and_explicit_options( + pipeline_type: type[SyncStoragePipeline | AsyncStoragePipeline], tmp_path: Path +) -> None: + registry = _ResolutionCountingRegistry() + pipeline = pipeline_type(registry=registry, storage_options={"write_options": {"delimiter": "|"}}) + destination = tmp_path / "file.parquet" + + assert pipeline.resolve_destination(destination) == pipeline.resolve_destination(destination, {}) + assert registry.calls == [(str(destination), None, {})] + pipeline.clear_cache() + assert pipeline.resolve_destination(destination).uri == str(destination) + assert len(registry.calls) == 2 + + for _ in range(2): + assert pipeline.resolve_destination(destination, {"backend": "local"}).uri == str(destination) + assert registry.calls[2:] == [(str(destination), "local", {}), (str(destination), "local", {})] + + class _CsvTestBackend: """Minimal backend for CSV pipeline tests."""