sqlbuddy-nwi 0.3.1__py3-none-any.whl

connections/__init__.py

@@ -0,0 +1,31 @@
+"""Connection management for sql-buddy.
+
+Public surface::
+
+    from connections import DBManager, LockedDBManager, connection_type
+
+    db = LockedDBManager()      # read-only guarded clients (the default)
+    ck = db.ck                  # a Locked ClickHouse client
+    mysql = db.mysql            # a Locked MySQL client
+
+`LockedDBManager` is what the MCP server uses: every client it hands out is
+wrapped by the read-only guard, so write keywords can never reach the network.
+`DBManager` returns raw clients (used by tooling that supplies its own guard).
+
+Implementation lives in focused submodules:
+    readonly  - the SQL read-only filter
+    runtime   - WSL / container / Windows-host-IP detection
+    wrappers  - raw + Locked client wrappers
+    manager   - DBManager / LockedDBManager + connection_type
+"""
+
+from .manager import DBManager, LockedDBManager, connection_type
+from .runtime import is_container, is_wsl
+
+__all__ = [
+    'DBManager',
+    'LockedDBManager',
+    'connection_type',
+    'is_wsl',
+    'is_container',
+]

connections/config_util.py

@@ -0,0 +1,159 @@
+"""Shared config helpers for reading db.ini with cross-platform encoding support.
+
+All modules that touch db.ini resolve its location and parse it through the
+helpers here, never duplicating the path arithmetic.
+
+Config-location resolution
+(see docs/adr/0002-config-discovery-and-invocation-mode-on-publish.md):
+  config_path()        - first existing candidate, for READ
+  config_write_path()  - the detected home, for WRITE (the init wizard)
+  precedence: $SQL_BUDDY_CONFIG > platformdirs user dir > checkout-relative
+"""
+
+import configparser
+import os
+import sys
+
+APP_NAME = "sql-buddy"
+_CONFIG_FILENAME = "db.ini"
+_ENV_OVERRIDE = "SQL_BUDDY_CONFIG"
+
+# config_util.py ships inside the `connections` package, so the checkout-relative
+# config lives at <this dir>/config/db.ini.
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_CHECKOUT_CONFIG_PATH = os.path.join(_HERE, "config", _CONFIG_FILENAME)
+# A checkout is signalled by a `.git` entry at the repo root (a directory for a
+# plain clone, a file for a linked worktree), or by a writable config dir
+# holding the example. Wheels/sdists include neither, so installed copies never
+# match here and correctly fall through to the per-user config dir.
+_REPO_ROOT = os.path.dirname(_HERE)
+
+
+def _env_config_path():
+    """$SQL_BUDDY_CONFIG may name the file directly or a directory holding it."""
+    raw = os.environ.get(_ENV_OVERRIDE)
+    if not raw:
+        return None
+    p = os.path.expanduser(raw)
+    if os.path.isdir(p):
+        return os.path.join(p, _CONFIG_FILENAME)
+    return p
+
+
+def _user_config_path():
+    """Per-user config dir, or None if platformdirs isn't importable.
+
+    Returning None lets config_path()/config_write_path() degrade to the
+    checkout-relative path when the dependency is somehow absent, so a
+    platformdirs install failure can never break an existing checkout user.
+    """
+    try:
+        import platformdirs
+    except ImportError:
+        return None
+    return os.path.join(platformdirs.user_config_dir(APP_NAME), _CONFIG_FILENAME)
+
+
+def is_checkout() -> bool:
+    """Whether sql-buddy is running from a source checkout vs an install.
+
+    Public so the CLI's invocation-mode detection can reuse the same signal
+    (see cli._mcp_entry). See _is_checkout docstring location for the rationale.
+    """
+    if os.path.exists(os.path.join(_REPO_ROOT, ".git")):
+        return True
+    cfg_dir = os.path.join(_HERE, "config")
+    example = os.path.join(cfg_dir, "db.ini.example")
+    return os.path.isfile(example) and os.access(cfg_dir, os.W_OK)
+
+
+def config_write_path() -> str:
+    """Where the init wizard should write db.ini.
+
+    $SQL_BUDDY_CONFIG wins; otherwise a checkout writes checkout-relative (dev
+    back-compat) and an install writes to the per-user config dir.
+    """
+    env = _env_config_path()
+    if env is not None:
+        return env
+    if is_checkout():
+        return _CHECKOUT_CONFIG_PATH
+    return _user_config_path() or _CHECKOUT_CONFIG_PATH
+
+
+def config_path() -> str:
+    """Where to READ db.ini from.
+
+    $SQL_BUDDY_CONFIG is strict (an explicit override). Otherwise the first
+    existing candidate among the per-user dir and the checkout-relative path. If
+    none exist yet, returns config_write_path() so callers can report the place
+    the user should create.
+    """
+    env = _env_config_path()
+    if env is not None:
+        return env
+    for candidate in (_user_config_path(), _CHECKOUT_CONFIG_PATH):
+        if candidate is not None and os.path.isfile(candidate):
+            return candidate
+    return config_write_path()
+
+
+# Historical env override kept for back-compat: db_query used it as the root for
+# ALL its caches (query_history + db_info_cache + db_schema_cache + table_list).
+_CACHE_ENV_OVERRIDE = "DB_QUERY_HISTORY_DIR"
+_CACHE_DIRNAME = "database"
+
+
+def _user_cache_root():
+    """Per-user cache dir, or None if platformdirs isn't importable."""
+    try:
+        import platformdirs
+    except ImportError:
+        return None
+    return platformdirs.user_cache_dir(APP_NAME)
+
+
+def cache_dir_path() -> str:
+    """Where sql-buddy writes runtime caches (query history, schema/info caches).
+
+    $DB_QUERY_HISTORY_DIR (historical override) > per-user cache dir (when not a
+    checkout) > the checkout-relative `database/` dir. Installed copies must
+    never write into the read-only package cache, so a non-checkout resolves to
+    the per-user cache dir; if platformdirs is somehow absent it falls back to
+    the checkout-relative path so a missing dep never crashes the server.
+    """
+    env = os.environ.get(_CACHE_ENV_OVERRIDE)
+    if env:
+        return os.path.expanduser(env)
+    if is_checkout():
+        return os.path.join(_REPO_ROOT, _CACHE_DIRNAME)
+    return _user_cache_root() or os.path.join(_REPO_ROOT, _CACHE_DIRNAME)
+
+
+def resolve_env(value: str) -> str:
+    if isinstance(value, str) and value.startswith('$env:'):
+        var_name = value[5:]
+        resolved = os.environ.get(var_name)
+        if resolved is None:
+            raise EnvironmentError(
+                f'[配置缺失] $env:{var_name} 未在环境中设置。请确认 .env 文件存在（参考 .env.example）或手动设置该环境变量。'
+            )
+        return resolved
+    return value
+
+
+def read_config(config_file: str) -> configparser.RawConfigParser:
+    cfg = configparser.RawConfigParser()
+    with open(config_file, "rb") as f:
+        raw = f.read()
+    encodings = ["utf-8-sig", "utf-8", "gbk"]
+    if sys.platform == "win32":
+        encodings.append("mbcs")
+    for enc in encodings:
+        try:
+            cfg.read_string(raw.decode(enc))
+            return cfg
+        except UnicodeDecodeError:
+            continue
+    cfg.read_string(raw.decode("utf-8", errors="replace"))
+    return cfg

