sqlide 0.1.0__py3-none-any.whl

sqlide/db/session.py

@@ -0,0 +1,339 @@
+"""One JDBC connection = one DbSession.
+
+JDBC connections are not thread-safe, so every JDBC call of a session runs on that
+session's own single worker thread. The UI only awaits; cancel() may be called from
+any thread (Statement.cancel is the one call JDBC allows concurrently).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import threading
+import time
+from collections.abc import Callable, Iterator
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, TypeVar, cast
+
+import jpype
+
+from sqlide.db.result import Column, DbError, Execution, ResultItem
+from sqlide.db.types import make_converter
+from sqlide.drivers.loader import LoadedDriver
+
+T = TypeVar("T")
+
+DEFAULT_PAGE = 500
+
+
+def _db_error(e: Exception, cancelled: bool = False) -> DbError:
+    if cancelled:
+        return DbError("Cancelled by user")
+    if isinstance(e, jpype.JException):
+        state = getattr(e, "getSQLState", lambda: "")() or ""
+        code = getattr(e, "getErrorCode", lambda: 0)()
+        return DbError(str(e.getMessage() or e), str(state), int(code))
+    return DbError(str(e))
+
+
+class _Cursor:
+    """Open JDBC ResultSet; satisfies result.RowSource."""
+
+    def __init__(self, session: DbSession, rs: Any, stmt: Any, convs: list, pending: bool) -> None:
+        self._s, self._rs, self._stmt, self._convs = session, rs, stmt, convs
+        self._pending = pending  # rs is already positioned on an unread row
+        self.closed = False
+        self.owns_stmt = False  # True once handed to the caller; before that execute() owns it
+
+    def _read(self, n: int) -> tuple[list[tuple], bool]:
+        if self.closed:
+            raise DbError("Cursor is closed (another query ran). Re-run to load more rows.")
+        rows: list[tuple] = []
+        try:
+            while len(rows) < n:
+                if self._pending:
+                    self._pending = False
+                elif not self._rs.next():
+                    self._close_sync()
+                    return rows, True
+                rows.append(tuple(c(self._rs) for c in self._convs))
+            if self._rs.next():
+                self._pending = True
+                return rows, False
+            self._close_sync()
+            return rows, True
+        except jpype.JException as e:
+            self._close_sync()
+            raise _db_error(e, self._s._cancelled) from e
+
+    def _close_sync(self) -> None:
+        if self.closed:
+            return
+        self.closed = True
+        for o in (self._rs, self._stmt) if self.owns_stmt else (self._rs,):
+            with contextlib.suppress(jpype.JException):
+                o.close()
+        self._s._cursors.discard(self)
+        if self.owns_stmt:
+            self._s._leave_paging_tx()
+
+    async def fetch_more(self, n: int) -> tuple[list[tuple], bool]:
+        return await self._s.call(lambda _c: self._read(n))
+
+    async def close(self) -> None:
+        await self._s.call(lambda _c: self._close_sync())
+
+
+class DbSession:
+    def __init__(
+        self,
+        loaded: LoadedDriver,
+        url: str,
+        user: str = "",
+        password: str | None = None,
+        properties: dict[str, str] | None = None,
+        dialect: str = "generic",
+        autocommit: bool = True,
+    ) -> None:
+        self._loaded, self._url, self._dialect = loaded, url, dialect
+        self._props = dict(properties or {})
+        if user:
+            self._props["user"] = user
+        if password is not None:
+            self._props["password"] = password
+        self._want_autocommit = autocommit
+        self._conn: Any = None
+        self._pool = ThreadPoolExecutor(1, "sqlide-jdbc", initializer=self._init_thread)
+        self._lock = threading.Lock()
+        self._stmt: Any = None
+        self._cancelled = False
+        self._cursors: set[_Cursor] = set()
+        self._paging_tx = False  # we turned autocommit off only to get server-side cursors
+        self.pending_tx = False  # manual mode: statements since last commit/rollback
+        self.product = ""
+
+    # --- plumbing ---
+    def _init_thread(self) -> None:
+        # drivers using ServiceLoader (ClickHouse) look classes up via the context loader
+        jpype.JClass("java.lang.Thread").currentThread().setContextClassLoader(self._loaded.loader)
+
+    async def call(self, fn: Callable[[Any], T]) -> T:
+        """Run fn(java_connection) on the session thread. Use for metadata and misc JDBC."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self._pool, lambda: fn(self._conn))
+
+    # --- lifecycle ---
+    async def open(self) -> None:
+        def connect(_: Any) -> None:
+            props = jpype.JClass("java.util.Properties")()
+            for k, v in self._props.items():
+                props.setProperty(k, v)
+            try:
+                conn = self._loaded.driver.connect(self._url, props)
+                if conn is None:
+                    raise DbError(f"Driver does not accept URL: {self._url}")
+                conn.setAutoCommit(self._want_autocommit)
+                md = conn.getMetaData()
+                self.product = f"{md.getDatabaseProductName()} {md.getDatabaseProductVersion()}"
+            except jpype.JException as e:
+                raise _db_error(e) from e
+            self._conn = conn
+
+        await self.call(connect)
+
+    async def close(self) -> None:
+        if self._conn is None:
+            self._pool.shutdown(wait=False)
+            return
+        self.cancel()
+
+        def shut(conn: Any) -> None:
+            for c in list(self._cursors):
+                c._close_sync()
+            if self.pending_tx:  # some drivers (Oracle) would commit on close
+                with contextlib.suppress(jpype.JException):
+                    conn.rollback()
+            with contextlib.suppress(jpype.JException):
+                conn.close()
+
+        await self.call(shut)
+        self._conn = None
+        self._pool.shutdown(wait=False)
+
+    @property
+    def connected(self) -> bool:
+        return self._conn is not None
+
+    # --- execution ---
+    async def execute(self, sql: str, page_size: int = DEFAULT_PAGE) -> Execution:
+        return await self.call(lambda _c: self._execute_sync(sql, page_size))
+
+    async def stream(
+        self,
+        sql: str,
+        consumer: Callable[[list[Column], Iterator[tuple]], T],
+        page_size: int = 2000,
+    ) -> T:
+        """Run `sql`, hand its first result set to `consumer` as a lazy row iterator.
+
+        `consumer` runs on the session thread, so it may block (write a file). Pages are
+        fetched on demand: memory stays flat for any result size.
+        """
+        return await self.call(lambda _c: self._stream_sync(sql, page_size, consumer))
+
+    def _stream_sync(
+        self, sql: str, page_size: int, consumer: Callable[[list[Column], Iterator[tuple]], T]
+    ) -> T:
+        ex = self._execute_sync(sql, page_size)
+        item = next((i for i in ex.items if i.has_rows), None)
+        if item is None:
+            raise DbError("The statement did not return a result set")
+
+        cur = cast("_Cursor | None", item.cursor)
+
+        def rows() -> Iterator[tuple]:
+            yield from item.rows
+            while cur is not None and not cur.closed:
+                page, done = cur._read(page_size)
+                yield from page
+                if done:
+                    break
+
+        try:
+            return consumer(item.columns, rows())
+        finally:
+            if cur is not None:
+                cur._close_sync()
+
+    def cancel(self) -> None:
+        """Thread-safe. Aborts the running statement, if any."""
+        with self._lock:
+            stmt = self._stmt
+        if stmt is not None:
+            self._cancelled = True
+            with contextlib.suppress(jpype.JException):
+                stmt.cancel()
+
+    def _execute_sync(self, sql: str, page_size: int) -> Execution:
+        for c in list(self._cursors):  # one live result per session, see module docs
+            c._close_sync()
+        self._cancelled = False
+        self._enter_paging_tx()
+        t0 = time.perf_counter()
+        stmt = self._conn.createStatement()
+        keep_open = False
+        try:
+            stmt.setFetchSize(page_size)
+            with self._lock:
+                self._stmt = stmt
+            items: list[ResultItem] = []
+            has_rs = bool(stmt.execute(sql))
+            while True:
+                if has_rs:
+                    item, more = self._read_first_page(stmt.getResultSet(), stmt, page_size)
+                    items.append(item)
+                    if more:  # cursor owns the statement now; further results are not reachable
+                        keep_open = True
+                        break
+                else:
+                    n = int(stmt.getUpdateCount())
+                    if n == -1:
+                        break
+                    items.append(ResultItem(update_count=n))
+                has_rs = bool(stmt.getMoreResults())
+            warnings = self._warnings(stmt)
+            if not self._want_autocommit:
+                self.pending_tx = True
+            return Execution(sql, items, warnings, time.perf_counter() - t0)
+        except jpype.JException as e:
+            err = _db_error(e, self._cancelled)
+            keep_open = False
+            self._rollback_paging_tx()
+            raise err from e
+        finally:
+            with self._lock:
+                self._stmt = None
+            if not keep_open:
+                with contextlib.suppress(jpype.JException):
+                    stmt.close()
+                self._leave_paging_tx()
+
+    def _read_first_page(self, rs: Any, stmt: Any, page_size: int) -> tuple[ResultItem, bool]:
+        md = rs.getMetaData()
+        n = int(md.getColumnCount())
+        cols = [
+            Column(
+                str(md.getColumnLabel(i)), str(md.getColumnTypeName(i)), int(md.getColumnType(i))
+            )
+            for i in range(1, n + 1)
+        ]
+        convs = [make_converter(i, c.jdbc_type) for i, c in zip(range(1, n + 1), cols, strict=True)]
+        cur = _Cursor(self, rs, stmt, convs, pending=False)
+        rows, done = cur._read(page_size)
+        item = ResultItem(columns=cols, rows=rows)
+        if done:
+            return item, False
+        cur.owns_stmt = True
+        self._cursors.add(cur)
+        item.cursor = cur
+        return item, True
+
+    @staticmethod
+    def _warnings(stmt: Any) -> list[str]:
+        out, w = [], stmt.getWarnings()
+        while w is not None:
+            out.append(str(w.getMessage()))
+            w = w.getNextWarning()
+        return out
+
+    # --- postgres: fetchSize is ignored in autocommit mode, so page inside a transaction ---
+    def _enter_paging_tx(self) -> None:
+        if self._dialect == "postgres" and self._want_autocommit and not self._paging_tx:
+            self._conn.setAutoCommit(False)
+            self._paging_tx = True
+
+    def _leave_paging_tx(self) -> None:
+        if self._paging_tx and not self._cursors and self._conn is not None:
+            self._paging_tx = False
+            with contextlib.suppress(jpype.JException):
+                self._conn.commit()
+                self._conn.setAutoCommit(True)
+
+    def _rollback_paging_tx(self) -> None:
+        if self._paging_tx:
+            self._paging_tx = False
+            with contextlib.suppress(jpype.JException):
+                self._conn.rollback()
+                self._conn.setAutoCommit(True)
+
+    # --- transactions ---
+    async def set_autocommit(self, on: bool) -> None:
+        def f(conn: Any) -> None:
+            try:
+                conn.setAutoCommit(on)
+            except jpype.JException as e:
+                raise _db_error(e) from e
+            self._want_autocommit = on
+            self.pending_tx = False
+
+        await self.call(f)
+
+    async def commit(self) -> None:
+        await self._finish_tx("commit")
+
+    async def rollback(self) -> None:
+        await self._finish_tx("rollback")
+
+    async def _finish_tx(self, how: str) -> None:
+        def f(conn: Any) -> None:
+            try:
+                getattr(conn, how)()
+            except jpype.JException as e:
+                raise _db_error(e) from e
+            self.pending_tx = False
+
+        await self.call(f)
+
+    @property
+    def autocommit(self) -> bool:
+        return self._want_autocommit

