Praneshrajan15's picture
Deploy DataForge playground API
13fe504 verified
Raw
History Blame Contribute Delete
3.2 kB
"""Table-store URI parsing and adapter selection."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlparse
from dataforge.stores.base import TableStore, TableStoreError
from dataforge.stores.cloud import CloudWarehouseStore
from dataforge.stores.duckdb import DuckDBStore
_CLOUD_BACKENDS = {"snowflake", "bigquery", "databricks", "databricks_delta"}
@dataclass(frozen=True)
class TableStoreSpec:
"""Parsed table-store URI."""
backend: str
target: str
relation: str
database_path: Path | None
row_identity_columns: tuple[str, ...]
def is_table_store_uri(raw: str) -> bool:
"""Return whether a CLI target string names a DataForge table store."""
return raw.startswith("warehouse://")
def parse_table_store_uri(raw: str, *, row_ids: tuple[str, ...] = ()) -> TableStoreSpec:
"""Parse a ``warehouse://`` URI into an adapter spec.
Supported local form:
``warehouse://duckdb?database=/tmp/dev.duckdb&relation=main.model&row_id=id``.
"""
parsed = urlparse(raw)
if parsed.scheme != "warehouse":
raise TableStoreError("Table-store URIs must use the warehouse:// scheme.")
backend = (parsed.netloc or parsed.path.strip("/").split("/", 1)[0]).lower()
if not backend:
raise TableStoreError("Warehouse URI must include a backend name.")
query = {key: values[-1] for key, values in parse_qs(parsed.query).items() if values}
relation = unquote(query.get("relation", ""))
if not relation:
raise TableStoreError("Warehouse URI must include relation=<schema.table>.")
database = query.get("database") or query.get("path")
row_id_query = query.get("row_id") or query.get("key")
resolved_row_ids = row_ids
if row_id_query:
resolved_row_ids = tuple(part.strip() for part in row_id_query.split(",") if part.strip())
return TableStoreSpec(
backend=backend,
target=raw,
relation=relation,
database_path=Path(unquote(database)).expanduser() if database else None,
row_identity_columns=resolved_row_ids,
)
def store_from_uri(raw: str, *, row_ids: tuple[str, ...] = ()) -> TableStore:
"""Create a table-store adapter from a CLI URI."""
spec = parse_table_store_uri(raw, row_ids=row_ids)
if spec.backend == "duckdb":
if spec.database_path is None:
raise TableStoreError("DuckDB warehouse URI requires database=<path>.")
return DuckDBStore(
database_path=spec.database_path,
relation=spec.relation,
row_identity_columns=spec.row_identity_columns,
target=spec.target,
)
if spec.backend in _CLOUD_BACKENDS:
backend = "databricks" if spec.backend == "databricks_delta" else spec.backend
return CloudWarehouseStore(
backend=backend,
target=spec.target,
relation=spec.relation,
row_identity_columns=spec.row_identity_columns,
)
raise TableStoreError(f"Unsupported warehouse backend: {spec.backend}")