replypier/backend/app/services/db_connector_service.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Execution layer for guardrailed ticket DB connectors. All functions here are
synchronous (drivers are sync) — callers run them via asyncio.to_thread.

The read-only guarantee is enforced at the SESSION level, independently of
the SQL guardrails: default_transaction_read_only for Postgres, SESSION
TRANSACTION READ ONLY for MySQL, plus statement timeouts on both.
"""

import io
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from uuid import UUID

from app.core.logger import get_logger
from app.core.outbound_policy import pin_outbound_target
from app.core.security import decrypt_api_key

logger = get_logger(__name__)

CONNECT_TIMEOUT_SECONDS = 8
SSH_CONNECT_TIMEOUT_SECONDS = 10
# Schemas that are never discovered or queried.
SYSTEM_SCHEMAS = (
    "pg_catalog", "information_schema", "pg_toast",
    "mysql", "performance_schema", "sys",
)


@dataclass
class DBConnectorConfig:
    """Plain snapshot of a TicketDBConnector row (secrets decrypted) so the
    executor never holds ORM objects across sessions/threads."""
    id: Optional[UUID]
    organization_id: Optional[UUID]
    name: str
    engine: str
    host: str
    port: int
    database: str
    username: str
    password: str
    allowed_tables: List[str] = field(default_factory=list)
    masked_columns: List[str] = field(default_factory=list)
    # {"table": "customer_column"} — tables restricted to the ticket
    # customer's own rows. See sql_guardrails._apply_row_scope.
    row_scope: dict = field(default_factory=dict)
    row_scope_key: str = "email"
    max_rows: int = 100
    statement_timeout_ms: int = 5000
    # SSH tunnel (bastion/jump host) — the common production topology where
    # the database isn't directly reachable.
    ssh_enabled: bool = False
    ssh_host: Optional[str] = None
    ssh_port: int = 22
    ssh_username: Optional[str] = None
    ssh_password: Optional[str] = None
    ssh_private_key: Optional[str] = None
    ssh_private_key_passphrase: Optional[str] = None

    @classmethod
    def from_model(cls, connector) -> "DBConnectorConfig":
        return cls(
            id=connector.id,
            organization_id=connector.organization_id,
            name=connector.name,
            engine=str(connector.engine),
            host=connector.host,
            port=connector.port,
            database=connector.database,
            username=connector.username,
            password=decrypt_api_key(connector.encrypted_password),
            allowed_tables=list(connector.allowed_tables or []),
            masked_columns=list(connector.masked_columns or []),
            row_scope=dict(connector.row_scope or {}),
            row_scope_key=connector.row_scope_key or "email",
            max_rows=connector.max_rows or 100,
            statement_timeout_ms=connector.statement_timeout_ms or 5000,
            ssh_enabled=bool(connector.ssh_enabled),
            ssh_host=connector.ssh_host,
            ssh_port=connector.ssh_port or 22,
            ssh_username=connector.ssh_username,
            ssh_password=(
                decrypt_api_key(connector.encrypted_ssh_password)
                if connector.encrypted_ssh_password else None
            ),
            ssh_private_key=(
                decrypt_api_key(connector.encrypted_ssh_private_key)
                if connector.encrypted_ssh_private_key else None
            ),
            ssh_private_key_passphrase=(
                decrypt_api_key(connector.encrypted_ssh_key_passphrase)
                if connector.encrypted_ssh_key_passphrase else None
            ),
        )


def _load_ssh_key(key_text: str, passphrase: Optional[str]):
    """Parse a PEM/OpenSSH private key of any common type from a string,
    without writing it to disk."""
    import paramiko

    last_error = None
    # DSA is obsolete — Ed25519/RSA/ECDSA cover every key type a modern host
    # issues.
    for key_class in (paramiko.Ed25519Key, paramiko.RSAKey, paramiko.ECDSAKey):
        try:
            return key_class.from_private_key(
                io.StringIO(key_text), password=passphrase or None
            )
        except Exception as e:  # wrong type or bad passphrase — try the next
            last_error = e
    raise ValueError(f"Unsupported or unreadable SSH private key: {last_error}")


@contextmanager
def _ssh_tunnel(config: DBConnectorConfig, ssh_address: Optional[str] = None):
    """Open an SSH local-forward tunnel to the database and yield the local
    (host, port) to connect through. ``ssh_address`` is the bastion's vetted
    IP when the outbound policy pinned one. Always torn down on exit."""
    from sshtunnel import SSHTunnelForwarder

    kwargs = {
        "ssh_username": config.ssh_username,
        "remote_bind_address": (config.host, config.port),
        # 127.0.0.1:0 — the OS assigns a free local port per tunnel.
        "local_bind_address": ("127.0.0.1", 0),
        "set_keepalive": 15.0,
    }
    if config.ssh_private_key:
        kwargs["ssh_pkey"] = _load_ssh_key(
            config.ssh_private_key, config.ssh_private_key_passphrase
        )
    elif config.ssh_password:
        kwargs["ssh_password"] = config.ssh_password
    else:
        raise ValueError("SSH tunnel needs a private key or password")

    tunnel = SSHTunnelForwarder((ssh_address or config.ssh_host, config.ssh_port), **kwargs)
    tunnel.start()
    try:
        yield "127.0.0.1", tunnel.local_bind_port
    finally:
        try:
            tunnel.stop()
        except Exception as e:
            logger.warning(f"SSH tunnel teardown warning (non-critical): {e}")


def _driver_connect(config: DBConnectorConfig, host: str, port: int,
                    hostaddr: Optional[str] = None):
    """Open a read-only DB session with a statement timeout, to (host, port)
    — which is the tunnel's local endpoint when tunneling. ``hostaddr`` is
    the vetted IP to dial instead of resolving ``host`` again."""
    if config.engine == "mysql":
        import pymysql
        conn = pymysql.connect(
            # pymysql has no separate address field; no TLS is configured
            # here, so the name isn't needed for verification.
            host=hostaddr or host,
            port=port,
            user=config.username,
            password=config.password,
            database=config.database,
            connect_timeout=CONNECT_TIMEOUT_SECONDS,
            read_timeout=max(1, config.statement_timeout_ms // 1000 + 1),
        )
        with conn.cursor() as cursor:
            cursor.execute("SET SESSION TRANSACTION READ ONLY")
            # max_execution_time bounds SELECTs server-side (milliseconds).
            cursor.execute(f"SET SESSION max_execution_time = {int(config.statement_timeout_ms)}")
        return conn
    import psycopg
    # libpq dials hostaddr and keeps host for TLS verification.
    pinned = {"hostaddr": hostaddr} if hostaddr else {}
    return psycopg.connect(
        host=host,
        **pinned,
        port=port,
        dbname=config.database,
        user=config.username,
        password=config.password,
        connect_timeout=CONNECT_TIMEOUT_SECONDS,
        options=(
            f"-c statement_timeout={int(config.statement_timeout_ms)} "
            "-c default_transaction_read_only=on"
        ),
    )


def first_hop(ssh_enabled: bool, ssh_host: Optional[str], host: Optional[str]) -> Tuple[Optional[str], str]:
    """The host this server dials first, with its field label, for the
    outbound policy: the SSH bastion when tunnelling, else the database. A
    database behind a bastion is reached from the bastion, and may be private."""
    if ssh_enabled:
        return ssh_host, "SSH host"
    return host, "Database host"


@contextmanager
def _connection(config: DBConnectorConfig):
    """Yield a read-only DB connection, opening an SSH tunnel first when the
    connector is configured for one. Tunnel and connection are both closed on
    exit — a fresh pair per call keeps the (short-lived) investigation queries
    isolated."""
    # Vetted on every connect, not only on save, and the vetted address is
    # the one dialled: DNS can change, even between two lookups.
    pinned = pin_outbound_target(*first_hop(config.ssh_enabled, config.ssh_host, config.host))
    if config.ssh_enabled:
        if not config.ssh_host or not config.ssh_username:
            raise ValueError("SSH tunnel is enabled but host/username are missing")
        with _ssh_tunnel(config, pinned) as (host, port):
            conn = _driver_connect(config, host, port)
            try:
                yield conn
            finally:
                try:
                    conn.close()
                except Exception:
                    pass
    else:
        conn = _driver_connect(config, config.host, config.port, pinned)
        try:
            yield conn
        finally:
            try:
                conn.close()
            except Exception:
                pass


def run_readonly_query(
    config: DBConnectorConfig, sql: str
) -> Tuple[List[str], List[tuple]]:
    """Execute validated SQL in a read-only session; fetches at most
    max_rows + 1 (the validator already forced a LIMIT — this is belt and
    braces)."""
    with _connection(config) as conn:
        with conn.cursor() as cursor:
            cursor.execute(sql)
            columns = [d[0] for d in (cursor.description or [])]
            rows = cursor.fetchmany(config.max_rows + 1)
        return columns, list(rows)


def test_connection(config: DBConnectorConfig) -> Dict:
    """Connect and discover user tables + columns (for the allowlist picker).
    Returns {"tables": [{"schema", "table", "columns": [{"name", "type"}]}]}."""
    placeholders = ", ".join(["%s"] * len(SYSTEM_SCHEMAS))
    sql = (
        "SELECT table_schema, table_name, column_name, data_type "
        "FROM information_schema.columns "
        f"WHERE table_schema NOT IN ({placeholders}) "
        "ORDER BY table_schema, table_name, ordinal_position"
    )
    with _connection(config) as conn:
        with conn.cursor() as cursor:
            cursor.execute(sql, SYSTEM_SCHEMAS)
            rows = cursor.fetchall()

    tables: Dict[Tuple[str, str], List[Dict]] = {}
    for schema, table, column, data_type in rows:
        tables.setdefault((schema, table), []).append({"name": column, "type": data_type})
    return {
        "tables": [
            {"schema": schema, "table": table, "columns": columns}
            for (schema, table), columns in tables.items()
        ]
    }


def describe_columns(config: DBConnectorConfig, schema: str, table: str) -> List[Dict]:
    """Column names/types of one table via information_schema."""
    with _connection(config) as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT column_name, data_type FROM information_schema.columns "
                "WHERE table_schema = %s AND table_name = %s ORDER BY ordinal_position",
                (schema, table),
            )
            return [{"name": name, "type": data_type} for name, data_type in cursor.fetchall()]


def format_result_table(
    columns: List[str], rows: List[tuple], max_rows: int, max_value_chars: int = 200
) -> str:
    """Compact text rendering of a result set for the model."""
    if not columns:
        return "(no result)"
    shown = rows[:max_rows]
    lines = [" | ".join(columns)]
    for row in shown:
        lines.append(
            " | ".join(
                (str(value)[:max_value_chars] if value is not None else "NULL")
                for value in row
            )
        )
    suffix = f"({len(shown)} row{'s' if len(shown) != 1 else ''}"
    if len(rows) > max_rows:
        suffix += ", truncated"
    suffix += ")"
    lines.append(suffix)
    return "\n".join(lines)