sqlide/db/types.py

@@ -0,0 +1,84 @@
+"""JDBC value -> Python value. Values stay typed so sorting and xlsx export are correct."""
+
+from __future__ import annotations
+
+import datetime as dt
+from collections.abc import Callable
+from decimal import Decimal
+from typing import Any
+
+# java.sql.Types constants (stable by spec)
+BIT, TINYINT, SMALLINT, INTEGER, BIGINT = -7, -6, 5, 4, -5
+FLOAT, REAL, DOUBLE, NUMERIC, DECIMAL = 6, 7, 8, 2, 3
+DATE, TIME, TIMESTAMP, BOOLEAN = 91, 92, 93, 16
+BINARY, VARBINARY, LONGVARBINARY, BLOB = -2, -3, -4, 2004
+
+_INTS = {TINYINT, SMALLINT, INTEGER, BIGINT}
+_FLOATS = {FLOAT, REAL, DOUBLE}
+_DECIMALS = {NUMERIC, DECIMAL}
+_BOOLS = {BIT, BOOLEAN}
+_BYTES = {BINARY, VARBINARY, LONGVARBINARY, BLOB}
+
+Converter = Callable[[Any], Any]  # (java ResultSet) -> python value
+
+
+def make_converter(index: int, jdbc_type: int) -> Converter:
+    """Converter for column `index` (1-based). Any conversion failure falls back to text."""
+    primary = _primary(index, jdbc_type)
+
+    def convert(rs: Any) -> Any:
+        try:
+            return primary(rs)
+        except Exception:  # java exceptions too: out-of-range dates, unsigned bigint, ...
+            try:
+                s = rs.getString(index)
+                return None if s is None else str(s)
+            except Exception:
+                return "<unreadable>"
+
+    return convert
+
+
+def _primary(i: int, t: int) -> Converter:
+    if t in _INTS:
+        return lambda rs: _prim(rs, int(rs.getLong(i)))
+    if t in _FLOATS:
+        return lambda rs: _prim(rs, float(rs.getDouble(i)))
+    if t in _BOOLS:
+        return lambda rs: _prim(rs, bool(rs.getBoolean(i)))
+    if t in _DECIMALS:
+        return lambda rs: _obj(rs.getBigDecimal(i), lambda v: Decimal(str(v.toPlainString())))
+    if t == DATE:
+        return lambda rs: _obj(rs.getDate(i), lambda v: _date(v.toLocalDate()))
+    if t == TIME:
+        return lambda rs: _obj(rs.getTime(i), lambda v: _time(v.toLocalTime()))
+    if t == TIMESTAMP:
+        return lambda rs: _obj(rs.getTimestamp(i), lambda v: _datetime(v.toLocalDateTime()))
+    if t in _BYTES:
+        return lambda rs: _obj(rs.getBytes(i), bytes)
+    return lambda rs: _obj(rs.getString(i), str)
+
+
+def _prim(rs: Any, value: Any) -> Any:
+    return None if rs.wasNull() else value
+
+
+def _obj(value: Any, f: Callable[[Any], Any]) -> Any:
+    return None if value is None else f(value)
+
+
+def _date(d: Any) -> dt.date:
+    return dt.date(int(d.getYear()), int(d.getMonthValue()), int(d.getDayOfMonth()))
+
+
+def _time(t: Any) -> dt.time:
+    return dt.time(
+        int(t.getHour()), int(t.getMinute()), int(t.getSecond()), int(t.getNano()) // 1000
+    )
+
+
+def _datetime(v: Any) -> dt.datetime:
+    return dt.datetime(
+        int(v.getYear()), int(v.getMonthValue()), int(v.getDayOfMonth()),
+        int(v.getHour()), int(v.getMinute()), int(v.getSecond()), int(v.getNano()) // 1000,
+    )  # fmt: skip

