mirror of
https://github.com/microsoft/data-formulator.git
synced 2026-11-04 12:01:37 +00:00
204 lines
7.7 KiB
Python
204 lines
7.7 KiB
Python
"""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)
|