dbt-ml 0.1.0__py3-none-any.whl

dbt_ml/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

dbt_ml/adapters/__init__.py

@@ -0,0 +1,12 @@
+from . import duckdb  # noqa: F401  # side-effect: registers DuckDBAdapter
+from .base import AdapterError, WarehouseAdapter
+from .registry import UnknownAdapterError, create_adapter, list_adapter_types, register
+
+__all__ = [
+    "AdapterError",
+    "UnknownAdapterError",
+    "WarehouseAdapter",
+    "create_adapter",
+    "list_adapter_types",
+    "register",
+]

dbt_ml/adapters/base.py

@@ -0,0 +1,142 @@
+"""Warehouse adapter base class.
+
+Each adapter wraps a warehouse-specific connection and exposes a uniform
+interface for the runner: connect/close, schema management, materialization
+(full + incremental), querying, and incremental-state CRUD. The point is
+that runner.py / manifest.py / dbt_export.py / cli.py never speak DuckDB
+SQL directly — they call adapter methods.
+
+Today: DuckDB. v0.2.2: LanceDB. Beyond: Postgres / Snowflake / BigQuery /
+Databricks / Redshift, matching the dbt-core set.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Self
+
+import polars as pl
+
+from ..config.profile import WarehouseConfig
+
+
+class AdapterError(Exception):
+    pass
+
+
+class WarehouseAdapter(ABC):
+    """Lifecycle-managed warehouse driver."""
+
+    def __init__(self, config: WarehouseConfig, *, project_dir: Path | None = None) -> None:
+        self.config = config
+        self.project_dir = project_dir
+
+    # ─── classification ────────────────────────────────────────────────────
+
+    @classmethod
+    @abstractmethod
+    def adapter_type(cls) -> str:
+        """Short name used in profiles.yml `warehouse.type`."""
+
+    # ─── lifecycle ────────────────────────────────────────────────────────
+
+    def __enter__(self) -> Self:
+        self._connect()
+        self._ensure_schema()
+        self._ensure_state_table()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self._close()
+
+    @abstractmethod
+    def _connect(self) -> None: ...
+
+    @abstractmethod
+    def _close(self) -> None: ...
+
+    @abstractmethod
+    def _ensure_schema(self) -> None: ...
+
+    @abstractmethod
+    def _ensure_state_table(self) -> None: ...
+
+    # ─── identity / SQL references ────────────────────────────────────────
+
+    @property
+    @abstractmethod
+    def catalog(self) -> str:
+        """Catalog name used in SQL references and emitted dbt sources."""
+
+    @property
+    def schema(self) -> str:
+        return self.config.schema_name
+
+    @property
+    @abstractmethod
+    def schema_ref(self) -> str:
+        """Quoted, fully-qualified schema reference for use in SQL."""
+
+    def table_ref(self, table: str) -> str:
+        return f"{self.schema_ref}.{table}"
+
+    # ─── materialization ──────────────────────────────────────────────────
+
+    @abstractmethod
+    def materialize_full(self, table: str, df: pl.DataFrame) -> int:
+        """Replace `table` with `df`. Returns row count written."""
+
+    @abstractmethod
+    def materialize_incremental(
+        self, table: str, df: pl.DataFrame, *, key_col: str
+    ) -> int:
+        """Upsert rows in `df` into `table`, keyed on `key_col`. Returns rows written."""
+
+    @abstractmethod
+    def drop_table(self, table: str) -> None: ...
+
+    # ─── querying ─────────────────────────────────────────────────────────
+
+    @abstractmethod
+    def execute(self, sql: str, params: list[Any] | None = None) -> Any: ...
+
+    @abstractmethod
+    def query_df(self, sql: str) -> pl.DataFrame: ...
+
+    @abstractmethod
+    def scalar(self, sql: str, params: list[Any] | None = None) -> Any:
+        """First column of first row, or None."""
+
+    @abstractmethod
+    def rows(self, sql: str, params: list[Any] | None = None) -> list[tuple]: ...
+
+    @abstractmethod
+    def list_tables(self) -> list[str]: ...
+
+    # ─── lifecycle-bypass operations ──────────────────────────────────────
+
+    @abstractmethod
+    def clean(self) -> str:
+        """Remove everything this adapter has materialized. Returns a human-readable
+        description of what was removed. Implementations handle their own short-lived
+        connection if needed — does not require __enter__ to be called first."""
+
+    # ─── incremental state CRUD ───────────────────────────────────────────
+
+    @abstractmethod
+    def fetch_state(self, model_name: str) -> dict[str, tuple[str, str]]:
+        """Return {document_id: (content_hash, code_version)} for `model_name`."""
+
+    @abstractmethod
+    def upsert_state(
+        self, model_name: str, records: list[tuple[str, str, str]]
+    ) -> None: ...
+
+    @abstractmethod
+    def clear_model_state(self, model_name: str) -> None: ...

dbt_ml/adapters/duckdb.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import duckdb
+import polars as pl
+
+from ..config.profile import WarehouseConfig
+from .base import AdapterError, WarehouseAdapter
+from .registry import register
+
+
+@register
+class DuckDBAdapter(WarehouseAdapter):
+    """The reference implementation. Wraps a single DuckDB connection.
+
+    DuckDB-specific wrinkle: the catalog name comes from the database
+    filename's stem; if the schema and the catalog collide (both `dbt_ml`)
+    we have to fully-qualify SQL references as `"catalog"."schema"`.
+    """
+
+    def __init__(self, config: WarehouseConfig, *, project_dir: Path | None = None) -> None:
+        super().__init__(config, project_dir=project_dir)
+        self._con: duckdb.DuckDBPyConnection | None = None
+        self._catalog: str = ""
+
+    @classmethod
+    def adapter_type(cls) -> str:
+        return "duckdb"
+
+    # ─── lifecycle ────────────────────────────────────────────────────────
+
+    def _connect(self) -> None:
+        db_path = self._resolved_path()
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._con = duckdb.connect(str(db_path))
+        row = self._con.execute("SELECT current_database()").fetchone()
+        self._catalog = row[0] if row else "memory"
+
+    def _close(self) -> None:
+        if self._con is not None:
+            self._con.close()
+            self._con = None
+
+    def _ensure_schema(self) -> None:
+        self.connection.execute(f"CREATE SCHEMA IF NOT EXISTS {self.schema_ref}")
+
+    def _ensure_state_table(self) -> None:
+        self.connection.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.schema_ref}.dbt_ml_state (
+                model_name VARCHAR NOT NULL,
+                document_id VARCHAR NOT NULL,
+                content_hash VARCHAR NOT NULL,
+                code_version VARCHAR NOT NULL,
+                last_run_at TIMESTAMP NOT NULL,
+                PRIMARY KEY (model_name, document_id)
+            )
+            """
+        )
+
+    # ─── identity ────────────────────────────────────────────────────────
+
+    @property
+    def connection(self) -> duckdb.DuckDBPyConnection:
+        if self._con is None:
+            raise AdapterError("Adapter must be used as a context manager")
+        return self._con
+
+    @property
+    def raw_connection(self) -> duckdb.DuckDBPyConnection:
+        """The underlying warehouse driver. Handed to custom python tests."""
+        return self.connection
+
+    @property
+    def catalog(self) -> str:
+        if not self._catalog:
+            raise AdapterError("Adapter must be used as a context manager")
+        return self._catalog
+
+    @property
+    def schema_ref(self) -> str:
+        return f'"{self.catalog}"."{self.schema}"'
+
+    # ─── materialization ─────────────────────────────────────────────────
+
+    def materialize_full(self, table: str, df: pl.DataFrame) -> int:
+        full = self.table_ref(table)
+        self.connection.register("dbt_ml_staging", df)
+        try:
+            self.connection.execute(
+                f"CREATE OR REPLACE TABLE {full} AS SELECT * FROM dbt_ml_staging"
+            )
+        finally:
+            self.connection.unregister("dbt_ml_staging")
+        return df.height
+
+    def materialize_incremental(
+        self, table: str, df: pl.DataFrame, *, key_col: str
+    ) -> int:
+        if df.height == 0:
+            return 0
+        full = self.table_ref(table)
+        self.connection.register("dbt_ml_staging", df)
+        try:
+            self.connection.execute(
+                f"CREATE TABLE IF NOT EXISTS {full} AS "
+                f"SELECT * FROM dbt_ml_staging LIMIT 0"
+            )
+            if key_col in df.columns:
+                ids = df[key_col].to_list()
+                if ids:
+                    placeholders = ",".join(["?"] * len(ids))
+                    self.connection.execute(
+                        f'DELETE FROM {full} WHERE "{key_col}" IN ({placeholders})',
+                        ids,
+                    )
+            self.connection.execute(f"INSERT INTO {full} SELECT * FROM dbt_ml_staging")
+        finally:
+            self.connection.unregister("dbt_ml_staging")
+        return df.height
+
+    def drop_table(self, table: str) -> None:
+        self.connection.execute(f"DROP TABLE IF EXISTS {self.table_ref(table)}")
+
+    # ─── querying ────────────────────────────────────────────────────────
+
+    def execute(self, sql: str, params: list[Any] | None = None) -> Any:
+        if params is None:
+            return self.connection.execute(sql)
+        return self.connection.execute(sql, params)
+
+    def query_df(self, sql: str) -> pl.DataFrame:
+        return self.connection.execute(sql).pl()
+
+    def scalar(self, sql: str, params: list[Any] | None = None) -> Any:
+        row = self.execute(sql, params).fetchone()
+        return row[0] if row else None
+
+    def rows(self, sql: str, params: list[Any] | None = None) -> list[tuple]:
+        return self.execute(sql, params).fetchall()
+
+    def clean(self) -> str:
+        """Delete the DuckDB file. Closes the connection if open."""
+        if self._con is not None:
+            self._close()
+        path = self._resolved_path()
+        if path.exists():
+            path.unlink()
+        return str(path)
+
+    def list_tables(self) -> list[str]:
+        rows = self.connection.execute(
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE table_catalog = ? AND table_schema = ? AND table_name != 'dbt_ml_state' "
+            "ORDER BY table_name",
+            [self.catalog, self.schema],
+        ).fetchall()
+        return [r[0] for r in rows]
+
+    # ─── state CRUD ──────────────────────────────────────────────────────
+
+    def fetch_state(self, model_name: str) -> dict[str, tuple[str, str]]:
+        rows = self.connection.execute(
+            f"SELECT document_id, content_hash, code_version "
+            f"FROM {self.schema_ref}.dbt_ml_state WHERE model_name = ?",
+            [model_name],
+        ).fetchall()
+        return {r[0]: (r[1], r[2]) for r in rows}
+
+    def upsert_state(
+        self, model_name: str, records: list[tuple[str, str, str]]
+    ) -> None:
+        if not records:
+            return
+        self.connection.executemany(
+            f"""
+            INSERT INTO {self.schema_ref}.dbt_ml_state
+                (model_name, document_id, content_hash, code_version, last_run_at)
+            VALUES (?, ?, ?, ?, current_timestamp)
+            ON CONFLICT (model_name, document_id) DO UPDATE SET
+                content_hash = excluded.content_hash,
+                code_version = excluded.code_version,
+                last_run_at  = excluded.last_run_at
+            """,
+            [[model_name, doc_id, ch, cv] for doc_id, ch, cv in records],
+        )
+
+    def clear_model_state(self, model_name: str) -> None:
+        self.connection.execute(
+            f"DELETE FROM {self.schema_ref}.dbt_ml_state WHERE model_name = ?",
+            [model_name],
+        )
+
+    # ─── internals ───────────────────────────────────────────────────────
+
+    def _resolved_path(self) -> Path:
+        path = self.config.path
+        if path.is_absolute() or self.project_dir is None:
+            return path.resolve()
+        return (self.project_dir / path).resolve()