sqlide/doctor.py

@@ -0,0 +1,60 @@
+"""`sqlide doctor`: environment checks. Each check is independent; add new ones to CHECKS."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from sqlide.config import paths
+from sqlide.drivers.registry import DriverRegistry
+from sqlide.jvm.locate import JvmNotFound, locate_jvm
+
+Check = Callable[[], tuple[bool, str]]
+
+
+def check_jvm() -> tuple[bool, str]:
+    try:
+        info = locate_jvm()
+    except JvmNotFound as e:
+        return False, str(e)
+    return True, f"Java {info.major} at {info.home}"
+
+
+def check_drivers() -> tuple[bool, str]:
+    reg = DriverRegistry()
+    have = [d for d in reg.all() if reg.is_installed(d)]
+    return bool(have), ", ".join(have) if have else "none installed: sqlide driver install postgres"
+
+
+def check_dirs() -> tuple[bool, str]:
+    return True, f"config {paths.config_dir()}, data {paths.data_dir()}"
+
+
+def check_clipboard() -> tuple[bool, str]:
+    from sqlide import clipboard
+
+    # OSC52 still works without a native tool, so this is informational, never a failure
+    return True, clipboard.available()
+
+
+def check_exporters() -> tuple[bool, str]:
+    from sqlide import export
+
+    return True, ", ".join(e.name for e in export.all_exporters())
+
+
+CHECKS: list[tuple[str, Check]] = [
+    ("java", check_jvm),
+    ("drivers", check_drivers),
+    ("clipboard", check_clipboard),
+    ("export", check_exporters),
+    ("dirs", check_dirs),
+]
+
+
+def run() -> int:
+    bad = 0
+    for label, fn in CHECKS:
+        ok, detail = fn()
+        bad += not ok
+        print(f"[{'ok' if ok else '!!'}] {label:<10} {detail}")
+    return 1 if bad else 0

sqlide/drivers/__init__.py

@@ -0,0 +1 @@
+"""Drivers block: catalog of JDBC drivers, Maven download, class loading."""

sqlide/drivers/catalog.toml

@@ -0,0 +1,103 @@
+# Built-in driver presets. Users add/override entries in drivers.toml (same keys).
+# classifiers: tried in order, first that exists on Maven Central wins ("" = plain jar).
+# version_suffix: keep only versions ending with it (mssql ships jre8/jre11 builds).
+
+[[driver]]
+id = "postgres"
+name = "PostgreSQL / Greenplum"
+class_name = "org.postgresql.Driver"
+url_template = "jdbc:postgresql://{host}:{port}/{database}"
+default_port = 5432
+dialect = "postgres"
+group = "org.postgresql"
+artifact = "postgresql"
+classifiers = [""]
+
+[[driver]]
+id = "clickhouse"
+name = "ClickHouse"
+class_name = "com.clickhouse.jdbc.ClickHouseDriver"
+url_template = "jdbc:clickhouse://{host}:{port}/{database}"
+default_port = 8123
+dialect = "clickhouse"
+group = "com.clickhouse"
+artifact = "clickhouse-jdbc"
+classifiers = ["all-dependencies", "all"]
+
+[[driver]]
+id = "mysql"
+name = "MySQL"
+class_name = "com.mysql.cj.jdbc.Driver"
+url_template = "jdbc:mysql://{host}:{port}/{database}"
+default_port = 3306
+dialect = "mysql"
+group = "com.mysql"
+artifact = "mysql-connector-j"
+classifiers = [""]
+
+[[driver]]
+id = "mariadb"
+name = "MariaDB"
+class_name = "org.mariadb.jdbc.Driver"
+url_template = "jdbc:mariadb://{host}:{port}/{database}"
+default_port = 3306
+dialect = "mysql"
+group = "org.mariadb.jdbc"
+artifact = "mariadb-java-client"
+classifiers = [""]
+
+[[driver]]
+id = "oracle"
+name = "Oracle"
+class_name = "oracle.jdbc.OracleDriver"
+url_template = "jdbc:oracle:thin:@//{host}:{port}/{database}"
+default_port = 1521
+dialect = "oracle"
+group = "com.oracle.database.jdbc"
+artifact = "ojdbc11"
+classifiers = [""]
+
+[[driver]]
+id = "mssql"
+name = "SQL Server"
+class_name = "com.microsoft.sqlserver.jdbc.SQLServerDriver"
+url_template = "jdbc:sqlserver://{host}:{port};databaseName={database}"
+default_port = 1433
+dialect = "mssql"
+group = "com.microsoft.sqlserver"
+artifact = "mssql-jdbc"
+classifiers = [""]
+version_suffix = "jre11"
+
+[[driver]]
+id = "h2"
+name = "H2"
+class_name = "org.h2.Driver"
+url_template = "jdbc:h2:{database}"
+default_port = 0
+dialect = "generic"
+group = "com.h2database"
+artifact = "h2"
+classifiers = [""]
+
+[[driver]]
+id = "sqlite"
+name = "SQLite"
+class_name = "org.sqlite.JDBC"
+url_template = "jdbc:sqlite:{database}"
+default_port = 0
+dialect = "sqlite"
+group = "org.xerial"
+artifact = "sqlite-jdbc"
+classifiers = [""]
+
+[[driver]]
+id = "duckdb"
+name = "DuckDB"
+class_name = "org.duckdb.DuckDBDriver"
+url_template = "jdbc:duckdb:{database}"
+default_port = 0
+dialect = "generic"
+group = "org.duckdb"
+artifact = "duckdb_jdbc"
+classifiers = [""]

sqlide/drivers/cli.py

@@ -0,0 +1,74 @@
+"""`sqlide driver ...` subcommands."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from sqlide.config._toml import ConfigError
+from sqlide.drivers.custom import build_custom_driver
+from sqlide.drivers.maven import MavenError
+from sqlide.drivers.registry import DriverRegistry
+
+
+def register(sub: argparse._SubParsersAction) -> None:
+    p = sub.add_parser("driver", help="manage JDBC drivers")
+    ds = p.add_subparsers(dest="driver_cmd", required=True)
+    ds.add_parser("list", help="show known drivers and install state")
+    ins = ds.add_parser("install", help="download a driver from Maven Central")
+    ins.add_argument("id")
+    ins.add_argument("--version", help="default: latest stable")
+    add = ds.add_parser("add", help="add a custom driver (local jars or Maven coordinates)")
+    add.add_argument("id")
+    add.add_argument("--name")
+    add.add_argument("--jar", action="append", default=[], help="local jar (repeatable)")
+    add.add_argument("--maven", help="group:artifact[:classifier]")
+    add.add_argument("--class", dest="class_name", default="", help="default: auto-detect")
+    add.add_argument("--url-template", default="")
+    add.add_argument("--dialect", default="generic")
+    p.set_defaults(handler=run)
+
+
+def run(args: argparse.Namespace) -> int:
+    reg = DriverRegistry()
+    try:
+        return {"list": _list, "install": _install, "add": _add}[args.driver_cmd](reg, args)
+    except (ConfigError, MavenError) as e:
+        print(f"error: {e}", file=sys.stderr)
+        return 1
+
+
+def _list(reg: DriverRegistry, _: argparse.Namespace) -> int:
+    for d in reg.all().values():
+        vs = reg.installed_versions(d.id)
+        state = "local jars" if d.jars else (vs[-1] if vs else "-")
+        print(f"{d.id:<12} {d.name:<24} {state}")
+    return 0
+
+
+def _install(reg: DriverRegistry, args: argparse.Namespace) -> int:
+    def progress(done: int, total: int | None) -> None:
+        size = f"{done / 1e6:.1f}" + (f"/{total / 1e6:.1f}" if total else "") + " MB"
+        print(f"\r  {size}", end="", file=sys.stderr, flush=True)
+
+    jar = reg.install(args.id, version=args.version, progress=progress)
+    print(f"\ninstalled {jar}")
+    return 0
+
+
+def _add(reg: DriverRegistry, args: argparse.Namespace) -> int:
+    reg.add_custom(
+        build_custom_driver(
+            args.id,
+            args.name,
+            maven=args.maven or "",
+            jars=args.jar,
+            class_name=args.class_name,
+            url_template=args.url_template,
+            dialect=args.dialect,
+        )
+    )
+    print(f"driver '{args.id}' saved to {reg.user_file}")
+    if args.maven:
+        print(f"download it with: sqlide driver install {args.id}")
+    return 0