connections/manager.py

@@ -0,0 +1,205 @@
+"""DBManager: builds and caches DB clients from db.ini sections.
+
+`DBManager` returns raw driver clients; `LockedDBManager` returns clients
+wrapped by the read-only guard (the "Locked" layer). Both honor the
+self-managed SSH tunnel (`tunnel.py`): a tunneled connection connects to
+`localhost:local_port`, and the cached client is rebuilt whenever the tunnel's
+generation bumps (the previous TCP socket is dead after a restart).
+"""
+
+import configparser
+import os
+import sys
+import uuid
+from contextlib import contextmanager
+from typing import Any, Dict, Tuple
+
+import pymysql
+from clickhouse_driver import Client
+
+from .config_util import read_config, resolve_env, config_path
+from .runtime import get_windows_host_ip, is_wsl
+from .wrappers import (
+    LockedClickHouseClient,
+    LockedMySQLConnectionWrapper,
+    MySQLConnectionWrapper,
+)
+
+# tunnel.py lives at the repo root; make sure it's importable regardless of cwd.
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _REPO_ROOT not in sys.path:
+    sys.path.insert(0, _REPO_ROOT)
+from tunnel import get_manager as _get_tunnel_manager
+
+
+DEFAULT_REMOTE_DB_CREDS: Dict[str, Tuple[str, str]] = {}
+
+
+def connection_type(cfg, section: str) -> str:
+    """Resolve a connection's DBMS type.
+
+    Reads an explicit `type = clickhouse|mysql` from the section; falls back to
+    inferring from the section-name prefix so older configs keep working.
+    """
+    explicit = cfg.get(section, 'type', fallback=None)
+    if explicit:
+        return explicit.strip().lower()
+    if section.startswith("ck"):
+        return 'clickhouse'
+    if section.startswith("db") or section.startswith("mysql"):
+        return 'mysql'
+    raise ValueError(
+        f"Unknown connection type for [{section}]; set `type = clickhouse|mysql` in db.ini"
+    )
+
+
+class DBManager:
+    """Unlocked DBManager - returns raw connections."""
+
+    def __init__(self, config_file: str = None, remote_db_creds: Dict[str, Tuple[str, str]] = None, wsl_host_override: bool = True):
+        if config_file is None:
+            config_file = config_path()
+        if not os.path.exists(config_file):
+            raise FileNotFoundError(f"Config file not found: {config_file}")
+
+        cfg = read_config(config_file)
+        self.config_file = config_file
+        self._cfg = cfg
+        self._cache = {}
+        self._remote_db_creds = remote_db_creds or DEFAULT_REMOTE_DB_CREDS
+        self._wsl_host_override = wsl_host_override
+
+    @classmethod
+    def get(cls, name: str, config_file: str = None, **kwargs):
+        return cls(config_file, **kwargs).getattr(name)
+
+    @staticmethod
+    def load_config(config_file: str) -> configparser.RawConfigParser:
+        if config_file is None:
+            config_file = config_path()
+        if not os.path.exists(config_file):
+            raise FileNotFoundError(f"Config file not found: {config_file}")
+        return read_config(config_file)
+
+    def __getattr__(self, name: str):
+        section = name.replace("_", "-")
+        if section not in self._cfg:
+            raise AttributeError(f"No section [{section}] in config")
+
+        tm = _get_tunnel_manager()
+        block = tm.block_for(section)
+        if block is not None:
+            # Bring up (or confirm) the self-managed SSH tunnel before connecting.
+            tm.ensure_for(section)
+        # Generation changes when the tunnel was (re)started -> the cached client's
+        # TCP socket is dead and must be rebuilt. 0 for non-tunneled connections.
+        gen = tm.generation(block) if block is not None else 0
+
+        cached = self._cache.get(name)
+        if cached is not None and cached['gen'] == gen:
+            return cached['client']
+
+        ctype = connection_type(self._cfg, section)
+        if ctype == 'clickhouse':
+            client = self._create_clickhouse_client(section)
+        elif ctype in ('mysql', 'mariadb'):
+            client = self._create_mysql_client(section)
+        else:
+            raise ValueError(f"Unsupported connection type '{ctype}' for [{section}]")
+
+        self._cache[name] = {'client': client, 'gen': gen}
+        return client
+
+    def _resolve_conn_config(self, section: str, default_port: int) -> dict:
+        """Build driver kwargs for a section: strip meta keys, resolve host/port.
+
+        Self-tunneled connections connect to localhost:local_port (the SSH forward).
+        Otherwise the legacy WSL host-override applies when running inside WSL.
+        """
+        conn_config = dict(self._cfg[section])
+        # Meta keys consumed elsewhere, not by the DB driver.
+        for meta in ('type', 'tunnel', 'remote_host', 'remote_port', 'local_port'):
+            conn_config.pop(meta, None)
+        for k, v in list(conn_config.items()):
+            conn_config[k] = resolve_env(str(v))
+
+        if _get_tunnel_manager().block_for(section) is not None:
+            conn_config['host'] = 'localhost'
+            conn_config['port'] = int(_get_tunnel_manager().local_port_for(section))
+        else:
+            if self._wsl_host_override and is_wsl():
+                windows_ip = get_windows_host_ip()
+                if windows_ip:
+                    conn_config['host'] = windows_ip
+            conn_config['port'] = int(conn_config.get('port', default_port))
+        return conn_config
+
+    def _create_clickhouse_client(self, section: str) -> Client:
+        return Client(**self._resolve_conn_config(section, 9000))
+
+    def _create_mysql_client(self, section: str) -> MySQLConnectionWrapper:
+        return MySQLConnectionWrapper(pymysql.connect(**self._resolve_conn_config(section, 3306)))
+
+    def all_clients(self) -> Dict[str, Any]:
+        return {name: entry['client'] for name, entry in self._cache.items()}
+
+    def close_all(self):
+        for entry in self._cache.values():
+            client = entry['client']
+            if hasattr(client, "disconnect"):
+                client.disconnect()
+            elif hasattr(client, "close"):
+                client.close()
+
+    @contextmanager
+    def session_scope(self, client_name: str = 'ck'):
+        client = getattr(self, client_name)
+        session_id = str(uuid.uuid4())
+        original_execute = client.execute
+        original_query_dataframe = getattr(client, 'query_dataframe', None)
+
+        def scoped_execute(query, params=None, settings=None, **kwargs):
+            if settings is None:
+                settings = {}
+            settings['session_id'] = session_id
+            return original_execute(query, params=params, settings=settings, **kwargs)
+
+        def scoped_query_dataframe(query, params=None, settings=None, **kwargs):
+            if settings is None:
+                settings = {}
+            settings['session_id'] = session_id
+            return original_query_dataframe(query, params=params, settings=settings, **kwargs)
+
+        class SessionClient:
+            def execute(self, *args, **kwargs):
+                return scoped_execute(*args, **kwargs)
+            def query_dataframe(self, *args, **kwargs):
+                if original_query_dataframe:
+                    return scoped_query_dataframe(*args, **kwargs)
+                raise AttributeError("query_dataframe not available on this client")
+
+        yield SessionClient()
+
+    @contextmanager
+    def transaction_scope(self, connection: str):
+        conn = getattr(self, connection)
+        cursor = conn.cursor()
+        cursor.execute("SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED")
+        cursor.execute("SET SESSION innodb_lock_wait_timeout = 600")
+        cursor.close()
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+
+
+class LockedDBManager(DBManager):
+    """Returns read-only wrapped connections."""
+
+    def _create_clickhouse_client(self, section: str) -> LockedClickHouseClient:
+        return LockedClickHouseClient(super()._create_clickhouse_client(section))
+
+    def _create_mysql_client(self, section: str) -> LockedMySQLConnectionWrapper:
+        return LockedMySQLConnectionWrapper(super()._create_mysql_client(section))