dbt_ml/adapters/registry.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from ..config.profile import WarehouseConfig
+from .base import AdapterError, WarehouseAdapter
+
+
+class UnknownAdapterError(AdapterError):
+    pass
+
+
+_REGISTRY: dict[str, type[WarehouseAdapter]] = {}
+
+
+def register(cls: type[WarehouseAdapter]) -> type[WarehouseAdapter]:
+    _REGISTRY[cls.adapter_type()] = cls
+    return cls
+
+
+def create_adapter(
+    config: WarehouseConfig, *, project_dir: Path | None = None
+) -> WarehouseAdapter:
+    cls = _REGISTRY.get(config.type)
+    if cls is None:
+        raise UnknownAdapterError(
+            f"No adapter registered for warehouse.type='{config.type}'. "
+            f"Known: {sorted(_REGISTRY)}"
+        )
+    return cls(config, project_dir=project_dir)
+
+
+def list_adapter_types() -> list[str]:
+    return sorted(_REGISTRY)

dbt_ml/backends/__init__.py

@@ -0,0 +1,19 @@
+from . import (
+    email_backend,  # noqa: F401  # side-effect: registers EmailBackend
+    html_backend,  # noqa: F401  # side-effect: registers HtmlBackend
+    json_backend,  # noqa: F401  # side-effect: registers JsonBackend
+    llm_backend,  # noqa: F401  # side-effect: registers LLMBackend
+    markdown_backend,  # noqa: F401  # side-effect: registers MarkdownBackend
+    pdf_backend,  # noqa: F401  # side-effect: registers PdfBackend
+)
+from .base import BaseBackend, ExtractionResult
+from .registry import BackendNotFoundError, get_backend, list_backends, register
+
+__all__ = [
+    "BackendNotFoundError",
+    "BaseBackend",
+    "ExtractionResult",
+    "get_backend",
+    "list_backends",
+    "register",
+]

