File size: 7,378 Bytes
2ba0613
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
"""Pydantic credential schemas for user-registered external databases.

Imported by the `/database-clients` API router (`src/api/v1/db_client.py`) and,
via `DbType`, by the db pipeline connector (`src/pipeline/db_pipeline/connector.py`).

Sensitive fields (`password`, `service_account_json`) are Fernet-encrypted by
the database_client service before being stored in the JSONB column; these
schemas describe the plaintext wire format, not the stored shape.
"""

from typing import Literal, Optional, Union

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Supported DB types
# ---------------------------------------------------------------------------

DbType = Literal["postgres", "mysql", "sqlserver", "supabase", "bigquery", "snowflake"]


# ---------------------------------------------------------------------------
# Typed credential schemas per DB type
# ---------------------------------------------------------------------------


class PostgresCredentials(BaseModel):
    """Connection credentials for PostgreSQL."""

    host: str = Field(..., description="Hostname or IP address of the PostgreSQL server.", examples=["db.example.com"])
    port: int = Field(5432, description="Port number (default: 5432).", examples=[5432])
    database: str = Field(..., description="Name of the target database.", examples=["mydb"])
    username: str = Field(..., description="Database username.", examples=["admin"])
    password: str = Field(..., description="Database password. Will be encrypted at rest.", examples=["s3cr3t!"])
    ssl_mode: Literal["disable", "require", "verify-ca", "verify-full"] = Field(
        "require",
        description="SSL mode for the connection.",
        examples=["require"],
    )


class MysqlCredentials(BaseModel):
    """Connection credentials for MySQL."""

    host: str = Field(..., description="Hostname or IP address of the MySQL server.", examples=["db.example.com"])
    port: int = Field(3306, description="Port number (default: 3306).", examples=[3306])
    database: str = Field(..., description="Name of the target database.", examples=["mydb"])
    username: str = Field(..., description="Database username.", examples=["admin"])
    password: str = Field(..., description="Database password. Will be encrypted at rest.", examples=["s3cr3t!"])
    ssl: bool = Field(True, description="Enable SSL for the connection.", examples=[True])


class SqlServerCredentials(BaseModel):
    """Connection credentials for Microsoft SQL Server."""

    host: str = Field(..., description="Hostname or IP address of the SQL Server.", examples=["sqlserver.example.com"])
    port: int = Field(1433, description="Port number (default: 1433).", examples=[1433])
    database: str = Field(..., description="Name of the target database.", examples=["mydb"])
    username: str = Field(..., description="Database username.", examples=["sa"])
    password: str = Field(..., description="Database password. Will be encrypted at rest.", examples=["s3cr3t!"])
    driver: Optional[str] = Field(
        None,
        description="ODBC driver name. Leave empty to use the default driver.",
        examples=["ODBC Driver 17 for SQL Server"],
    )


class SupabaseCredentials(BaseModel):
    """Connection credentials for Supabase (PostgreSQL-based).

    Use the connection string details from your Supabase project dashboard
    under Settings > Database.
    """

    host: str = Field(
        ...,
        description="Supabase database host (e.g. db.<project-ref>.supabase.co, or the pooler host).",
        examples=["db.xxxx.supabase.co"],
    )
    port: int = Field(
        5432,
        description="Port number. Use 5432 for direct connection, 6543 for the connection pooler.",
        examples=[5432],
    )
    database: str = Field("postgres", description="Database name (always 'postgres' for Supabase).", examples=["postgres"])
    username: str = Field(
        ...,
        description="Database user. Use 'postgres' for direct connection, or 'postgres.<project-ref>' for the pooler.",
        examples=["postgres"],
    )
    password: str = Field(..., description="Database password (set in Supabase dashboard). Will be encrypted at rest.", examples=["s3cr3t!"])
    ssl_mode: Literal["require", "verify-ca", "verify-full"] = Field(
        "require",
        description="SSL mode. Supabase always requires SSL.",
        examples=["require"],
    )


class BigQueryCredentials(BaseModel):
    """Connection credentials for Google BigQuery.

    Requires a GCP Service Account with at least BigQuery Data Viewer
    and BigQuery Job User roles.
    """

    project_id: str = Field(..., description="GCP project ID where the BigQuery dataset resides.", examples=["my-gcp-project"])
    dataset_id: str = Field(..., description="BigQuery dataset name to connect to.", examples=["my_dataset"])
    location: Optional[str] = Field(
        "US",
        description="Dataset location/region (default: US).",
        examples=["US", "EU", "asia-southeast1"],
    )
    service_account_json: str = Field(
        ...,
        description=(
            "Full content of the GCP Service Account key JSON file as a string. "
            "Will be encrypted at rest."
        ),
        examples=['{"type":"service_account","project_id":"my-gcp-project","private_key_id":"..."}'],
    )


class SnowflakeCredentials(BaseModel):
    """Connection credentials for Snowflake."""

    account: str = Field(
        ...,
        description="Snowflake account identifier, including region if applicable (e.g. myaccount.us-east-1).",
        examples=["myaccount.us-east-1"],
    )
    warehouse: str = Field(..., description="Name of the virtual warehouse to use for queries.", examples=["COMPUTE_WH"])
    database: str = Field(..., description="Name of the target Snowflake database.", examples=["MY_DB"])
    db_schema: Optional[str] = Field("PUBLIC", alias="schema", description="Schema name (default: PUBLIC).", examples=["PUBLIC"])
    username: str = Field(..., description="Snowflake username.", examples=["admin"])
    password: str = Field(..., description="Snowflake password. Will be encrypted at rest.", examples=["s3cr3t!"])
    role: Optional[str] = Field(None, description="Snowflake role to assume for the session.", examples=["SYSADMIN"])


# Union of all credential shapes — reserved for future typed validation on
# DatabaseClientCreate.credentials (currently Dict[str, Any]). Kept exported
# so downstream code can reference it without re-declaring.
CredentialsUnion = Union[
    PostgresCredentials,
    MysqlCredentials,
    SqlServerCredentials,
    SupabaseCredentials,
    BigQueryCredentials,
    SnowflakeCredentials,
]


# Doc-only helper: surfaces per-type credential shapes in the Swagger "Schemas"
# panel so API consumers can discover the exact field set for each db_type.
# Not referenced by any endpoint — importing it in db_client.py is enough for
# FastAPI's OpenAPI generator to pick it up.
class CredentialSchemas(BaseModel):
    """Reference schemas for `credentials` per `db_type` (Swagger-only, not used by endpoints)."""

    postgres: PostgresCredentials
    mysql: MysqlCredentials
    sqlserver: SqlServerCredentials
    supabase: SupabaseCredentials
    bigquery: BigQueryCredentials
    snowflake: SnowflakeCredentials