connections/readonly.py

@@ -0,0 +1,45 @@
+"""Read-only SQL guard.
+
+The enforcement layer that makes every connection read-only by construction:
+write/mutating keywords are rejected here, before any SQL reaches the network.
+"""
+
+import re
+
+# SQL verbs that mutate state. Matched as whole words (case-insensitive).
+WRITE_KEYWORDS = [
+    'INSERT', 'UPDATE', 'DELETE', 'DROP', 'CREATE', 'ALTER',
+    'TRUNCATE', 'REPLACE', 'MERGE', 'CALL', 'EXEC', 'EXECUTE',
+    'GRANT', 'REVOKE', 'LOCK', 'UNLOCK', 'SET GLOBAL',
+    'KILL', 'LOAD DATA', 'START TRANSACTION', 'BEGIN'
+]
+
+# Verbs a read-only query may begin with. Anything else is treated as a write.
+READ_ONLY_KEYWORDS = [
+    'SELECT', 'SHOW', 'DESCRIBE', 'DESC', 'EXPLAIN',
+    'USE', 'WITH', 'EXISTS', 'CHECK', 'HELP'
+]
+
+
+def _is_write_operation(sql: str) -> bool:
+    sql_upper = sql.strip().upper()
+    for keyword in READ_ONLY_KEYWORDS:
+        if sql_upper.startswith(keyword):
+            return False
+    for keyword in WRITE_KEYWORDS:
+        pattern = r'\b' + keyword + r'\b'
+        if re.search(pattern, sql_upper):
+            if keyword == 'CREATE' and 'SHOW' in sql_upper:
+                continue
+            return True
+    return False
+
+
+def _filter_sql(sql: str) -> str:
+    """Pass through `sql` unchanged, or raise PermissionError if it mutates."""
+    if _is_write_operation(sql):
+        raise PermissionError(
+            "Write operation blocked by sql-buddy's read-only guard. "
+            f"Only READ queries are allowed. Blocked SQL: {sql[:100]}..."
+        )
+    return sql