dbt_ml/backends/base.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class ExtractionResult:
+    """Output of a single document extraction.
+
+    `fields` holds the projected field values. `warnings` collects
+    non-fatal issues surfaced by the backend.
+    """
+
+    fields: dict[str, Any]
+    warnings: list[str] = field(default_factory=list)
+
+
+class BaseBackend(ABC):
+    """Contract every extraction backend implements."""
+
+    @abstractmethod
+    def name(self) -> str: ...
+
+    @abstractmethod
+    def supported_formats(self) -> list[str]: ...
+
+    @abstractmethod
+    def extract(self, path: Path, options: dict[str, Any]) -> ExtractionResult: ...
+
+    def validate(self) -> None:
+        """Raise if the backend's runtime deps are missing. Default: no-op."""
+        return None

dbt_ml/backends/email_backend.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from email import message_from_bytes
+from email.message import Message
+from email.utils import parsedate_to_datetime
+from pathlib import Path
+from typing import Any
+
+from .base import BaseBackend, ExtractionResult
+from .registry import register
+
+
+@register
+class EmailBackend(BaseBackend):
+    """Read .eml files via stdlib `email`. No external dependencies.
+
+    Options:
+        include_body:      Emit the plaintext body (default True).
+        body_field:        Field name for the body (default "body").
+        include_html:      Emit the HTML alternative if present (default False).
+        include_headers:   Emit a `headers` dict of all headers (default False).
+    """
+
+    def name(self) -> str:
+        return "email"
+
+    def supported_formats(self) -> list[str]:
+        return [".eml"]
+
+    def extract(self, path: Path, options: dict[str, Any]) -> ExtractionResult:
+        msg: Message = message_from_bytes(path.read_bytes())
+        warnings: list[str] = []
+        fields: dict[str, Any] = {
+            "from": msg.get("From"),
+            "to": msg.get("To"),
+            "cc": msg.get("Cc"),
+            "subject": msg.get("Subject"),
+            "date": _parse_date(msg.get("Date")),
+            "message_id": msg.get("Message-ID"),
+        }
+
+        text_body, html_body = _walk_parts(msg)
+        if options.get("include_body", True):
+            body_field = options.get("body_field", "body")
+            if text_body is None and html_body is not None:
+                warnings.append(f"{path.name}: no text/plain part, fell back to text/html")
+                fields[body_field] = html_body
+            else:
+                fields[body_field] = text_body or ""
+                if not (text_body or html_body):
+                    warnings.append(f"{path.name}: no text or html body found")
+
+        if options.get("include_html", False):
+            fields["html_body"] = html_body
+
+        if options.get("include_headers", False):
+            fields["headers"] = {k: v for k, v in msg.items()}
+
+        return ExtractionResult(fields=fields, warnings=warnings)
+
+
+def _parse_date(raw: str | None) -> str | None:
+    if not raw:
+        return None
+    try:
+        return parsedate_to_datetime(raw).isoformat()
+    except (TypeError, ValueError):
+        return raw
+
+
+def _walk_parts(msg: Message) -> tuple[str | None, str | None]:
+    """Return (text_body, html_body) — preferring multipart/alternative parts."""
+    text_body: str | None = None
+    html_body: str | None = None
+    for part in msg.walk():
+        ctype = part.get_content_type()
+        if part.is_multipart():
+            continue
+        try:
+            payload = part.get_payload(decode=True)
+        except Exception:
+            continue
+        if payload is None:
+            continue
+        try:
+            decoded = payload.decode(part.get_content_charset() or "utf-8", errors="replace")
+        except (LookupError, ValueError):
+            decoded = payload.decode("utf-8", errors="replace")
+        if ctype == "text/plain" and text_body is None:
+            text_body = decoded
+        elif ctype == "text/html" and html_body is None:
+            html_body = decoded
+    return text_body, html_body

