sqlite_data_loader.py

"""SQLite data loader — example plugin for Data Formulator.

This file is a complete, working example of a Data Formulator plugin.
SQLite is part of the Python standard library, so this loader has **no
extra dependencies** and is a good template for building your own.

How to install
--------------
1. Make sure Data Formulator runs in single-user mode
   (the default — ``WORKSPACE_BACKEND`` unset or ``local``).
2. Copy this file to ``~/.data_formulator/plugins/`` (or whatever
   directory ``DF_PLUGIN_DIR`` points to).  Filename must end in
   ``_data_loader.py``.
3. Restart Data Formulator.  A new "sqlite" connector should appear
   alongside the built-ins.

How to test it quickly
----------------------
::

    sqlite3 /tmp/demo.db <<'SQL'
    CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER);
    INSERT INTO users VALUES (1, 'Alice', 30), (2, 'Bob', 25);
    SQL

Then in DF, add a SQLite connector pointing at ``/tmp/demo.db``.

What this example demonstrates
------------------------------
* Implementing every abstract method on :class:`ExternalDataLoader`.
* Declaring connection params via :meth:`list_params` so the UI
  auto-renders a config form.
* Identifier quoting to avoid SQL-injection in table/column names.
* Returning data directly as a PyArrow Table (no pandas in the hot path).
* Read-only access (``mode=ro`` URI) — the loader never writes to the DB.
"""

from __future__ import annotations

import logging
import os
import sqlite3
import threading
from typing import Any
from urllib.parse import quote as _url_quote

import pyarrow as pa

from data_formulator.data_loader.external_data_loader import (
    ExternalDataLoader,
    MAX_IMPORT_ROWS,
)

logger = logging.getLogger(__name__)


def _quote_ident(name: str) -> str:
    """Quote a SQLite identifier safely (table or column name)."""
    return '"' + str(name).replace('"', '""') + '"'


class SQLiteDataLoader(ExternalDataLoader):
    """Read tables from a local SQLite database file."""

    # Override the default title-casing of the registry key ("Sqlite") with
    # a properly-cased product name.  Optional; remove if you don't care.
    DISPLAY_NAME = "SQLite"

    # ------------------------------------------------------------------ #
    # Static metadata: what the UI shows + how to configure              #
    # ------------------------------------------------------------------ #

    @staticmethod
    def list_params() -> list[dict[str, Any]]:
        return [
            {
                "name": "database_path",
                "type": "string",
                "required": True,
                "default": "",
                "tier": "connection",
                "description": "Absolute path to a .db / .sqlite file",
            },
        ]

    @staticmethod
    def auth_instructions() -> str:
        return (
            "**Example:** `/tmp/demo.db`\n\n"
            "Provide the absolute path to a local SQLite database file. "
            "The file must exist and be readable. The loader opens it "
            "read-only — your data is never modified.\n\n"
            "**Create a test DB:**\n"
            "```\nsqlite3 /tmp/demo.db 'CREATE TABLE t(a,b); "
            "INSERT INTO t VALUES (1,2),(3,4);'\n```"
        )

    # ------------------------------------------------------------------ #
    # Connection                                                          #
    # ------------------------------------------------------------------ #

    def __init__(self, params: dict[str, Any]):
        self.params = params or {}
        self.database_path = (self.params.get("database_path") or "").strip()

        if not self.database_path:
            raise ValueError("SQLite database_path is required")
        if not os.path.isfile(self.database_path):
            raise ValueError(f"SQLite database not found: {self.database_path}")

        # Open read-only via URI so we cannot accidentally mutate the DB.
        uri = f"file:{_url_quote(self.database_path)}?mode=ro"
        try:
            self._conn = sqlite3.connect(
                uri, uri=True, check_same_thread=False, isolation_level=None,
            )
        except sqlite3.Error as e:
            raise ValueError(
                f"Failed to open SQLite database '{self.database_path}': {e}"
            ) from e

        self._lock = threading.Lock()
        logger.info("SQLite plugin connected: %s", self.database_path)

    # ------------------------------------------------------------------ #
    # Helpers                                                             #
    # ------------------------------------------------------------------ #

    def _read_sql(self, query: str) -> pa.Table:
        """Run a query and return the result as a PyArrow Table."""
        with self._lock:
            cur = self._conn.execute(query)
            if cur.description is None:
                return pa.table({})
            columns = [d[0] for d in cur.description]
            rows = cur.fetchall()
        # Transpose rows to per-column lists; pyarrow infers types.
        col_data: dict[str, list[Any]] = {c: [] for c in columns}
        for row in rows:
            for col, value in zip(columns, row):
                col_data[col].append(value)
        return pa.table(col_data)

    # ------------------------------------------------------------------ #
    # Required loader API                                                 #
    # ------------------------------------------------------------------ #

    def list_tables(self, table_filter: str | None = None) -> list[dict[str, Any]]:
        """List every table + view in the database."""
        names_tbl = self._read_sql(
            "SELECT name, type FROM sqlite_master "
            "WHERE type IN ('table', 'view') "
            "AND name NOT LIKE 'sqlite_%' "
            "ORDER BY name"
        )
        if names_tbl.num_rows == 0:
            return []

        results: list[dict[str, Any]] = []
        for name in names_tbl.column("name").to_pylist():
            if table_filter and table_filter.lower() not in name.lower():
                continue
            # PRAGMA table_info gives column name + declared type.
            info = self._read_sql(f"PRAGMA table_info({_quote_ident(name)})")
            columns = [
                {"name": n, "type": t or "ANY"}
                for n, t in zip(
                    info.column("name").to_pylist() if info.num_rows else [],
                    info.column("type").to_pylist() if info.num_rows else [],
                )
            ]
            count_tbl = self._read_sql(
                f"SELECT COUNT(*) AS n FROM {_quote_ident(name)}"
            )
            row_count = int(count_tbl.column("n")[0].as_py()) if count_tbl.num_rows else 0
            results.append({
                "name": name,
                "metadata": {"columns": columns, "row_count": row_count},
            })
        return results

    def fetch_data_as_arrow(
        self,
        source_table: str,
        import_options: dict[str, Any] | None = None,
    ) -> pa.Table:
        """Fetch rows from ``source_table`` as a PyArrow Table."""
        if not source_table:
            raise ValueError("source_table must be provided")

        opts = import_options or {}
        size = min(int(opts.get("size", MAX_IMPORT_ROWS)), MAX_IMPORT_ROWS)
        sort_columns = opts.get("sort_columns") or []
        sort_order = "DESC" if str(opts.get("sort_order", "asc")).lower() == "desc" else "ASC"

        query = f"SELECT * FROM {_quote_ident(source_table)}"
        if sort_columns:
            order_by = ", ".join(f"{_quote_ident(c)} {sort_order}" for c in sort_columns)
            query += f" ORDER BY {order_by}"
        query += f" LIMIT {int(size)}"

        logger.info("SQLite plugin query: %s", query)
        return self._read_sql(query)