connections/runtime.py

@@ -0,0 +1,51 @@
+"""Runtime environment detection.
+
+Decides whether a connection's `host` should be rewritten to the Windows host
+IP: when sql-buddy runs inside WSL but the user's tunnel is on the Windows side,
+the driver must connect to the Windows host rather than `localhost`.
+"""
+
+import os
+from typing import Optional
+
+
+def is_wsl() -> bool:
+    if is_container():
+        return False
+    try:
+        with open('/proc/version', 'r') as f:
+            version = f.read().lower()
+            return 'microsoft' in version or 'wsl' in version
+    except (FileNotFoundError, IOError):
+        return False
+
+
+def is_container() -> bool:
+    if os.path.exists('/.dockerenv'):
+        return True
+    try:
+        with open('/proc/1/cgroup', 'r') as f:
+            cgroup = f.read()
+            if any(sig in cgroup for sig in ['docker', 'containerd', 'kubepods', 'lxc']):
+                return True
+    except (FileNotFoundError, IOError):
+        pass
+    for var in ['DOCKER_CONTAINER', 'DOCKER_ID', 'CONTAINER', 'KUBERNETES_SERVICE_HOST']:
+        if os.environ.get(var):
+            return True
+    return False
+
+
+def get_windows_host_ip() -> Optional[str]:
+    if not is_wsl():
+        return None
+    try:
+        with open('/etc/resolv.conf', 'r') as f:
+            for line in f:
+                if line.strip().startswith('nameserver'):
+                    ip = line.split()[1]
+                    if ip and ip != '127.0.0.1':
+                        return ip
+    except (FileNotFoundError, IOError):
+        pass
+    return None