dbt_ml/backends/html_backend.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+from .base import BaseBackend, ExtractionResult
+from .registry import register
+
+
+@register
+class HtmlBackend(BaseBackend):
+    """Read .html files via BeautifulSoup.
+
+    Options:
+        text_field:        Field name for the plain-text body (default "text").
+        include_text:      Emit body text with tags stripped (default True).
+        selectors:         dict of {field_name: css_selector}. First match's text
+                           per selector is emitted; missing selectors yield None
+                           with a warning.
+        include_meta:      Emit a `meta` dict of <meta> name→content pairs.
+        include_opengraph: Emit `og` dict of OpenGraph properties (og:*).
+        include_links:     Emit `links` as a list of href strings.
+        parser:            "html.parser" (default, stdlib) or "lxml" if installed.
+    """
+
+    def name(self) -> str:
+        return "html"
+
+    def supported_formats(self) -> list[str]:
+        return [".html", ".htm"]
+
+    def extract(self, path: Path, options: dict[str, Any]) -> ExtractionResult:
+        parser = options.get("parser", "html.parser")
+        soup = BeautifulSoup(path.read_text(), parser)
+
+        warnings: list[str] = []
+        fields: dict[str, Any] = {}
+
+        if options.get("include_text", True):
+            text_field = options.get("text_field", "text")
+            body = soup.body or soup
+            fields[text_field] = body.get_text(separator="\n", strip=True)
+
+        selectors = options.get("selectors") or {}
+        for field_name, selector in selectors.items():
+            match = soup.select_one(selector)
+            if match is None:
+                warnings.append(
+                    f"selector {selector!r} for field '{field_name}' matched nothing"
+                )
+                fields[field_name] = None
+            else:
+                fields[field_name] = match.get_text(strip=True)
+
+        if options.get("include_meta", False):
+            meta: dict[str, str] = {}
+            for tag in soup.find_all("meta"):
+                key = tag.get("name") or tag.get("property")
+                content = tag.get("content")
+                if key and content:
+                    meta[str(key)] = str(content)
+            fields["meta"] = meta
+
+        if options.get("include_opengraph", False):
+            og: dict[str, str] = {}
+            for tag in soup.find_all("meta"):
+                prop = tag.get("property")
+                if prop and isinstance(prop, str) and prop.startswith("og:"):
+                    og[prop[3:]] = str(tag.get("content") or "")
+            fields["og"] = og
+
+        if options.get("include_links", False):
+            fields["links"] = [
+                str(a.get("href"))
+                for a in soup.find_all("a", href=True)
+            ]
+
+        return ExtractionResult(fields=fields, warnings=warnings)

dbt_ml/backends/json_backend.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from .base import BaseBackend, ExtractionResult
+from .registry import register
+
+
+@register
+class JsonBackend(BaseBackend):
+    def name(self) -> str:
+        return "json"
+
+    def supported_formats(self) -> list[str]:
+        return [".json"]
+
+    def extract(self, path: Path, options: dict[str, Any]) -> ExtractionResult:
+        with path.open() as f:
+            data = json.load(f)
+
+        if not isinstance(data, dict):
+            raise ValueError(
+                f"json backend expects each document to be a JSON object; "
+                f"got {type(data).__name__} from {path}"
+            )
+
+        fields_to_project = options.get("fields")
+        warnings: list[str] = []
+        if fields_to_project:
+            extracted: dict[str, Any] = {}
+            for key in fields_to_project:
+                if key not in data:
+                    warnings.append(f"Field '{key}' missing in {path.name}")
+                extracted[key] = data.get(key)
+        else:
+            extracted = dict(data)
+
+        return ExtractionResult(fields=extracted, warnings=warnings)