dataface 0.1.2__py3-none-any.whl

dataface/core/execute/adapters/dbt_adapter.py

@@ -0,0 +1,283 @@
+"""dbt adapter integration for SQL query execution.
+
+Stage: EXECUTE (inside RENDER stage)
+Purpose: Execute SQL queries via dbt's adapter system.
+
+This adapter leverages dbt's adapter API to execute SQL queries,
+supporting dbt-specific features like ref() and source() resolution.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import yaml
+
+if TYPE_CHECKING:
+    from dataface.core.compile.models.source import SourceConfig
+
+from dataface.core.compile.jinja import resolve_jinja_template
+from dataface.core.compile.models.face.compiled import VariableValues
+from dataface.core.compile.models.query.compiled import (
+    AnyQuery,
+    is_sql_query,
+)
+from dataface.core.execute.adapters.base import (
+    BaseAdapter,
+    QueryParams,
+    QueryResult,
+    handle_adapter_error,
+)
+from dataface.core.execute.dbt_jinja import has_dbt_jinja
+from dataface.core.execute.sql_guard import validate_select_only
+
+
+def _read_profiles_yml(dbt_project_path: Path) -> dict[str, Any]:
+    """Read and return the profiles.yml nearest to the project.
+
+    Checks project-local profiles.yml first, then ~/.dbt/profiles.yml.
+    Raises FileNotFoundError if neither exists.
+    """
+    for candidate in (
+        dbt_project_path / "profiles.yml",
+        Path.home() / ".dbt" / "profiles.yml",
+    ):
+        if candidate.exists():
+            with open(candidate) as f:
+                data = yaml.safe_load(f)
+            if not isinstance(data, dict):
+                raise ValueError(
+                    f"profiles.yml at {candidate} did not parse to a mapping. "
+                    f"Got: {type(data).__name__}"
+                )
+            return data
+    raise FileNotFoundError(f"No profiles.yml found in {dbt_project_path} or ~/.dbt/")
+
+
+def _read_target_dict(
+    dbt_project_path: Path, profile_name: str, target_name: str
+) -> dict[str, Any]:
+    """Return the target dict for the given profile+target from profiles.yml.
+
+    Relative `path:` entries (DuckDB) are resolved against `dbt_project_path`
+    so the adapter opens the file dbt would have opened, regardless of CWD.
+    """
+    profiles = _read_profiles_yml(dbt_project_path)
+    profile = profiles.get(profile_name)
+    if profile is None:
+        available = [k for k in profiles if k != "config"]
+        raise ValueError(
+            f"Profile '{profile_name}' not found in profiles.yml "
+            f"(checked {dbt_project_path}). "
+            f"Available profiles: {available}"
+        )
+    outputs = profile.get("outputs", {})
+    target = outputs.get(target_name)
+    if target is None:
+        raise ValueError(
+            f"Target '{target_name}' not found in profile '{profile_name}'. "
+            f"Available targets: {list(outputs.keys())}"
+        )
+    target = dict(target)
+    raw_path = target.get("path")
+    if isinstance(raw_path, str) and raw_path and raw_path != ":memory:":
+        path_obj = Path(raw_path)
+        if not path_obj.is_absolute():
+            target["path"] = str((dbt_project_path / path_obj).resolve())
+    return target
+
+
+class DbtAdapter(BaseAdapter):
+    """Adapter for executing SQL queries via dbt's adapter system.
+
+    Supported query types: sql
+
+    Leverages dbt's adapter API to execute SQL queries with support for
+    dbt-specific features like ref(), source(), etc.
+
+    This adapter requires dbt-adapters and a warehouse-specific dbt package
+    (e.g. dbt-duckdb, dbt-postgres) — not dbt-core. A valid dbt project
+    with profiles.yml configuration is also required.
+
+    Example:
+        >>> adapter = DbtAdapter(dbt_project_path=Path("./my_dbt_project"))
+        >>> query = SqlQuery(sql="SELECT * FROM {{ ref('customers') }}")
+        >>> result = adapter.execute(query)
+    """
+
+    def __init__(
+        self,
+        dbt_project_path: Path | None = None,
+        profile_name: str | None = None,
+        target_name: str | None = None,
+    ):
+        """Initialize dbt adapter.
+
+        Args:
+            dbt_project_path: Path to dbt project (default: current directory)
+            profile_name: dbt profile name (default: from dbt_project.yml)
+            target_name: dbt target name (default: DBT_TARGET env var, then 'dev')
+        """
+        import os
+
+        self.dbt_project_path = (
+            Path(dbt_project_path).resolve()
+            if dbt_project_path
+            else Path.cwd().resolve()
+        )
+        self.profile_name = profile_name
+        self.target_name = target_name or os.environ.get("DBT_TARGET") or "dev"
+        self._adapter: Any = None
+        # _dialect is populated by _get_dbt_adapter() — which always runs before
+        # any validate_select_only() call site — and `build_adapter` raises if
+        # the target dict lacks 'type', so the empty-string default is never
+        # observed by the guard in practice.
+        self._dialect: str = ""
+        self._manifest: dict[str, Any] | None = None
+        self._manifest_loaded: bool = False
+
+    @property
+    def supported_types(self) -> set[str]:
+        """Return supported query types."""
+        return {"sql"}
+
+    def _can_execute(self, query: AnyQuery) -> bool:
+        """Return True if this adapter is certain it can handle the query.
+
+        Pure predicate — never constructs the adapter, never reads profiles.yml.
+        Misconfiguration errors surface in _execute() as QueryResult.error,
+        not here in routing.
+
+        Routing contract:
+          - Explicit source (dict or named string) → SqlAdapter owns it
+            (it has the inline-source resolver and the named-source registry).
+          - No source AND SQL has `{{ ref() }}` / `{{ source() }}` jinja →
+            DbtAdapter (needs manifest resolution).
+          - No source, no jinja → SqlAdapter (default DuckDB).
+        """
+        if not is_sql_query(query):
+            return False
+        if query.source is not None:
+            return False
+        return has_dbt_jinja(query.sql)
+
+    def _execute(
+        self,
+        query: AnyQuery,
+        variables: VariableValues | None = None,
+        params: QueryParams = None,
+        source_config: SourceConfig | None = None,
+    ) -> QueryResult:
+        """Execute a SQL query via dbt adapter.
+
+        Args:
+            query: AnyQuery object (SqlQuery expected)
+            variables: Variable values for Jinja resolution
+            params: Accepted for interface compatibility; not used by dbt adapter.
+
+        Returns:
+            QueryResult with data or error
+        """
+        if not is_sql_query(query):
+            return QueryResult(
+                data=[],
+                error=f"Expected SQL query, got {query.query_type}",
+            )
+
+        try:
+            if self._adapter is None:
+                self._get_dbt_adapter()
+
+            resolved_sql = self._resolve_dbt_sql(query.sql, variables)
+            validate_select_only(resolved_sql, dialect=self._dialect)
+
+            with self._adapter.connection_named("dataface_query"):
+                _, table = self._adapter.execute(
+                    resolved_sql, auto_begin=True, fetch=True
+                )
+
+            columns = list(table.column_names)
+            rows = list(table.rows)
+            if query.limit and query.limit > 0:
+                rows = rows[: query.limit]
+            data = [dict(zip(columns, row, strict=False)) for row in rows]
+            return QueryResult(data=data, columns=columns)
+
+        except Exception as e:  # noqa: BLE001
+            return handle_adapter_error("Query execution", e)
+
+    def _get_dbt_adapter(self) -> Any:
+        """Get dbt adapter instance (lazy-loaded).
+
+        Always loads the manifest first (doesn't require dbt).
+        Then delegates adapter construction to build_adapter() in the factory —
+        single point of truth for source_config → dbt adapter.
+        """
+        if self._adapter is not None:
+            return self._adapter
+
+        # Always try to load manifest first — this doesn't require dbt
+        self._load_manifest()
+
+        if not self.profile_name:
+            project_config_path = self.dbt_project_path / "dbt_project.yml"
+            if project_config_path.exists():
+                with open(project_config_path) as f:
+                    project_dict = yaml.safe_load(f)
+                profile_from_project = project_dict.get("profile")
+                if not profile_from_project:
+                    raise ValueError(
+                        f"dbt_project.yml at {project_config_path} has no 'profile:' key"
+                    )
+                self.profile_name = profile_from_project
+            else:
+                raise FileNotFoundError(
+                    f"No dbt_project.yml found at {project_config_path} and no "
+                    f"profile_name was provided to DbtAdapter"
+                )
+
+        target_dict = _read_target_dict(
+            self.dbt_project_path, self.profile_name, self.target_name
+        )
+        from dataface.core.execute.adapters.dbt_adapter_factory import build_adapter
+
+        # build_adapter raises ValueError if target_dict has no 'type'; reordering
+        # so the build happens first means `target_dict["type"]` is guaranteed
+        # populated when we read it.
+        self._adapter = build_adapter(target_dict)
+        self._dialect = target_dict["type"]
+        return self._adapter
+
+    def _load_manifest(self) -> None:
+        """Load dbt manifest from target/ or snapshot file.
+
+        Tries target/manifest.json first (for development), then falls back
+        to manifest.snapshot.json (for production deployments).
+        """
+        if self._manifest_loaded:
+            return  # Already attempted (result may be None if no manifest exists)
+
+        from dataface.core.execute.adapters.dbt_utils import load_dbt_manifest
+
+        self._manifest = load_dbt_manifest(self.dbt_project_path)
+        self._manifest_loaded = True
+
+    def _resolve_dbt_sql(
+        self, sql: str, variables: dict[str, Any] | None = None
+    ) -> str:
+        """Resolve SQL with dbt-specific Jinja functions."""
+        from dataface.core.execute.adapters.dbt_utils import resolve_dbt_refs
+
+        if has_dbt_jinja(sql) and not self._manifest:
+            kind = "ref()" if "ref(" in sql else "source()"
+            raise ValueError(
+                f"SQL contains {kind} but no dbt manifest found. "
+                f"Checked paths: {self.dbt_project_path / 'target' / 'manifest.json'}, "
+                f"{self.dbt_project_path / 'manifest.snapshot.json'}"
+            )
+
+        if self._manifest:
+            sql = resolve_dbt_refs(sql, self._manifest)
+
+        return resolve_jinja_template(sql, variables)

dataface/core/execute/adapters/dbt_adapter_factory.py

@@ -0,0 +1,212 @@
+"""Single point of truth for source_config → dbt adapter instantiation.
+
+Every consumer that needs a live dbt adapter calls build_adapter().
+The adapter lives as long as the caller holds a reference; GC reclaims
+it automatically via weakref.finalize, which deletes the temp target dir.
+"""
+
+import importlib
+import logging
+import multiprocessing
+import shutil
+import tempfile
+import weakref
+from types import SimpleNamespace
+from typing import Any, cast
+
+from dataface.core.compile.models.source import infer_bq_method
+
+logger = logging.getLogger(__name__)
+
+# Mapping: source_config type → (module_path, AdapterClass, CredentialsClass)
+# Add entries here when new warehouse adapters are needed.
+# This is the single source of truth — dbt_adapter.py imports from here.
+_ADAPTER_TYPE_MAP: dict[str, tuple[str, str, str]] = {
+    "duckdb": ("dbt.adapters.duckdb", "DuckDBAdapter", "DuckDBCredentials"),
+    "postgres": ("dbt.adapters.postgres", "PostgresAdapter", "PostgresCredentials"),
+    "postgresql": ("dbt.adapters.postgres", "PostgresAdapter", "PostgresCredentials"),
+    "redshift": ("dbt.adapters.redshift", "RedshiftAdapter", "RedshiftCredentials"),
+    "snowflake": ("dbt.adapters.snowflake", "SnowflakeAdapter", "SnowflakeCredentials"),
+    "bigquery": ("dbt.adapters.bigquery", "BigQueryAdapter", "BigQueryCredentials"),
+    "spark": ("dbt.adapters.spark", "SparkAdapter", "SparkCredentials"),
+    "databricks": (
+        "dbt.adapters.databricks",
+        "DatabricksAdapter",
+        "DatabricksCredentials",
+    ),
+}
+
+
+_DUCKDB_INITIALIZE_DB_PATCHED = False
+_ORIGINAL_DUCKDB_INITIALIZE_DB: Any = None
+
+
+def _patched_duckdb_initialize_db(
+    cls: type[Any],
+    creds: Any,
+    plugins: dict[str, Any] | None = None,
+) -> Any:
+    if getattr(creds, "_dft_read_only", False) and creds.path != ":memory:":
+        import duckdb
+
+        config = dict(creds.config_options or {})
+        config.setdefault("enable_external_access", False)
+        return duckdb.connect(creds.path, read_only=True, config=config)
+    return _ORIGINAL_DUCKDB_INITIALIZE_DB(cls, creds, plugins)
+
+
+def _ensure_duckdb_readonly_initialize_db_patch() -> None:
+    """Patch dbt-duckdb's hardcoded read_only=False for schema/introspection paths."""
+    global _DUCKDB_INITIALIZE_DB_PATCHED, _ORIGINAL_DUCKDB_INITIALIZE_DB
+    if _DUCKDB_INITIALIZE_DB_PATCHED:
+        return
+
+    from dbt.adapters.duckdb import environments as duckdb_envs
+
+    _orig_cm = cast(Any, duckdb_envs.Environment.initialize_db)
+    _ORIGINAL_DUCKDB_INITIALIZE_DB = _orig_cm.__func__
+    duckdb_envs.Environment.initialize_db = classmethod(  # type: ignore[method-assign,assignment]
+        _patched_duckdb_initialize_db
+    )
+    _DUCKDB_INITIALIZE_DB_PATCHED = True
+
+
+def build_adapter(source_config: dict[str, Any], *, read_only: bool = False) -> Any:
+    """Construct a fresh dbt adapter instance.
+
+    Args:
+        source_config: Dict with at minimum a 'type' key matching a known
+            warehouse adapter. Other keys are translated via the credentials
+            class's alias map (e.g. BigQuery project→database) and deserialized
+            through from_dict, so profile-level extras like threads are dropped.
+        read_only: When True, open DuckDB file sources read-only so pure-metadata
+            verbs coexist with other readers (e.g. a running ``dft serve``).
+            In-memory DuckDB stays read-write. Other warehouses ignore this
+            flag until a driver-specific read-only hook exists.
+
+    Returns:
+        A dbt adapter instance. Callers use adapter.connection_named("name")
+        as a context manager before calling adapter.execute(...). The temp
+        target directory is deleted automatically when the adapter is GC'd.
+
+    Raises:
+        ValueError: 'type' key is missing or names an unsupported adapter.
+        ImportError: The dbt adapter package for this warehouse is not installed.
+    """
+    adapter_type = source_config.get("type")
+    if not adapter_type:
+        raise ValueError(
+            "source_config must have a 'type' field. "
+            f"Got: {list(source_config.keys())}"
+        )
+
+    adapter_type_lower = str(adapter_type).lower()
+    entry = _ADAPTER_TYPE_MAP.get(adapter_type_lower)
+    if entry is None:
+        raise ValueError(
+            f"Unsupported adapter type '{adapter_type}'. "
+            f"Supported: {', '.join(sorted(_ADAPTER_TYPE_MAP))}. "
+            f"For other warehouses, install the relevant dbt-<warehouse> package "
+            f"and add an entry to _ADAPTER_TYPE_MAP in dbt_adapter_factory.py."
+        )
+
+    module_path, adapter_name, creds_name = entry
+
+    try:
+        mod = importlib.import_module(module_path)
+    except ImportError as e:
+        raise ImportError(
+            f"dbt-{adapter_type_lower} is not installed. "
+            f"Install it with: pip install dbt-{adapter_type_lower}"
+        ) from e
+
+    adapter_cls = getattr(mod, adapter_name)
+    creds_cls = getattr(mod, creds_name)
+
+    # Strip 'type', translate profile-style aliases (e.g. BigQuery project→database),
+    # then use from_dict so extras like threads are silently dropped.
+    creds_kwargs = {k: v for k, v in source_config.items() if k != "type"}
+    if adapter_type_lower == "bigquery" and "method" not in creds_kwargs:
+        creds_kwargs["method"] = infer_bq_method(
+            creds_kwargs.get("keyfile"), creds_kwargs.get("keyfile_json")
+        )
+    creds_kwargs = creds_cls.translate_aliases(creds_kwargs)
+    creds = creds_cls.from_dict(creds_kwargs)
+    if read_only and adapter_type_lower == "duckdb":
+        _ensure_duckdb_readonly_initialize_db_patch()
+        creds._dft_read_only = True
+
+    # Build a minimal duck-typed config satisfying AdapterRequiredConfig.
+    # No dbt_project.yml, no manifest — but enough for dbt's macro context to
+    # construct itself, since list_schemas / list_relations / get_columns route
+    # through execute_macro -> generate_runtime_macro_context, which reads
+    # quoting / dependencies / args / vars / load_dependencies / project_root /
+    # get_macro_search_order off the config.
+    profile_name = f"dft_{adapter_type_lower}"
+    target_path = tempfile.mkdtemp(prefix="dft-target-")
+    cfg = SimpleNamespace(
+        credentials=creds,
+        profile_name=profile_name,
+        target_name="dev",
+        threads=1,
+        project_name=profile_name,
+        query_comment=SimpleNamespace(comment="", append=False, job_label=False),
+        cli_vars={},
+        target_path=target_path,
+        project_root=target_path,
+        log_cache_events=False,
+        quoting={"database": True, "schema": True, "identifier": True},
+        dependencies={},
+        args=SimpleNamespace(vars={}),
+        vars=_NoVars(),
+        flags={},
+        load_dependencies=dict,
+        get_macro_search_order=lambda namespace: None,
+        to_target_dict=lambda c=source_config: dict(c),
+    )
+
+    adapter = adapter_cls(cfg, multiprocessing.get_context("spawn"))
+    # Register cleanup to fire when the adapter is GC'd.
+    # shutil.rmtree is called with ignore_errors=True so a missing dir is skipped.
+    weakref.finalize(adapter, shutil.rmtree, target_path, True)
+    _bootstrap_macros(adapter, adapter_type_lower)
+    return adapter
+
+
+def _bootstrap_macros(adapter: Any, adapter_type: str) -> None:
+    """Wire dbt-bundled macros onto the adapter so introspection works.
+
+    Without this, every BaseAdapter introspection method (list_schemas,
+    list_relations, get_columns_in_relation, calculate_freshness_from_metadata,
+    get_partitions_metadata) raises "Macro resolver was None" because
+    execute_macro() needs a populated _macro_resolver.
+
+    Also registers a weakref.proxy of the adapter in dbt's FACTORY so the
+    macro context can resolve `adapter` via get_adapter(config) — the macro
+    context is constructed in dbt-core's MacroContext.__init__, which calls
+    get_adapter(self.config). A weakref keeps GC working: the FACTORY entry
+    doesn't pin the adapter, so weakref.finalize cleanup still fires when the
+    caller drops their reference. Last build_adapter() call wins per type.
+    """
+    from dbt.adapters.factory import FACTORY
+    from dbt.context.providers import generate_runtime_macro_context
+
+    from dataface.core.execute.adapters.dbt_macro_loader import load_macro_manifest
+
+    with FACTORY.lock:
+        # weakref.proxy: FACTORY.adapters[type] doesn't pin the adapter, so
+        # weakref.finalize cleanup still fires when the caller drops their ref.
+        FACTORY.adapters[adapter_type] = weakref.proxy(adapter)  # type: ignore[assignment]
+    adapter.set_macro_resolver(load_macro_manifest(adapter_type))
+    adapter.set_macro_context_generator(generate_runtime_macro_context)
+
+
+class _NoVars:
+    """Stub for `RuntimeConfig.vars` accessed during macro context build.
+
+    Macros we care about (list_schemas, list_relations, get_columns_in_relation)
+    don't reference any project vars, so an empty `vars_for` is the right shape.
+    """
+
+    def vars_for(self, lookup: Any, adapter_type: str) -> dict:
+        return {}

dataface/core/execute/adapters/dbt_macro_loader.py

@@ -0,0 +1,95 @@
+"""Bootstrap a dbt MacroManifest from pip-installed dbt packages.
+
+dbt-core's BaseAdapter introspection API (`list_schemas`, `list_relations`,
+`get_columns_in_relation`, ...) routes through `execute_macro()`, which needs
+a populated `_macro_resolver`. Without it, every adapter built by
+`build_adapter()` raises "no macro registered" the moment we try to introspect.
+
+This module loads the bundled macros — dbt-core's global project plus the
+active warehouse adapter's package — into a fresh `MacroManifest` that can be
+attached to an adapter via `set_macro_resolver()`. No user files required,
+no `dbt_project.yml`, no `target/manifest.json`. Just pip-installed packages.
+
+Coupled to dbt-internal API (`MacroParser`, `MacroManifest`,
+`load_source_file`, `FACTORY.packages`); the smoke tests are the canary for
+upgrades.
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import cache
+from pathlib import Path
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dbt.contracts.graph.manifest import MacroManifest
+
+logger = logging.getLogger(__name__)
+
+
+class DbtMacroLoaderError(RuntimeError):
+    """Raised when bootstrapping the dbt MacroManifest fails."""
+
+
+def _resolve_macro_packages(adapter_type: str) -> list[tuple[str, Path]]:
+    """Return [(package_name, package_root), ...] for warehouse + global macros.
+
+    Uses dbt's adapter FACTORY to resolve include paths the same way dbt-core
+    does — by loading the plugin and reading the registered include directory.
+    Warehouse package comes first so its macros parse before the global ones,
+    matching dbt-core's lookup precedence (warehouse-specific overrides win).
+    """
+    from dbt.adapters.factory import FACTORY
+
+    FACTORY.load_plugin(adapter_type)
+    plugin = FACTORY.get_plugin_by_name(adapter_type)
+    return [
+        (plugin.project_name, FACTORY.packages[plugin.project_name]),
+        ("dbt", FACTORY.packages["dbt"]),
+    ]
+
+
+@cache
+def load_macro_manifest(adapter_type: str) -> MacroManifest:
+    """Build a MacroManifest containing dbt-core + warehouse macros.
+
+    Cached per adapter_type for the process lifetime — the bundled macros
+    don't change between calls within a single session.
+    """
+    from dbt.contracts.files import ParseFileType
+    from dbt.contracts.graph.manifest import MacroManifest, Manifest
+    from dbt.parser.macros import MacroParser
+    from dbt.parser.read_files import load_source_file
+    from dbt.parser.search import FileBlock
+
+    packages = _resolve_macro_packages(adapter_type)
+    manifest = Manifest()
+    for package_name, package_root in packages:
+        macros_dir = Path(package_root) / "macros"
+        if not macros_dir.is_dir():
+            raise DbtMacroLoaderError(
+                f"dbt {package_name} macros not found at {macros_dir}. "
+                f"This usually means the dbt-core or dbt-{adapter_type} "
+                f"package layout changed; run the smoke test in "
+                f"test_dbt_macro_loader.py to see what dbt-core version "
+                f"this code was last validated against."
+            )
+        # MacroParser only reads project_root / project_name / macro_paths off
+        # the project arg, so a SimpleNamespace duck-type is sufficient.
+        fake_project = SimpleNamespace(
+            project_root=str(package_root),
+            project_name=package_name,
+            macro_paths=["macros"],
+        )
+        parser = MacroParser(fake_project, manifest)  # type: ignore[arg-type]
+        for path in parser.get_paths():
+            source_file = load_source_file(path, ParseFileType.Macro, package_name, {})
+            assert source_file is not None
+            parser.parse_file(FileBlock(source_file))
+
+    logger.debug(
+        "Loaded %d dbt macros for adapter_type=%s", len(manifest.macros), adapter_type
+    )
+    return MacroManifest(manifest.macros)