connections/wrappers.py

@@ -0,0 +1,105 @@
+"""Connection wrappers: raw and read-only-guarded DB clients.
+
+`MySQLConnectionWrapper` adapts a raw pymysql connection to the dataframe/
+execute interface the rest of sql-buddy expects. The `Locked*` variants wrap a
+client and route every statement through the read-only guard (see `readonly`)
+so write keywords can never reach the network.
+
+See UBIQUITOUS_LANGUAGE.md: "Locked" == the wrapper layer enforcing Read-only.
+"""
+
+import pymysql
+from clickhouse_driver import Client
+import pandas as pd
+
+from .readonly import _filter_sql
+
+
+class MySQLConnectionWrapper:
+    def __init__(self, conn):
+        self._conn = conn
+
+    def __getattr__(self, name):
+        return getattr(self._conn, name)
+
+    def db_execute(self, sql: str, params: tuple = None) -> int:
+        cursor = self._conn.cursor()
+        try:
+            cursor.execute(sql, params)
+            self._conn.commit()
+            return cursor.rowcount
+        finally:
+            cursor.close()
+
+    def query_dataframe(self, sql: str) -> pd.DataFrame:
+        cursor = self._conn.cursor(pymysql.cursors.DictCursor)
+        try:
+            cursor.execute(sql)
+            rows = cursor.fetchall()
+            return pd.DataFrame(rows)
+        finally:
+            cursor.close()
+
+
+class LockedMySQLConnectionWrapper:
+    def __init__(self, wrapped: MySQLConnectionWrapper):
+        self._wrapped = wrapped
+
+    def __getattr__(self, name):
+        return getattr(self._wrapped, name)
+
+    def cursor(self, *args, **kwargs):
+        return self._wrapped.cursor(*args, **kwargs)
+
+    def commit(self):
+        return self._wrapped.commit()
+
+    def rollback(self):
+        return self._wrapped.rollback()
+
+    def close(self):
+        return self._wrapped.close()
+
+    @property
+    def open(self):
+        return self._wrapped.open
+
+    def db_execute(self, sql: str, params: tuple = None) -> int:
+        _filter_sql(sql)
+        return self._wrapped.db_execute(sql, params)
+
+    def query_dataframe(self, sql: str) -> pd.DataFrame:
+        _filter_sql(sql)
+        return self._wrapped.query_dataframe(sql)
+
+
+class LockedClickHouseClient:
+    def __init__(self, wrapped: Client):
+        self._wrapped = wrapped
+
+    def __getattr__(self, name):
+        return getattr(self._wrapped, name)
+
+    def execute(self, query, params=None, with_column_types=False, settings=None, **kwargs):
+        _filter_sql(query)
+        return self._wrapped.execute(
+            query, params=params,
+            with_column_types=with_column_types,
+            settings=settings, **kwargs
+        )
+
+    def query_dataframe(self, query, params=None, settings=None, **kwargs):
+        _filter_sql(query)
+        result = self._wrapped.execute(query, params=params, settings=settings, with_column_types=True, **kwargs)
+        if result:
+            data, columns = result
+            if data:
+                return pd.DataFrame(data, columns=[c[0] for c in columns])
+        return pd.DataFrame()
+
+    def disconnect(self):
+        return self._wrapped.disconnect()
+
+    @property
+    def connected(self):
+        return self._wrapped.connected