replypier/backend/app/models/schemas/ticket_db_connector.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.
"""

import re
from datetime import datetime
from typing import Dict, List, Optional
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field, field_validator

ENGINES = ("postgresql", "mysql")
ROW_SCOPE_KEYS = ("email", "phone")


class DBConnectorColumnOut(BaseModel):
    name: str
    type: str


class DBConnectorTableOut(BaseModel):
    schema_name: str = Field(alias="schema")
    table: str
    columns: List[DBConnectorColumnOut]

    model_config = ConfigDict(populate_by_name=True)


class SSHTunnelFields(BaseModel):
    """Optional SSH bastion/jump-host tunnel. Auth is by private key
    (preferred) or password; secrets are write-only."""
    ssh_enabled: bool = False
    ssh_host: Optional[str] = Field(default=None, max_length=500)
    ssh_port: int = Field(default=22, ge=1, le=65535)
    ssh_username: Optional[str] = Field(default=None, max_length=200)
    ssh_password: Optional[str] = Field(default=None, max_length=1000)
    ssh_private_key: Optional[str] = Field(default=None, max_length=20000)
    ssh_private_key_passphrase: Optional[str] = Field(default=None, max_length=1000)


class DBConnectorDiscoverRequest(SSHTunnelFields):
    """Pre-save connection test: credentials go over TLS, are used once for
    discovery, and are never persisted by this endpoint."""
    engine: str = Field(pattern=f"^({'|'.join(ENGINES)})$")
    host: str = Field(min_length=1, max_length=500)
    port: int = Field(ge=1, le=65535)
    database: str = Field(min_length=1, max_length=200)
    username: str = Field(min_length=1, max_length=200)
    password: str = Field(min_length=1, max_length=500)


class DBConnectorDiscoverResponse(BaseModel):
    ok: bool
    error: Optional[str] = None
    tables: List[DBConnectorTableOut] = []


_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
_TABLE_REF = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*(\.[A-Za-z_][A-Za-z0-9_$]*)?$")

ROW_SCOPE_DESCRIPTION = (
    'Tables restricted to the ticket customer\'s own rows, as '
    '{"schema.table": "customer_column"}. Listed tables are rewritten into a '
    "pre-filtered subquery; tables left out are readable in full."
)


def validate_row_scope(value: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
    """Both halves land in generated SQL as identifiers. sqlglot quotes them,
    but a bad value should be refused at the edge rather than quoted and
    executed — a typo'd column silently scopes nothing useful."""
    if value is None:
        return None
    if len(value) > 500:
        raise ValueError("Too many row-scope entries")
    cleaned: Dict[str, str] = {}
    for table, column in value.items():
        table = (table or "").strip()
        column = (column or "").strip()
        if not table or not column:
            continue
        if not _TABLE_REF.match(table):
            raise ValueError(f"Invalid table reference in row scope: {table!r}")
        if not _IDENTIFIER.match(column):
            raise ValueError(f"Invalid column name in row scope: {column!r}")
        cleaned[table.lower()] = column
    return cleaned


class DBConnectorCreate(SSHTunnelFields):
    name: str = Field(min_length=1, max_length=200)
    engine: str = Field(pattern=f"^({'|'.join(ENGINES)})$")
    host: str = Field(min_length=1, max_length=500)
    port: int = Field(ge=1, le=65535)
    database: str = Field(min_length=1, max_length=200)
    username: str = Field(min_length=1, max_length=200)
    password: str = Field(min_length=1, max_length=500)
    enabled: bool = True
    # Nothing is queryable unless explicitly listed ("schema.table").
    allowed_tables: List[str] = Field(default_factory=list, max_length=500)
    masked_columns: List[str] = Field(default_factory=list, max_length=500)
    row_scope: Dict[str, str] = Field(
        default_factory=dict, description=ROW_SCOPE_DESCRIPTION
    )
    row_scope_key: str = Field(
        default="email", pattern=f"^({'|'.join(ROW_SCOPE_KEYS)})$"
    )
    max_rows: int = Field(default=100, ge=1, le=1000)
    statement_timeout_ms: int = Field(default=5000, ge=100, le=30000)

    @field_validator("row_scope")
    @classmethod
    def _check_row_scope(cls, value):
        return validate_row_scope(value) or {}


class DBConnectorUpdate(BaseModel):
    name: Optional[str] = Field(default=None, min_length=1, max_length=200)
    host: Optional[str] = Field(default=None, min_length=1, max_length=500)
    port: Optional[int] = Field(default=None, ge=1, le=65535)
    database: Optional[str] = Field(default=None, min_length=1, max_length=200)
    username: Optional[str] = Field(default=None, min_length=1, max_length=200)
    # Omit to keep the stored password.
    password: Optional[str] = Field(default=None, min_length=1, max_length=500)
    enabled: Optional[bool] = None
    allowed_tables: Optional[List[str]] = Field(default=None, max_length=500)
    masked_columns: Optional[List[str]] = Field(default=None, max_length=500)
    row_scope: Optional[Dict[str, str]] = Field(
        default=None, description=ROW_SCOPE_DESCRIPTION
    )
    row_scope_key: Optional[str] = Field(
        default=None, pattern=f"^({'|'.join(ROW_SCOPE_KEYS)})$"
    )
    max_rows: Optional[int] = Field(default=None, ge=1, le=1000)
    statement_timeout_ms: Optional[int] = Field(default=None, ge=100, le=30000)
    # SSH tunnel — omit a secret to keep the stored one.
    ssh_enabled: Optional[bool] = None
    ssh_host: Optional[str] = Field(default=None, max_length=500)
    ssh_port: Optional[int] = Field(default=None, ge=1, le=65535)
    ssh_username: Optional[str] = Field(default=None, max_length=200)
    ssh_password: Optional[str] = Field(default=None, max_length=1000)
    ssh_private_key: Optional[str] = Field(default=None, max_length=20000)
    ssh_private_key_passphrase: Optional[str] = Field(default=None, max_length=1000)

    @field_validator("row_scope")
    @classmethod
    def _check_row_scope(cls, value):
        return validate_row_scope(value)


class DBConnectorOut(BaseModel):
    """Never includes the password in any form."""
    model_config = ConfigDict(from_attributes=True)
    id: UUID
    name: str
    engine: str
    host: str
    port: int
    database: str
    username: str
    enabled: bool
    allowed_tables: Optional[List[str]] = None
    masked_columns: Optional[List[str]] = None
    row_scope: Optional[Dict[str, str]] = None
    row_scope_key: str = "email"
    max_rows: int
    statement_timeout_ms: int
    # SSH tunnel — connection details only, never the key/password.
    ssh_enabled: bool = False
    ssh_host: Optional[str] = None
    ssh_port: int = 22
    ssh_username: Optional[str] = None
    last_test_at: Optional[datetime] = None
    last_test_ok: Optional[bool] = None
    created_at: Optional[datetime] = None