From aa18118734222581f5b7db5a88ec7f96ca9c5d46 Mon Sep 17 00:00:00 2001 From: edwardvaneechoud Date: Wed, 5 Aug 2026 14:25:05 +0200 Subject: [PATCH 1/4] add support for snowflake --- build_backends/build_backends/main.py | 3 + .../integrations/database_read_snowflake.py | 40 ++ docs/index.html | 2 +- docs/users/connect/index.md | 4 +- docs/users/data-elsewhere.md | 4 +- .../python-api/reference/reading-data.md | 30 ++ .../python-api/reference/writing-data.md | 6 + docs/users/visual-editor/connections.md | 11 +- docs/what-is-flowfile-technical.md | 2 +- .../030_database_connection_extra_params.py | 37 ++ .../flowfile_core/database/models.py | 1 + .../db_connections.py | 21 ++ .../flowfile_core/flowfile/flow_graph.py | 3 + .../external_sources/sql_source/sql_source.py | 1 + .../flowfile_core/project/importer.py | 1 + .../flowfile_core/project/manifest_entries.py | 1 + .../flowfile_core/project/projection.py | 3 + flowfile_core/flowfile_core/routes/routes.py | 5 + .../flowfile_core/schemas/input_schema.py | 25 ++ .../tests/docs_examples/test_docs_examples.py | 8 + .../test_dialect_vocabulary.py | 16 + .../external_sources/test_snowflake_source.py | 190 ++++++++++ flowfile_core/tests/project/test_roundtrip.py | 47 +++ .../tests/test_db_dialect_endpoints.py | 45 +++ .../database/connection_manager.py | 14 +- .../database/connection_manager.pyi | 4 +- .../src/renderer/app/api/dbDialects.test.ts | 13 + .../src/renderer/app/api/dbDialects.ts | 34 ++ .../DatabaseConnectionSettings.vue | 33 +- .../renderer/app/composables/useDbDialects.ts | 15 +- .../src/renderer/app/types/node.types.ts | 1 + .../DatabaseConnectionSettings.vue | 46 ++- .../app/views/DatabaseView/DatabaseView.vue | 1 + .../renderer/app/views/DatabaseView/api.ts | 4 + .../DatabaseView/databaseConnectionTypes.ts | 4 + .../external_sources/sql_source/models.py | 2 + .../external_sources/test_dialect_ports.py | 19 + poetry.lock | 162 +++++++- pyproject.toml | 4 +- shared/CLAUDE.md | 2 +- shared/db_dialects/__init__.py | 20 +- shared/db_dialects/base.py | 44 ++- shared/db_dialects/snowflake.py | 355 ++++++++++++++++++ .../db_dialects/test_dialect_contract.py | 9 +- .../db_dialects/test_snowflake_dialect.py | 181 +++++++++ 45 files changed, 1440 insertions(+), 33 deletions(-) create mode 100644 docs/examples/integrations/database_read_snowflake.py create mode 100644 flowfile_core/flowfile_core/alembic/versions/030_database_connection_extra_params.py create mode 100644 flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py create mode 100644 shared/db_dialects/snowflake.py create mode 100644 shared/tests/db_dialects/test_snowflake_dialect.py diff --git a/build_backends/build_backends/main.py b/build_backends/build_backends/main.py index 3c47bd540..7f10dd8be 100644 --- a/build_backends/build_backends/main.py +++ b/build_backends/build_backends/main.py @@ -435,6 +435,9 @@ def main(): # legs; imported lazily (sqlalchemy dialect adapter / is_available # probe), so list it explicitly like the other DB drivers. "pymssql", + # Snowflake driver: imported lazily inside the snowflake dialect, so + # PyInstaller's static scan misses it (hooks-contrib bundles its data). + "snowflake.connector", "alembic", # certifi ships cacert.pem; ssl uses it via certifi.where(). The # data_downloader builds its SSL context against this so urllib calls diff --git a/docs/examples/integrations/database_read_snowflake.py b/docs/examples/integrations/database_read_snowflake.py new file mode 100644 index 000000000..71cea0a3c --- /dev/null +++ b/docs/examples/integrations/database_read_snowflake.py @@ -0,0 +1,40 @@ +"""Read from Snowflake through a stored database connection. + +Runs only when FLOWFILE_TEST_SNOWFLAKE_* credentials are configured (see +flowfile_core/tests/docs_examples/test_docs_examples.py); the account, +warehouse, and role travel in the connection's extra_params. +""" + +import os + +import flowfile as ff + +ff.create_database_connection_if_not_exists( + "analytics-snowflake", + database_type="snowflake", + database=os.environ.get("FLOWFILE_TEST_SNOWFLAKE_DATABASE", "SNOWFLAKE_SAMPLE_DATA"), + username=os.environ["FLOWFILE_TEST_SNOWFLAKE_USER"], + password=os.environ["FLOWFILE_TEST_SNOWFLAKE_PASSWORD"], + extra_params={ + "account": os.environ["FLOWFILE_TEST_SNOWFLAKE_ACCOUNT"], + "warehouse": os.environ["FLOWFILE_TEST_SNOWFLAKE_WAREHOUSE"], + }, +) + +# --8<-- [start:example] +suppliers = ff.read_database( + "analytics-snowflake", + query=""" + SELECT s_name, s_acctbal + FROM tpch_sf1.supplier + WHERE s_acctbal > 9900 + """, +).collect() + +nations = ff.read_database( + "analytics-snowflake", schema_name="TPCH_SF1", table_name="NATION" +).collect() +# --8<-- [end:example] + +assert nations.height == 25 +assert suppliers.height > 0 diff --git a/docs/index.html b/docs/index.html index d9a550a6a..640e1be53 100644 --- a/docs/index.html +++ b/docs/index.html @@ -637,7 +637,7 @@

Scheduling & Triggers

Kafka, Databases & Cloud Storage

Ingest from Kafka/Redpanda as a canvas node. Read and write S3, Azure Data Lake, and GCS. - Connect PostgreSQL, MySQL, SQL Server, SQLite, and DuckDB. + Connect PostgreSQL, MySQL, SQL Server, Snowflake, SQLite, and DuckDB.

diff --git a/docs/users/connect/index.md b/docs/users/connect/index.md index 10954dbcc..1797413fc 100644 --- a/docs/users/connect/index.md +++ b/docs/users/connect/index.md @@ -10,7 +10,7 @@ Each connector below is read-only, write-only, or both. "Where configured" names | Connector | Reads | Writes | Where configured | Setup | |---|---|---|---|---| -| PostgreSQL / MySQL / SQLite / DuckDB / SQL Server | yes | yes | Connections → Database | [Databases](../visual-editor/tutorials/database-connectivity.md) | +| PostgreSQL / MySQL / SQLite / DuckDB / SQL Server / Snowflake | yes | yes | Connections → Database | [Databases](../visual-editor/tutorials/database-connectivity.md) | | Cloud storage (S3 / ADLS / GCS) | yes | yes | Connections → Cloud Storage | [Cloud storage](../visual-editor/tutorials/cloud-connections.md) | | Kafka / Redpanda | yes | no | Connections → Kafka | [Kafka](kafka.md) | | REST API | yes | no | REST API Reader node (inline) | [REST APIs](apis.md#rest-apis) | @@ -23,7 +23,7 @@ Each connector below is read-only, write-only, or both. "Where configured" names ## Databases -Typed connections to **PostgreSQL**, **MySQL**, **SQLite**, **DuckDB**, and **SQL Server**. Both directions are supported: the Database Reader node (`ff.read_database`) runs a query or reads a whole table, and the Database Writer node (`ff.write_database`) writes a frame back. Credentials are stored encrypted and referenced by name. +Typed connections to **PostgreSQL**, **MySQL**, **SQLite**, **DuckDB**, **SQL Server**, and **Snowflake**. Both directions are supported: the Database Reader node (`ff.read_database`) runs a query or reads a whole table, and the Database Writer node (`ff.write_database`) writes a frame back. Credentials are stored encrypted and referenced by name. See [Databases](../visual-editor/tutorials/database-connectivity.md) for the connection form and worked reader/writer examples. diff --git a/docs/users/data-elsewhere.md b/docs/users/data-elsewhere.md index 81f736ca8..ebb0280ab 100644 --- a/docs/users/data-elsewhere.md +++ b/docs/users/data-elsewhere.md @@ -10,7 +10,7 @@ None of this route requires code: every source below is a node in the visual edi ![One encrypted connection store in the middle holds named connections — warehouse, data-lake, events — while flows on the canvas and Python scripts around it reference those connections by name only; rotating a credential once updates every reference.](../assets/images/concepts/connection-store.svg) -Supported out of the box: databases (PostgreSQL, MySQL, SQLite, DuckDB, SQL Server), cloud storage (S3, Azure Data Lake, Google Cloud Storage — with auth methods from stored keys to ambient `aws-cli`/environment credentials), Kafka/Redpanda brokers, and Google Analytics 4 properties. [Connections](visual-editor/connections.md) covers every form field. +Supported out of the box: databases (PostgreSQL, MySQL, SQLite, DuckDB, SQL Server, Snowflake), cloud storage (S3, Azure Data Lake, Google Cloud Storage — with auth methods from stored keys to ambient `aws-cli`/environment credentials), Kafka/Redpanda brokers, and Google Analytics 4 properties. [Connections](visual-editor/connections.md) covers every form field. ## 2. Read from the source, not from copies @@ -18,7 +18,7 @@ Each source is a drag-and-drop node on the canvas — with a matching `ff.*` cal | Your data is in… | On the canvas | In Python | |---|---|---| -| PostgreSQL / MySQL / SQLite / DuckDB / SQL Server | Database Reader | `ff.read_database` | +| PostgreSQL / MySQL / SQLite / DuckDB / SQL Server / Snowflake | Database Reader | `ff.read_database` | | S3 / ADLS / GCS (CSV, Parquet, JSON, Delta, Iceberg) | Cloud Storage Reader | `ff.scan_parquet_from_cloud_storage`, … | | A Kafka / Redpanda topic | Kafka Source | `ff.read_kafka` | | A REST endpoint | REST API Reader | `ff.read_api` | diff --git a/docs/users/python-api/reference/reading-data.md b/docs/users/python-api/reference/reading-data.md index e65539f19..d849257ed 100644 --- a/docs/users/python-api/reference/reading-data.md +++ b/docs/users/python-api/reference/reading-data.md @@ -397,6 +397,36 @@ The tested SQL Server example reads a table and a query through a stored connect --8<-- "docs/examples/integrations/database_read_mssql.py:example" ``` +### Snowflake + +Snowflake connections use `database_type="snowflake"` with no host or port — the account +identifier (plus an optional warehouse and role) goes in `extra_params`: + +```python +ff.create_database_connection( + connection_name="analytics-snowflake", + database_type="snowflake", + database="ANALYTICS", + username="user", + password="pass", + extra_params={ + "account": "myorg-myaccount", + "warehouse": "COMPUTE_WH", + "role": "ANALYST", + }, +) + +df = ff.read_database("analytics-snowflake", schema_name="PUBLIC", table_name="EVENTS") +``` + +Semi-structured columns (`VARIANT`, `OBJECT`, `ARRAY`) are read as JSON text. The tested +Snowflake example reads a table and a query through a stored connection (it runs only when +Snowflake test credentials are configured): + +```python +--8<-- "docs/examples/integrations/database_read_snowflake.py:example" +``` + ## Connection Management Set up cloud and database connections once, then reference them by name. See [Cloud Connection Management](cloud-connections.md). diff --git a/docs/users/python-api/reference/writing-data.md b/docs/users/python-api/reference/writing-data.md index bbb39e144..41dd9d3d3 100644 --- a/docs/users/python-api/reference/writing-data.md +++ b/docs/users/python-api/reference/writing-data.md @@ -286,6 +286,12 @@ Returns a new child `FlowFrame`. natively rather than text-encoded. A DuckDB file accepts one writer at a time — close other tools using the file while a flow writes to it. +!!! note "Snowflake" + Connections with `database_type="snowflake"` write through the native Snowflake + connector with the same `if_exists` modes. The table is created from the frame's + schema; column names are written case-exact, and nested columns are stored as JSON + text. + ## Write Modes ### Overwrite vs Append diff --git a/docs/users/visual-editor/connections.md b/docs/users/visual-editor/connections.md index b1f0515ab..0e0ff8e40 100644 --- a/docs/users/visual-editor/connections.md +++ b/docs/users/visual-editor/connections.md @@ -26,6 +26,7 @@ and Cloud Storage Writer nodes without re-entering credentials each time. | **SQLite** | `sqlite` | | **DuckDB** | `duckdb` | | **SQL Server** | `mssql` | +| **Snowflake** | `snowflake` | !!! note "File-based connections (SQLite, DuckDB)" SQLite and DuckDB connect to a local database **file path** (e.g. `/path/to/database.db` @@ -40,6 +41,12 @@ and Cloud Storage Writer nodes without re-entering credentials each time. `INTERVAL` columns are read as text — calendar intervals (months) have no fixed length, so there is no matching Polars type. +!!! note "Snowflake connections" + Snowflake has no host or port: the form asks for the **Account** identifier + (e.g. `myorg-myaccount`) plus an optional **Warehouse** and **Role**, together with the + usual username, password, and database. Connections always use TLS, so there is no SSL + toggle. Semi-structured columns (`VARIANT`, `OBJECT`, `ARRAY`) are read as JSON text. + ### Creating a Database Connection 1. Open the **Connections** page from the left sidebar and select the **Database** tab @@ -49,8 +56,8 @@ and Cloud Storage Writer nodes without re-entering credentials each time. | Field | Description | Example | |-------|-------------|---------| | **Connection Name** | Unique identifier for this connection | `prod_postgres` | -| **Database Type** | PostgreSQL, MySQL, SQLite, DuckDB, or SQL Server | `postgresql` | -| **Host** | Database server hostname | `db.example.com` | +| **Database Type** | PostgreSQL, MySQL, SQLite, DuckDB, SQL Server, or Snowflake | `postgresql` | +| **Host** | Database server hostname (Snowflake asks for an account/warehouse/role instead) | `db.example.com` | | **Port** | Database port | `5432` | | **Database** | Database name | `analytics` | | **Username** | Database user | `readonly_user` | diff --git a/docs/what-is-flowfile-technical.md b/docs/what-is-flowfile-technical.md index ffaa58248..b1cee8ca0 100644 --- a/docs/what-is-flowfile-technical.md +++ b/docs/what-is-flowfile-technical.md @@ -12,7 +12,7 @@ Flowfile bundles a visual flow editor, a data catalog, a scheduler, and a Polars **Python environments.** [Kernels](users/visual-editor/kernels.md) are Docker containers with CPU/memory limits, optional GPU passthrough, and a pinned package set. Everyone who runs the flow or opens the [notebook](users/visual-editor/catalog/notebooks.md) executes against the same environment. User code accesses data through `flowfile_ctx`; raw credentials are not exposed to it. -**Source connections.** [Named connections](users/data-elsewhere.md) cover PostgreSQL/MySQL/SQLite/DuckDB/SQL Server, S3/ADLS/GCS (CSV, Parquet, JSON, Delta, Iceberg), Kafka, REST endpoints, and GA4. The database reader accepts a query, so filtering and pre-aggregation can run at the source. [Kafka consumption](users/connect/kafka.md) is incremental by consumer group: offsets commit only after a successful run, so a scheduled flow reads exactly the messages that arrived since its last successful run. +**Source connections.** [Named connections](users/data-elsewhere.md) cover PostgreSQL/MySQL/SQLite/DuckDB/SQL Server/Snowflake, S3/ADLS/GCS (CSV, Parquet, JSON, Delta, Iceberg), Kafka, REST endpoints, and GA4. The database reader accepts a query, so filtering and pre-aggregation can run at the source. [Kafka consumption](users/connect/kafka.md) is incremental by consumer group: offsets commit only after a successful run, so a scheduled flow reads exactly the messages that arrived since its last successful run. **Data freshness.** [Cron schedules and table triggers](users/visual-editor/catalog/schedules.md) are part of the catalog; the scheduler is embedded in the core service (a standalone mode exists). A table update triggers the flows watching it; set-triggers fire when all listed tables have updated. Dependent-pipeline execution is derived from trigger edges — there is no separate DAG definition to maintain. diff --git a/flowfile_core/flowfile_core/alembic/versions/030_database_connection_extra_params.py b/flowfile_core/flowfile_core/alembic/versions/030_database_connection_extra_params.py new file mode 100644 index 000000000..abb69c0a2 --- /dev/null +++ b/flowfile_core/flowfile_core/alembic/versions/030_database_connection_extra_params.py @@ -0,0 +1,37 @@ +"""Add extra_params to database_connections. + +JSON dict of dialect-specific connection parameters that don't fit the +host/port shape (e.g. Snowflake's account/warehouse/role). Keys that could +override credentials are rejected at the API boundary and dropped again at +URI-build time (shared.db_dialects.base.is_blocked_extra_param). NULL means +the connection has no dialect-specific parameters. + +Revision ID: 030 +Revises: 029 +Create Date: 2026-08-05 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import inspect + +revision: str = "030" +down_revision: str | None = "029" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _has_column(table: str, column: str) -> bool: + return column in {c["name"] for c in inspect(op.get_bind()).get_columns(table)} + + +def upgrade() -> None: + if not _has_column("database_connections", "extra_params"): + op.add_column("database_connections", sa.Column("extra_params", sa.Text, nullable=True)) + + +def downgrade() -> None: + if _has_column("database_connections", "extra_params"): + op.drop_column("database_connections", "extra_params") diff --git a/flowfile_core/flowfile_core/database/models.py b/flowfile_core/flowfile_core/database/models.py index 57dd8fd26..05e558645 100644 --- a/flowfile_core/flowfile_core/database/models.py +++ b/flowfile_core/flowfile_core/database/models.py @@ -58,6 +58,7 @@ class DatabaseConnection(Base): port = Column(Integer) database = Column(String, default=None) ssl_enabled = Column(Boolean, default=False) + extra_params = Column(Text, nullable=True) # JSON dict of dialect-specific params (e.g. snowflake account) password_id = Column(Integer, ForeignKey("secrets.id")) user_id = Column(Integer, ForeignKey("users.id")) diff --git a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py index 9d0b6182e..af5de49f1 100644 --- a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py +++ b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py @@ -1,3 +1,5 @@ +import json + from sqlalchemy.orm import Session from flowfile_core.auth import sharing @@ -13,6 +15,21 @@ _OWNER_ACCESS = AccessInfo(is_owner=True, access_level="owner") +def _dump_extra_params(extra_params: dict[str, str] | None) -> str | None: + return json.dumps(extra_params, sort_keys=True) if extra_params else None + + +def parse_extra_params(raw: str | None) -> dict[str, str] | None: + """Parse the JSON extra_params column; malformed or empty values resolve to None.""" + if not raw: + return None + try: + parsed = json.loads(raw) + except (TypeError, ValueError): + return None + return parsed or None + + def _project_sync_connection(kind: str, name: str, user_id: int, deleted: bool = False) -> None: """Mirror a connection change into the active project folder (no-op when none active).""" from flowfile_core.project import project_sync @@ -44,6 +61,7 @@ def store_database_connection(db: Session, connection: FullDatabaseConnection, u username=connection.username, password_id=password_id, ssl_enabled=connection.ssl_enabled, + extra_params=_dump_extra_params(connection.extra_params), user_id=user_id, ) @@ -73,6 +91,7 @@ def update_database_connection(db: Session, connection: FullDatabaseConnection, db_connection.database_type = connection.database_type db_connection.username = connection.username db_connection.ssl_enabled = connection.ssl_enabled + db_connection.extra_params = _dump_extra_params(connection.extra_params) password_value = connection.password.get_secret_value() if password_value: @@ -171,6 +190,7 @@ def get_database_connection_schema(db: Session, connection_name: str, user_id: i username=db_connection.username, password=password_secret.encrypted_value, ssl_enabled=db_connection.ssl_enabled, + extra_params=parse_extra_params(db_connection.extra_params), ) return None @@ -227,6 +247,7 @@ def database_connection_interface_from_db_connection( port=db_connection.port, database=db_connection.database, ssl_enabled=db_connection.ssl_enabled, + extra_params=parse_extra_params(db_connection.extra_params), id=db_connection.id, access=access, ) diff --git a/flowfile_core/flowfile_core/flowfile/flow_graph.py b/flowfile_core/flowfile_core/flowfile/flow_graph.py index d1320bce2..5e12e9979 100644 --- a/flowfile_core/flowfile_core/flowfile/flow_graph.py +++ b/flowfile_core/flowfile_core/flowfile/flow_graph.py @@ -4592,6 +4592,7 @@ def _func(df: FlowDataEngine): password=decrypt_secret(encrypted_password) if encrypted_password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + **(database_connection.extra_params or {}), ), table_name=table_name, if_exists=database_settings.if_exists or "append", @@ -4680,6 +4681,7 @@ def _func(): password=decrypt_secret(encrypted_password) if encrypted_password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + **(database_connection.extra_params or {}), ), query=None if database_settings.query_mode == "table" else database_settings.query, table_name=database_settings.table_name, @@ -4729,6 +4731,7 @@ def schema_callback(): password=decrypt_secret(encrypted_password) if encrypted_password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + **(database_connection.extra_params or {}), ), query=None if database_settings.query_mode == "table" else database_settings.query, table_name=database_settings.table_name, diff --git a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py index 3cbfbd880..f516beaf8 100644 --- a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py +++ b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py @@ -463,6 +463,7 @@ def _resolve_connection(database_settings: DatabaseSettings, user_id: int) -> Re password=password, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + **(database_connection.extra_params or {}), ) return ResolvedConnection(uri=uri, database_type=database_connection.database_type) diff --git a/flowfile_core/flowfile_core/project/importer.py b/flowfile_core/flowfile_core/project/importer.py index f9eaf161c..55d209d68 100644 --- a/flowfile_core/flowfile_core/project/importer.py +++ b/flowfile_core/flowfile_core/project/importer.py @@ -179,6 +179,7 @@ def _import_db_connection(data: dict, owner_id: int, dotenv: dict, result: Setup port=entry.port, database=entry.database, ssl_enabled=entry.ssl_enabled, + extra_params=entry.extra_params, ) with get_db_context() as db: if _get_own_database_connection(db, name, owner_id): diff --git a/flowfile_core/flowfile_core/project/manifest_entries.py b/flowfile_core/flowfile_core/project/manifest_entries.py index 970221487..57ee13a62 100644 --- a/flowfile_core/flowfile_core/project/manifest_entries.py +++ b/flowfile_core/flowfile_core/project/manifest_entries.py @@ -23,6 +23,7 @@ class DatabaseConnectionEntry(_Entry): database: str | None = None username: str = "" ssl_enabled: bool = False + extra_params: dict[str, str] | None = None password: str | None = None diff --git a/flowfile_core/flowfile_core/project/projection.py b/flowfile_core/flowfile_core/project/projection.py index ab613d646..57c08d9da 100644 --- a/flowfile_core/flowfile_core/project/projection.py +++ b/flowfile_core/flowfile_core/project/projection.py @@ -183,10 +183,13 @@ def remove_stale_flow_files(root: Path, flow_uuid: str, keep: Path) -> None: def _db_connection_dict(db: Session, conn: DatabaseConnection) -> dict: + from flowfile_core.flowfile.database_connection_manager.db_connections import parse_extra_params + secret_name = _secret_name(db, conn.password_id) d = {"kind": "database_connection", "connection_name": conn.connection_name} for f in _DB_PLAIN_FIELDS: d[f] = getattr(conn, f) + d["extra_params"] = parse_extra_params(conn.extra_params) d["password"] = make_placeholder(secret_name) if secret_name else None return d diff --git a/flowfile_core/flowfile_core/routes/routes.py b/flowfile_core/flowfile_core/routes/routes.py index 1f2ee5315..119fd1284 100644 --- a/flowfile_core/flowfile_core/routes/routes.py +++ b/flowfile_core/flowfile_core/routes/routes.py @@ -80,6 +80,7 @@ delete_database_connection, get_all_database_connections_interface, get_database_connection, + parse_extra_params, store_database_connection, update_database_connection, ) @@ -786,6 +787,10 @@ def update_db_connection( changed = changed_target_fields( db_connection, input_connection, ("host", "port", "database", "database_type", "ssl_enabled") ) + # extra_params is a target too (e.g. the snowflake account): compare the row's + # JSON against the incoming dict, normalized, so unchanged params don't trip the guard. + if (parse_extra_params(db_connection.extra_params) or {}) != (input_connection.extra_params or {}): + changed.append("extra_params") require_credentials_on_target_change( changed, has_new_credentials=bool(input_connection.password.get_secret_value()), diff --git a/flowfile_core/flowfile_core/schemas/input_schema.py b/flowfile_core/flowfile_core/schemas/input_schema.py index 43e13cb88..ea4f43791 100644 --- a/flowfile_core/flowfile_core/schemas/input_schema.py +++ b/flowfile_core/flowfile_core/schemas/input_schema.py @@ -917,6 +917,18 @@ def get_default_description(self) -> str: return f"{name} ({rf.file_type})" +def _validate_extra_params(v: dict[str, str] | None) -> dict[str, str] | None: + """Normalize empty to None and reject keys that could override auth/target settings.""" + if not v: + return None + from shared.db_dialects import is_blocked_extra_param + + blocked = sorted(key for key in v if is_blocked_extra_param(key)) + if blocked: + raise ValueError(f"extra_params may not override connection auth or target settings: {', '.join(blocked)}") + return v + + class DatabaseConnection(BaseModel): """Defines the connection parameters for a database.""" @@ -927,6 +939,7 @@ class DatabaseConnection(BaseModel): port: int | None = None database: str | None = None url: str | None = None + extra_params: dict[str, str] | None = None @field_validator("database_type") @classmethod @@ -945,6 +958,11 @@ def empty_string_to_none(cls, v): return None return v + @field_validator("extra_params") + @classmethod + def guard_extra_params(cls, v): + return _validate_extra_params(v) + class FullDatabaseConnection(BaseModel): """A complete database connection model including the secret password.""" @@ -958,6 +976,7 @@ class FullDatabaseConnection(BaseModel): database: str | None = None ssl_enabled: bool | None = False url: str | None = None + extra_params: dict[str, str] | None = None @field_validator("database_type") @classmethod @@ -965,6 +984,11 @@ def normalize_database_type(cls, v: str) -> str: # lowercase only, no vocabulary check: legacy stored types (e.g. redshift) must keep loading return v.lower() + @field_validator("extra_params") + @classmethod + def guard_extra_params(cls, v): + return _validate_extra_params(v) + class FullDatabaseConnectionInterface(BaseModel): """A database connection model intended for UI display, omitting the password.""" @@ -977,6 +1001,7 @@ class FullDatabaseConnectionInterface(BaseModel): database: str | None = None ssl_enabled: bool | None = False url: str | None = None + extra_params: dict[str, str] | None = None id: int | None = None access: AccessInfo | None = None diff --git a/flowfile_core/tests/docs_examples/test_docs_examples.py b/flowfile_core/tests/docs_examples/test_docs_examples.py index 19b6db0b0..ca5f4c539 100644 --- a/flowfile_core/tests/docs_examples/test_docs_examples.py +++ b/flowfile_core/tests/docs_examples/test_docs_examples.py @@ -83,10 +83,18 @@ def _mssql_available() -> bool: return mssql_fixtures.can_connect_to_db() +def _snowflake_available() -> bool: + import os + + required = ("ACCOUNT", "USER", "PASSWORD", "WAREHOUSE") + return all(os.environ.get(f"FLOWFILE_TEST_SNOWFLAKE_{name}") for name in required) + + INTEGRATION_GATES = { "database_read": _postgres_available, "database_read_duckdb": lambda: True, # in-process, no backing service "database_read_mssql": _mssql_available, + "database_read_snowflake": _snowflake_available, # live account only; skipped in CI "database_transform_write": _postgres_available, "cloud_storage_s3": _minio_available, "kafka_read": _kafka_available, diff --git a/flowfile_core/tests/flowfile/external_sources/test_dialect_vocabulary.py b/flowfile_core/tests/flowfile/external_sources/test_dialect_vocabulary.py index 94aee9025..87e326a9c 100644 --- a/flowfile_core/tests/flowfile/external_sources/test_dialect_vocabulary.py +++ b/flowfile_core/tests/flowfile/external_sources/test_dialect_vocabulary.py @@ -28,6 +28,15 @@ def test_database_connection_accepts_mixed_case_and_normalizes_it(): assert DatabaseConnection(database_type="SQLITE").database_type == "sqlite" +def test_database_connection_rejects_blocked_extra_params(): + with pytest.raises(ValidationError, match="extra_params may not override"): + DatabaseConnection(database_type="snowflake", extra_params={"account": "a", "password": "x"}) + + +def test_database_connection_normalizes_empty_extra_params_to_none(): + assert DatabaseConnection(database_type="snowflake", extra_params={}).extra_params is None + + class TestSqlSourceDialectInference: """SqlSource resolves its dialect from the explicit database_type, else the URI scheme.""" @@ -51,6 +60,13 @@ def test_mssql_scheme_resolves_to_top_based_limits(self): assert source.dialect.name == "mssql" assert source.get_sample_query() == "SELECT TOP 1 * FROM t" + def test_snowflake_scheme_resolves_to_native_dialect(self): + from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source import SqlSource + + source = SqlSource(connection_string="snowflake://u:p@acct/d?warehouse=WH", table_name="t") + assert source.dialect.name == "snowflake" + assert source.get_sample_query() == "SELECT * FROM t LIMIT 1" + def test_explicit_database_type_wins_over_scheme(self): from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source import SqlSource diff --git a/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py new file mode 100644 index 000000000..262f8a977 --- /dev/null +++ b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py @@ -0,0 +1,190 @@ +"""Snowflake connection-model tests: extra_params CRUD round-trip, the repoint +guard's normalized comparison, and live reads/writes against a real account. + +URI shapes, the type map, and the fakesnow behavioral legs live in +shared/tests/db_dialects/test_snowflake_dialect.py. The live classes here are +gated on FLOWFILE_TEST_SNOWFLAKE_* credentials (account/user/password/database/ +warehouse) per the roadmap's testing decision, so CI skips them. +""" + +import logging +import os +import uuid + +import polars as pl +import pytest +from pydantic import SecretStr, ValidationError + +from flowfile_core.database.connection import get_db_context +from flowfile_core.flowfile.database_connection_manager.db_connections import ( + delete_database_connection, + get_database_connection, + get_database_connection_schema, + parse_extra_params, + store_database_connection, + update_database_connection, +) +from flowfile_core.routes._connection_sharing import require_credentials_on_target_change +from flowfile_core.schemas.input_schema import FullDatabaseConnection +from shared.db_dialects import get_dialect + +logger = logging.getLogger(__name__) + +_LIVE_ENV = ("ACCOUNT", "USER", "PASSWORD", "DATABASE", "WAREHOUSE") + + +def _live_credentials_configured() -> bool: + return all(os.environ.get(f"FLOWFILE_TEST_SNOWFLAKE_{name}") for name in _LIVE_ENV) + + +snowflake_available = pytest.mark.skipif( + not _live_credentials_configured(), + reason="FLOWFILE_TEST_SNOWFLAKE_* credentials are not configured", +) + +USER_ID = 1 +EXTRA_PARAMS = {"account": "myorg-myaccount", "warehouse": "COMPUTE_WH", "role": "ANALYST"} + + +def _snowflake_connection(name: str, extra_params: dict[str, str] | None = EXTRA_PARAMS) -> FullDatabaseConnection: + return FullDatabaseConnection( + connection_name=name, + database_type="snowflake", + username="user", + password=SecretStr("pass"), + database="ANALYTICS", + extra_params=extra_params, + ) + + +def _cleanup(name: str) -> None: + with get_db_context() as db: + if get_database_connection(db, name, USER_ID) is not None: + delete_database_connection(db, name, USER_ID) + + +class TestExtraParamsCrud: + def test_store_and_reload_round_trips_extra_params(self): + name = "snowflake_crud_store" + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _snowflake_connection(name), USER_ID) + with get_db_context() as db: + reloaded = get_database_connection_schema(db, name, USER_ID) + assert reloaded is not None + assert reloaded.extra_params == EXTRA_PARAMS + finally: + _cleanup(name) + + def test_update_persists_changed_extra_params(self): + name = "snowflake_crud_update" + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _snowflake_connection(name), USER_ID) + updated = _snowflake_connection(name, {"account": "other-account"}) + with get_db_context() as db: + update_database_connection(db, updated, USER_ID) + with get_db_context() as db: + reloaded = get_database_connection_schema(db, name, USER_ID) + assert reloaded is not None and reloaded.extra_params == {"account": "other-account"} + finally: + _cleanup(name) + + def test_blocked_extra_params_rejected_at_model_boundary(self): + with pytest.raises(ValidationError, match="extra_params may not override"): + _snowflake_connection("snowflake_evil", {"account": "a", "private_key_file": "/tmp/k.p8"}) + + def test_parse_extra_params_normalizes(self): + assert parse_extra_params(None) is None + assert parse_extra_params("") is None + assert parse_extra_params("not json") is None + assert parse_extra_params("{}") is None + assert parse_extra_params('{"account": "a"}') == {"account": "a"} + + +class TestRepointGuardComparison: + """The PUT route treats a changed extra_params as a target change; the comparison + is normalized (row JSON vs incoming dict) so an unchanged dict never trips it.""" + + @staticmethod + def _changed(row_json: str | None, incoming: dict[str, str] | None) -> bool: + return (parse_extra_params(row_json) or {}) != (incoming or {}) + + def test_unchanged_extra_params_do_not_trip_the_guard(self): + assert self._changed('{"account": "a", "warehouse": "w"}', {"warehouse": "w", "account": "a"}) is False + assert self._changed(None, None) is False + assert self._changed(None, {}) is False + + def test_changed_account_trips_the_guard(self): + assert self._changed('{"account": "a"}', {"account": "attacker"}) is True + assert self._changed(None, {"account": "a"}) is True + + def test_guard_requires_credentials_on_change(self): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + require_credentials_on_target_change( + ["extra_params"], has_new_credentials=False, has_bundled_secrets=True + ) + assert exc.value.status_code == 422 + require_credentials_on_target_change(["extra_params"], has_new_credentials=True, has_bundled_secrets=True) + require_credentials_on_target_change([], has_new_credentials=False, has_bundled_secrets=True) + + +def _live_uri() -> str: + dialect = get_dialect("snowflake") + extra = { + "account": os.environ["FLOWFILE_TEST_SNOWFLAKE_ACCOUNT"], + "warehouse": os.environ["FLOWFILE_TEST_SNOWFLAKE_WAREHOUSE"], + } + if os.environ.get("FLOWFILE_TEST_SNOWFLAKE_ROLE"): + extra["role"] = os.environ["FLOWFILE_TEST_SNOWFLAKE_ROLE"] + if os.environ.get("FLOWFILE_TEST_SNOWFLAKE_SCHEMA"): + extra["schema"] = os.environ["FLOWFILE_TEST_SNOWFLAKE_SCHEMA"] + return dialect.build_uri( + username=os.environ["FLOWFILE_TEST_SNOWFLAKE_USER"], + password=os.environ["FLOWFILE_TEST_SNOWFLAKE_PASSWORD"], + database=os.environ["FLOWFILE_TEST_SNOWFLAKE_DATABASE"], + **extra, + ) + + +@snowflake_available +class TestSnowflakeLive: + def test_write_read_roundtrip_and_fast_schema_parity(self): + dialect = get_dialect("snowflake") + uri = _live_uri() + table = f"ff_test_{uuid.uuid4().hex[:12]}" + df = pl.DataFrame({"id": [1, 2, 3], "score": [0.5, 1.5, 2.5], "label": ["a", "b", "c"]}) + try: + dialect.write(df, uri=uri, table_name=table, if_exists="replace") + dialect.write(df, uri=uri, table_name=table, if_exists="append") + with pytest.raises(ValueError, match="already exists"): + dialect.write(df, uri=uri, table_name=table, if_exists="fail") + dialect.write(df, uri=uri, table_name=table, if_exists="replace") + + result = dialect.read(f"SELECT * FROM {table}", uri, logger) + assert result.height == df.height + assert result.columns == df.columns + + predicted = dialect.query_schema(uri, f"SELECT * FROM {table}") + assert predicted is not None and dict(predicted) == dict(result.schema) + predicted_table = dialect.table_schema(uri, table, None) + assert predicted_table is not None and dict(predicted_table) == dict(result.schema) + finally: + dialect._execute_rows(uri, f"DROP TABLE IF EXISTS {table}") + + def test_browse_lists_schemas_and_tables(self): + dialect = get_dialect("snowflake") + uri = _live_uri() + schemas = dialect.list_schemas(uri) + assert schemas, "expected at least one schema" + + def test_sql_source_schema_prediction(self): + from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source import SqlSource + + source = SqlSource(connection_string=_live_uri(), query="SELECT 1 AS one", database_type="snowflake") + columns = source.get_schema() + assert [c.name for c in columns] == ["ONE"] diff --git a/flowfile_core/tests/project/test_roundtrip.py b/flowfile_core/tests/project/test_roundtrip.py index 16cb818b7..bf8d6d657 100644 --- a/flowfile_core/tests/project/test_roundtrip.py +++ b/flowfile_core/tests/project/test_roundtrip.py @@ -1268,3 +1268,50 @@ def test_flow_writer_projects_portable_namespace(tmp_path): _delete_flow(flow_uuid) _cleanup_custom_namespaces(["WrCat"]) project_sync.close_project(OWNER) + + +def test_database_connection_extra_params_round_trip(tmp_path, monkeypatch): + """A snowflake connection's extra_params (account/warehouse) must survive + project → import → project, not silently drop out of the git projection.""" + project_sync.close_project(OWNER) + conn = "proj_db_snowflake" + extra_params = {"account": "myorg-myaccount", "warehouse": "COMPUTE_WH"} + with get_db_context() as db: + if get_database_connection(db, conn, OWNER) is None: + store_database_connection( + db, + input_schema.FullDatabaseConnection( + connection_name=conn, + database_type="snowflake", + username="etl_reader", + password="s3cr3t", + database="ANALYTICS", + extra_params=extra_params, + ), + OWNER, + ) + root = tmp_path / "project" + try: + project_sync.init_project(str(root), "Snowflake RT", OWNER) + conn_text = (root / "connections" / "database" / f"{conn}.yaml").read_text(encoding="utf-8") + assert "myorg-myaccount" in conn_text + assert "s3cr3t" not in conn_text + + project_sync.close_project(OWNER) + with get_db_context() as db: + delete_database_connection(db, conn, OWNER) + + monkeypatch.setenv("FLOWFILE_SECRET_PROJ_DB_SNOWFLAKE", "s3cr3t") + from flowfile_core.project.importer import import_project + + import_project(root, OWNER) + from flowfile_core.flowfile.database_connection_manager.db_connections import ( + get_database_connection_schema, + ) + + with get_db_context() as db: + restored = get_database_connection_schema(db, conn, OWNER) + assert restored is not None + assert restored.extra_params == extra_params + finally: + _cleanup([conn], []) diff --git a/flowfile_core/tests/test_db_dialect_endpoints.py b/flowfile_core/tests/test_db_dialect_endpoints.py index 8165d894d..1f654e710 100644 --- a/flowfile_core/tests/test_db_dialect_endpoints.py +++ b/flowfile_core/tests/test_db_dialect_endpoints.py @@ -212,6 +212,51 @@ def test_mssql_in_dialect_catalog(): assert mssql_entry["available"] is True +def test_snowflake_in_dialect_catalog(): + response = client.get("/db_dialects") + assert response.status_code == 200, response.text + entry = next(d for d in response.json() if d["name"] == "snowflake") + assert entry["display_name"] == "Snowflake" + assert entry["file_based"] is False + assert entry["default_port"] == 443 + assert entry["available"] is True + assert [f["name"] for f in entry["extra_fields"]] == ["account", "warehouse", "role"] + assert entry["extra_fields"][0]["required"] is True + assert entry["hidden_fields"] == ["host", "port", "ssl"] + + +def test_create_snowflake_connection_with_extra_params(): + _cleanup_connection("snowflake_conn") + payload = { + "connection_name": "snowflake_conn", + "database_type": "snowflake", + "username": "user", + "password": "pass", + "database": "ANALYTICS", + "extra_params": {"account": "myorg-myaccount", "warehouse": "COMPUTE_WH", "role": "ANALYST"}, + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 200, response.text + listed = client.get("/db_connection_lib").json() + entry = next(c for c in listed if c["connection_name"] == "snowflake_conn") + assert entry["extra_params"] == payload["extra_params"] + _cleanup_connection("snowflake_conn") + + +def test_create_connection_rejects_blocked_extra_params(): + payload = { + "connection_name": "snowflake_evil", + "database_type": "snowflake", + "username": "user", + "password": "pass", + "extra_params": {"account": "acct", "authenticator": "externalbrowser"}, + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 422, response.text + assert "authenticator" in response.text + _cleanup_connection("snowflake_evil") + + def test_create_duckdb_connection_without_credentials(): _cleanup_connection("duckdb_conn") payload = { diff --git a/flowfile_frame/flowfile_frame/database/connection_manager.py b/flowfile_frame/flowfile_frame/database/connection_manager.py index d33bd544e..d93edd1ca 100644 --- a/flowfile_frame/flowfile_frame/database/connection_manager.py +++ b/flowfile_frame/flowfile_frame/database/connection_manager.py @@ -40,13 +40,14 @@ def create_database_connection( password: str | SecretStr | None = None, ssl_enabled: bool = False, url: str | None = None, + extra_params: dict[str, str] | None = None, ) -> FullDatabaseConnection: """Create and store a new database connection. Args: connection_name: Unique name for this connection. database_type: Type of database (one of shared.db_dialects.KNOWN_DIALECT_NAMES, - e.g. postgresql, mysql, sqlite, duckdb, mssql). + e.g. postgresql, mysql, sqlite, duckdb, mssql, snowflake). host: Database server hostname. port: Database server port. database: Database name. @@ -54,6 +55,9 @@ def create_database_connection( password: Database password (not needed for file-based types like sqlite/duckdb). ssl_enabled: Whether to use SSL for the connection. url: Full database URL (overrides other connection parameters). + extra_params: Dialect-specific connection parameters, e.g. for Snowflake + ``{"account": "myorg-myaccount", "warehouse": "COMPUTE_WH", "role": "ANALYST"}``. + Keys that could override credentials (user, password, host, ...) are rejected. Returns: FullDatabaseConnection: The created connection object. @@ -86,6 +90,7 @@ def create_database_connection( password=password, ssl_enabled=ssl_enabled, url=url, + extra_params=extra_params, ) with get_db_context() as db: @@ -105,13 +110,14 @@ def create_database_connection_if_not_exists( password: str | SecretStr | None = None, ssl_enabled: bool = False, url: str | None = None, + extra_params: dict[str, str] | None = None, ) -> FullDatabaseConnection: """Create a database connection if it doesn't already exist. Args: connection_name: Unique name for this connection. database_type: Type of database (one of shared.db_dialects.KNOWN_DIALECT_NAMES, - e.g. postgresql, mysql, sqlite, duckdb, mssql). + e.g. postgresql, mysql, sqlite, duckdb, mssql, snowflake). host: Database server hostname. port: Database server port. database: Database name. @@ -119,6 +125,7 @@ def create_database_connection_if_not_exists( password: Database password (not needed for file-based types like sqlite/duckdb). ssl_enabled: Whether to use SSL for the connection. url: Full database URL (overrides other connection parameters). + extra_params: Dialect-specific connection parameters (see create_database_connection). Returns: FullDatabaseConnection: The existing or newly created connection. @@ -139,6 +146,7 @@ def create_database_connection_if_not_exists( password=password, ssl_enabled=ssl_enabled, url=url, + extra_params=extra_params, ) @@ -163,6 +171,7 @@ def get_all_available_database_connections() -> list[FullDatabaseConnectionInter List of database connection interfaces (without passwords). """ from flowfile_core.database.models import DatabaseConnection as DBConnectionModel + from flowfile_core.flowfile.database_connection_manager.db_connections import parse_extra_params user_id = get_current_user_id() with get_db_context() as db: @@ -177,6 +186,7 @@ def get_all_available_database_connections() -> list[FullDatabaseConnectionInter port=conn.port, database=conn.database, ssl_enabled=conn.ssl_enabled, + extra_params=parse_extra_params(conn.extra_params), ) for conn in connections ] diff --git a/flowfile_frame/flowfile_frame/database/connection_manager.pyi b/flowfile_frame/flowfile_frame/database/connection_manager.pyi index cd58bfa5d..82dfe0a7f 100644 --- a/flowfile_frame/flowfile_frame/database/connection_manager.pyi +++ b/flowfile_frame/flowfile_frame/database/connection_manager.pyi @@ -6,8 +6,8 @@ from pydantic import SecretStr from flowfile_core.schemas.input_schema import FullDatabaseConnection, FullDatabaseConnectionInterface def get_current_user_id() -> int: ... -def create_database_connection(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None) -> FullDatabaseConnection: ... -def create_database_connection_if_not_exists(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None) -> FullDatabaseConnection: ... +def create_database_connection(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None) -> FullDatabaseConnection: ... +def create_database_connection_if_not_exists(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None) -> FullDatabaseConnection: ... def get_database_connection_by_name(connection_name: str) -> FullDatabaseConnection | None: ... def get_all_available_database_connections() -> list[FullDatabaseConnectionInterface]: ... def del_database_connection(connection_name: str) -> bool: ... diff --git a/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts b/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts index 6acef794b..f8acb0d72 100644 --- a/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts +++ b/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts @@ -80,6 +80,19 @@ describe("getDbDialects", () => { "sqlite", "duckdb", "mssql", + "snowflake", ]); }); + + it("carries the snowflake form metadata in the fallback entry", async () => { + const { FALLBACK_DIALECTS } = await loadModule(); + const snowflake = FALLBACK_DIALECTS.find((dialect) => dialect.name === "snowflake"); + expect(snowflake?.extra_fields?.map((field) => field.name)).toEqual([ + "account", + "warehouse", + "role", + ]); + expect(snowflake?.extra_fields?.[0].required).toBe(true); + expect(snowflake?.hidden_fields).toEqual(["host", "port", "ssl"]); + }); }); diff --git a/flowfile_frontend/src/renderer/app/api/dbDialects.ts b/flowfile_frontend/src/renderer/app/api/dbDialects.ts index 19604da7f..ec16b6d45 100644 --- a/flowfile_frontend/src/renderer/app/api/dbDialects.ts +++ b/flowfile_frontend/src/renderer/app/api/dbDialects.ts @@ -1,5 +1,11 @@ import axios from "axios"; +export interface DbDialectFieldInfo { + name: string; + label: string; + required: boolean; +} + export interface DbDialectInfo { name: string; display_name: string; @@ -7,6 +13,10 @@ export interface DbDialectInfo { default_port: number | null; supports_ssl: boolean; available: boolean; + // Dialect-specific connection fields (stored in extra_params) and standard + // form fields the dialect hides; absent on older cores, so keep optional. + extra_fields?: DbDialectFieldInfo[]; + hidden_fields?: string[]; } // Rendered when the catalog request fails (offline, older core). Must mirror the @@ -19,6 +29,8 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ default_port: 5432, supports_ssl: true, available: true, + extra_fields: [], + hidden_fields: [], }, { name: "mysql", @@ -27,6 +39,8 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ default_port: 3306, supports_ssl: false, available: true, + extra_fields: [], + hidden_fields: [], }, { name: "sqlite", @@ -35,6 +49,8 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ default_port: null, supports_ssl: false, available: true, + extra_fields: [], + hidden_fields: [], }, { name: "duckdb", @@ -43,6 +59,8 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ default_port: null, supports_ssl: false, available: true, + extra_fields: [], + hidden_fields: [], }, { name: "mssql", @@ -51,6 +69,22 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ default_port: 1433, supports_ssl: false, available: true, + extra_fields: [], + hidden_fields: [], + }, + { + name: "snowflake", + display_name: "Snowflake", + file_based: false, + default_port: 443, + supports_ssl: false, + available: true, + extra_fields: [ + { name: "account", label: "Account", required: true }, + { name: "warehouse", label: "Warehouse", required: false }, + { name: "role", label: "Role", required: false }, + ], + hidden_fields: ["host", "port", "ssl"], }, ]; diff --git a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue index d9c771274..9e45d4338 100644 --- a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue +++ b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue @@ -53,7 +53,20 @@ -
+
+ + +
+ +
-
+
(); -const { dialects, isFileBased } = useDbDialects(); +const { dialects, isFileBased, extraFields, isFieldHidden } = useDbDialects(); const isFileBasedConnection = computed(() => isFileBased(props.modelValue.database_type)); +const dialectExtraFields = computed(() => extraFields(props.modelValue.database_type)); + +const isHidden = (field: string) => isFieldHidden(props.modelValue.database_type, field); + const emit = defineEmits<{ (e: "update:modelValue", value: DatabaseConnection): void; }>(); @@ -143,6 +160,16 @@ const updateField = ( }); }; +const updateExtraParam = (name: string, value: string) => { + const params = { ...(props.modelValue.extra_params || {}) }; + if (value) { + params[name] = value; + } else { + delete params[name]; + } + updateField("extra_params", Object.keys(params).length ? params : null); +}; + const fetchSecrets = async () => { try { const secrets = await fetchSecretsApi(); diff --git a/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts b/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts index 5b611c07b..a17039166 100644 --- a/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts +++ b/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts @@ -1,6 +1,11 @@ import { ref } from "vue"; -import { DbDialectInfo, FALLBACK_DIALECTS, getDbDialects } from "../api/dbDialects"; +import { + DbDialectFieldInfo, + DbDialectInfo, + FALLBACK_DIALECTS, + getDbDialects, +} from "../api/dbDialects"; const dialects = ref(FALLBACK_DIALECTS); let fetchStarted = false; @@ -27,5 +32,11 @@ export function useDbDialects() { const defaultPort = (databaseType?: string): number | undefined => findDialect(databaseType)?.default_port ?? undefined; - return { dialects, findDialect, isFileBased, defaultPort }; + const extraFields = (databaseType?: string): DbDialectFieldInfo[] => + findDialect(databaseType)?.extra_fields ?? []; + + const isFieldHidden = (databaseType?: string, field?: string): boolean => + (findDialect(databaseType)?.hidden_fields ?? []).includes(field || ""); + + return { dialects, findDialect, isFileBased, defaultPort, extraFields, isFieldHidden }; } diff --git a/flowfile_frontend/src/renderer/app/types/node.types.ts b/flowfile_frontend/src/renderer/app/types/node.types.ts index 660d6d648..40d0392df 100644 --- a/flowfile_frontend/src/renderer/app/types/node.types.ts +++ b/flowfile_frontend/src/renderer/app/types/node.types.ts @@ -709,6 +709,7 @@ export interface DatabaseConnection extends BaseConnection { database_type: string; username?: string; password_ref?: string; // Unused by file-based databases (sqlite, duckdb) + extra_params?: Record | null; // Dialect-specific params (e.g. snowflake account) } export type ConnectionModeOption = "inline" | "reference"; export type IfExistAction = "append" | "replace" | "fail"; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue index 456ea6749..ef25415fd 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue @@ -25,7 +25,20 @@
-
+
+ + +
+ +
-
+
-
+
(); -const { dialects, isFileBased, defaultPort } = useDbDialects(); +const { dialects, isFileBased, defaultPort, extraFields, isFieldHidden } = useDbDialects(); const defaultConnection = (): FullDatabaseConnection => ({ connectionName: "", @@ -174,6 +187,8 @@ watch( connection.value.port = newDefault; } } + // Extra params are dialect-specific; a stale set must not leak into the new dialect. + connection.value.extraParams = undefined; } }, ); @@ -182,6 +197,26 @@ const showPassword = ref(false); const isFileBasedType = computed(() => isFileBased(connection.value.databaseType)); +const dialectExtraFields = computed(() => extraFields(connection.value.databaseType)); + +const isHidden = (field: string) => isFieldHidden(connection.value.databaseType, field); + +const extraParamValue = (name: string): string => connection.value.extraParams?.[name] ?? ""; + +const setExtraParam = (name: string, value: string) => { + const params = { ...(connection.value.extraParams || {}) }; + if (value) { + params[name] = value; + } else { + delete params[name]; + } + connection.value.extraParams = Object.keys(params).length ? params : undefined; +}; + +const requiredExtraFieldsFilled = computed(() => + dialectExtraFields.value.every((f) => !f.required || !!connection.value.extraParams?.[f.name]), +); + const isValid = computed(() => { if (isFileBasedType.value) { return !!connection.value.connectionName && !!connection.value.database; @@ -190,7 +225,8 @@ const isValid = computed(() => { !!connection.value.connectionName && !!connection.value.username && (props.isEditing || !!connection.value.password) && - !!connection.value.host + (isHidden("host") || !!connection.value.host) && + requiredExtraFieldsFilled.value ); }); diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue index 991c525d8..477aa6bab 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue @@ -225,6 +225,7 @@ const showEditModal = (connection: FullDatabaseConnectionInterface) => { database: connection.database || "", sslEnabled: connection.sslEnabled, url: connection.url || "", + extraParams: connection.extraParams || undefined, }; dialogVisible.value = true; }; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts index 893b4f79b..5a1990954 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts @@ -24,6 +24,7 @@ const toPythonFormat = (connection: FullDatabaseConnection): PythonFullDatabaseC database: connection.database, ssl_enabled: connection.sslEnabled, url: connection.url, + extra_params: connection.extraParams, }; }; @@ -55,6 +56,7 @@ export const convertConnectionInterfacePytoTs = ( sslEnabled: pythonConnectionInterface.ssl_enabled, url: pythonConnectionInterface.url, database: pythonConnectionInterface.database, + extraParams: pythonConnectionInterface.extra_params, id: pythonConnectionInterface.id, access: pythonConnectionInterface.access, }; @@ -69,8 +71,10 @@ export const convertConnectionInterfaceTstoPy = ( database_type: dbConnectionInterface.databaseType, host: dbConnectionInterface.host, port: dbConnectionInterface.port, + database: dbConnectionInterface.database, ssl_enabled: dbConnectionInterface.sslEnabled, url: dbConnectionInterface.url, + extra_params: dbConnectionInterface.extraParams, }; }; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts index a9488a83e..3185eb3a7 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts @@ -16,6 +16,7 @@ export interface PythonFullDatabaseConnection { database?: string; ssl_enabled: boolean; url?: string; + extra_params?: Record | null; } export interface FullDatabaseConnection { @@ -28,6 +29,7 @@ export interface FullDatabaseConnection { database?: string; sslEnabled: boolean; url?: string; + extraParams?: Record | null; } export interface PythonFullDatabaseConnectionInterface { @@ -39,6 +41,7 @@ export interface PythonFullDatabaseConnectionInterface { database?: string; ssl_enabled: boolean; url?: string; + extra_params?: Record | null; id?: number; access?: AccessInfo | null; } @@ -52,6 +55,7 @@ export interface FullDatabaseConnectionInterface { database?: string; sslEnabled: boolean; url?: string; + extraParams?: Record | null; id?: number; access?: AccessInfo | null; } diff --git a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py index 1a082838c..e3ea6bf72 100644 --- a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py +++ b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py @@ -17,6 +17,7 @@ class DataBaseConnection(BaseModel): database_type: str = "postgresql" # Database type (postgresql, mysql, etc.) ssl_enabled: bool | None = False url: str | None = None + extra_params: dict[str, str] | None = None # Dialect-specific params (e.g. snowflake account/warehouse) def get_decrypted_secret(self) -> SecretStr: return decrypt_secret(self.password.get_secret_value()) @@ -43,6 +44,7 @@ def create_uri(self) -> str: url=self.url, ssl_enabled=bool(self.ssl_enabled), connect_timeout=10, + **(self.extra_params or {}), ) def create_sqlalchemy_uri(self) -> str: diff --git a/flowfile_worker/tests/external_sources/test_dialect_ports.py b/flowfile_worker/tests/external_sources/test_dialect_ports.py index 0eab9dc25..613554bc8 100644 --- a/flowfile_worker/tests/external_sources/test_dialect_ports.py +++ b/flowfile_worker/tests/external_sources/test_dialect_ports.py @@ -57,3 +57,22 @@ def test_preflight_skips_file_based_dialects(captured_connections): def test_preflight_skips_url_connections(captured_connections): verify_database_reachable(DataBaseConnection(database_type="postgresql", url="postgresql://u@h/d")) assert captured_connections == [] + + +def test_preflight_skips_hostless_snowflake_connections(captured_connections): + """Snowflake's locator is the account in extra_params, so host is None and the + TCP pre-flight has nothing meaningful to probe.""" + verify_database_reachable( + DataBaseConnection(database_type="snowflake", extra_params={"account": "myorg-myaccount"}) + ) + assert captured_connections == [] + + +def test_snowflake_create_uri_carries_extra_params(): + connection = DataBaseConnection( + database_type="snowflake", + username="u", + database="ANALYTICS", + extra_params={"account": "myorg-myaccount", "warehouse": "COMPUTE_WH"}, + ) + assert connection.create_uri() == "snowflake://u@myorg-myaccount/ANALYTICS?warehouse=COMPUTE_WH" diff --git a/poetry.lock b/poetry.lock index 9396e1407..050c495e8 100644 --- a/poetry.lock +++ b/poetry.lock @@ -368,6 +368,17 @@ files = [ [package.dependencies] typing-extensions = {version = "*", markers = "python_full_version < \"3.12\""} +[[package]] +name = "asn1crypto" +version = "1.5.1" +description = "Fast ASN.1 parser and serializer with definitions for private keys, public keys, certificates, CRL, OCSP, CMS, PKCS#3, PKCS#7, PKCS#8, PKCS#12, PKCS#5, X.509 and TSP" +optional = false +python-versions = "*" +files = [ + {file = "asn1crypto-1.5.1-py2.py3-none-any.whl", hash = "sha256:db4e40728b728508912cbb3d44f19ce188f218e9eba635821bb4b68564f8fd67"}, + {file = "asn1crypto-1.5.1.tar.gz", hash = "sha256:13ae38502be632115abf8a24cbe5f4da52e3b5231990aff31123c805306ccb9c"}, +] + [[package]] name = "async-timeout" version = "5.0.1" @@ -1352,6 +1363,25 @@ files = [ [package.dependencies] python-dateutil = ">=2.4" +[[package]] +name = "fakesnow" +version = "0.11.11" +description = "Fake Snowflake Connector for Python. Run, mock and test Snowflake DB locally." +optional = false +python-versions = ">=3.10" +files = [ + {file = "fakesnow-0.11.11-py3-none-any.whl", hash = "sha256:5e681610f7371425e56538070a2bfcbe10a2dba1a021847eb5bef03c690325d5"}, +] + +[package.dependencies] +duckdb = ">=1.5.4,<1.6.0" +pyarrow = "*" +snowflake-connector-python = "*" +sqlglot = ">=30.12.0,<30.13.0" + +[package.extras] +server = ["starlette", "uvicorn"] + [[package]] name = "fastapi" version = "0.115.14" @@ -4355,6 +4385,23 @@ files = [ [package.dependencies] packaging = ">=22.0" +[[package]] +name = "pyjwt" +version = "2.13.0" +description = "JSON Web Token implementation in Python" +optional = false +python-versions = ">=3.9" +files = [ + {file = "pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728"}, + {file = "pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423"}, +] + +[package.dependencies] +typing_extensions = {version = ">=4.0", markers = "python_version < \"3.11\""} + +[package.extras] +crypto = ["cryptography (>=3.4.0)"] + [[package]] name = "pymdown-extensions" version = "11.0.1" @@ -4440,6 +4487,25 @@ files = [ ed25519 = ["PyNaCl (>=1.6.2)"] rsa = ["cryptography (>=46.0.7)"] +[[package]] +name = "pyopenssl" +version = "26.2.0" +description = "Python wrapper module around the OpenSSL library" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pyopenssl-26.2.0-py3-none-any.whl", hash = "sha256:4f9d971bc5298b8bc1fab282803da04bf000c755d4ad9d99b52de2569ca19a70"}, + {file = "pyopenssl-26.2.0.tar.gz", hash = "sha256:8c6fcecd1183a7fc897548dfe388b0cdb7f37e018200d8409cf33959dbe35387"}, +] + +[package.dependencies] +cryptography = ">=46.0.0,<49" +typing-extensions = {version = ">=4.9", markers = "python_version < \"3.13\" and python_version >= \"3.8\""} + +[package.extras] +docs = ["sphinx (!=5.2.0,!=5.2.0.post0,!=7.2.5)", "sphinx_rtd_theme"] +test = ["pretend", "pytest (>=3.0.1)", "pytest-rerunfailures"] + [[package]] name = "pyparsing" version = "3.3.2" @@ -4591,6 +4657,17 @@ files = [ {file = "python_multipart-0.0.32.tar.gz", hash = "sha256:be54b7f3fa167bb83e4fcd936b887b708f4e57fe75911c02aebf53efaf8d938e"}, ] +[[package]] +name = "pytz" +version = "2026.3.post1" +description = "World timezone definitions, modern and historical" +optional = false +python-versions = "*" +files = [ + {file = "pytz-2026.3.post1-py2.py3-none-any.whl", hash = "sha256:dd95840dd199baea12d9cc096a1d452caa6596a1c1e4b5f3dbd1541855d5e815"}, + {file = "pytz-2026.3.post1.tar.gz", hash = "sha256:2211d3fcf9a797d3405cac96ac7f61d80e6a644f72a3309607282fe8a2010c5d"}, +] + [[package]] name = "pytzdata" version = "2020.1" @@ -5187,6 +5264,67 @@ files = [ {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, ] +[[package]] +name = "snowflake-connector-python" +version = "4.7.1" +description = "Snowflake Connector for Python" +optional = false +python-versions = ">=3.10" +files = [ + {file = "snowflake_connector_python-4.7.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:ff5bb51c1c21cbbb5c90b9785cc3df9d649c64db26553668b1a1ea8a461a0d5b"}, + {file = "snowflake_connector_python-4.7.1-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:0e23a5aaa1e9eaa9ba88b8867247d55c24e40f3be6367e3aa4f8633b37f56317"}, + {file = "snowflake_connector_python-4.7.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:571eeeb5f3b034671c125186ecb02a2e54ebee1fad1b3af7f091773dc11f693d"}, + {file = "snowflake_connector_python-4.7.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:167ba97d5b615b507fc8234266c8736ea97f91b7324db7a305a28ec3505c0edf"}, + {file = "snowflake_connector_python-4.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:d5821dc73d305b0804aa415a4fafee018e611cc8cf339de36a5c8c44fb57beaa"}, + {file = "snowflake_connector_python-4.7.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4b79e818d83306babff9d0803e697a008e8ada961deff55e3c5da0a9c3505d9a"}, + {file = "snowflake_connector_python-4.7.1-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:f5a066d2c1db940740c49bafe1c8983395bc574bc31465263898c9bd050c5c52"}, + {file = "snowflake_connector_python-4.7.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4f2c6f40544739b43f2da262dd67a32d4ac7ca2b09cc6b7181f02b24c60d5754"}, + {file = "snowflake_connector_python-4.7.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e582f6ac2c5fa53170ae28303d361ea6cc64e695ba528f1b6ec8529476324327"}, + {file = "snowflake_connector_python-4.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:22afd6de9fec8ef2cc23b231af9fc0f352351cd53da0cc11c2495272eff04ccf"}, + {file = "snowflake_connector_python-4.7.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:b03ef22742fee88d387f2ec3969dfb032417a2119bd366414b2c5ac9acb38a45"}, + {file = "snowflake_connector_python-4.7.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:1ed85264edb186a39e641f1b744a1a1973e10d9d15a96a46a1e78ac67442c7b6"}, + {file = "snowflake_connector_python-4.7.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f4e3b1222bc53b56ba4aa88224afe1ad894ee6f7389b2c0105b829d13d559aa7"}, + {file = "snowflake_connector_python-4.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b7b27df86fc13afccabdb525973c0e9a730f8c25f07a4a55f95bc2ad88f8d35"}, + {file = "snowflake_connector_python-4.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:92b935a4f73651ea1306b8c3c58003e0e5a72fa3a86453087161296168c31ed3"}, + {file = "snowflake_connector_python-4.7.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:5011a5eb55dd80fed4198081743f75a8a56412b4623cb9cf61ce2813f600cbd9"}, + {file = "snowflake_connector_python-4.7.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:f2482d1b059fcaa45b7edf8ea97b0f35179ab781fb1ce4716809d83d9f5f4d4c"}, + {file = "snowflake_connector_python-4.7.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8a5e2bd3701176521577eea8c5b384178b404bbcd73d634f86b888929c15fb8a"}, + {file = "snowflake_connector_python-4.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7c67e735d38403109f6bc1a17022a4c950966e1c989f5e2376b9c28c5ece0fb5"}, + {file = "snowflake_connector_python-4.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:81a8f1ae86222e7b8561f41f688462687a435f5afc8aa34ae42f69c0c0f16a57"}, + {file = "snowflake_connector_python-4.7.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:052fd457fa79616d074f5b8f2e106f9f2ca36645526e4d0abddf7e3685d1bdb7"}, + {file = "snowflake_connector_python-4.7.1-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:171ef3515c44e804dd86ab2b288179037fec40b3f45e1aa28fba095b764f9a6a"}, + {file = "snowflake_connector_python-4.7.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:051e49157d955e5ba76577345169dd1869288ecef903dab17ec6180e1009b117"}, + {file = "snowflake_connector_python-4.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3e53626c600b59c181ce18e2b253487fb7e2b0d30ac0d3cfbd3ddffd13743b"}, + {file = "snowflake_connector_python-4.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:7703f33059daa6e30d824e5d88bd525a7feccd38412b888c6e35cf0847cc34b0"}, + {file = "snowflake_connector_python-4.7.1.tar.gz", hash = "sha256:fad8e1fb0c49eb1d93dad785ce738dac17473c6826f8cd36ee8d6f9675bceae1"}, +] + +[package.dependencies] +asn1crypto = ">0.24.0,<2.0.0" +boto3 = ">=1.24" +botocore = ">=1.24" +certifi = ">=2024.7.4" +charset_normalizer = ">=2,<4" +cryptography = ">=46.0.5" +filelock = ">=3.5,<4" +idna = ">=3.7,<4" +packaging = "*" +platformdirs = ">=2.6.0,<5.0.0" +pyjwt = ">=2.10.1,<3.0.0" +pyOpenSSL = ">=24.0.0" +pytz = "*" +requests = ">=2.32.4,<3.0.0" +sortedcontainers = ">=2.4.0" +tomlkit = "*" +typing_extensions = ">=4.3,<5" + +[package.extras] +azure = ["azure-identity (>=1.16)"] +boto = ["boto3 (>=1.24)", "botocore (>=1.24)"] +development = ["Cython", "azure-identity (>=1.16)", "coverage", "mitmproxy (>=12.0.0)", "more-itertools", "numpy (<=2.4.3)", "pendulum (!=2.1.1)", "pexpect", "pytest (<7.5.0)", "pytest-asyncio", "pytest-cov", "pytest-rerunfailures (<16.0)", "pytest-timeout", "pytest-xdist", "pytzdata", "responses"] +pandas = ["pandas (>=1.0.0,<3.0.0)", "pandas (>=2.1.2,<3.0.0)", "pyarrow (>=14.0.1)", "pyarrow (>=14.0.1,<24)"] +secure-local-storage = ["keyring (>=23.1.0,<26.0.0)"] + [[package]] name = "sortedcontainers" version = "2.4.0" @@ -5296,18 +5434,19 @@ sqlcipher = ["sqlcipher3_binary"] [[package]] name = "sqlglot" -version = "27.29.0" +version = "30.12.0" description = "An easily customizable SQL parser and transpiler" optional = false python-versions = ">=3.9" files = [ - {file = "sqlglot-27.29.0-py3-none-any.whl", hash = "sha256:9a5ea8ac61826a7763de10cad45a35f0aa9bfcf7b96ee74afb2314de9089e1cb"}, - {file = "sqlglot-27.29.0.tar.gz", hash = "sha256:2270899694663acef94fa93497971837e6fadd712f4a98b32aee1e980bc82722"}, + {file = "sqlglot-30.12.0-py3-none-any.whl", hash = "sha256:86cccc610073c645c03e72b55b60ae0518aa3253a7fc3bd56551370d003c6554"}, + {file = "sqlglot-30.12.0.tar.gz", hash = "sha256:6b8369704662d4f654bc934cea4dd31c916c2a571b389210cb9e951a275e5fd9"}, ] [package.extras] -dev = ["duckdb (>=0.6)", "maturin (>=1.4,<2.0)", "mypy", "pandas", "pandas-stubs", "pdoc", "pre-commit", "pyperf", "python-dateutil", "pytz", "ruff (==0.7.2)", "types-python-dateutil", "types-pytz", "typing_extensions"] -rs = ["sqlglotrs (==0.7.3)"] +c = ["sqlglotc (==30.12.0)"] +dev = ["duckdb (>=0.6)", "mypy", "pandas", "pandas-stubs", "pdoc", "pre-commit", "pyperf", "python-dateutil", "pytz", "ruff (==0.15.6)", "setuptools_scm", "sqlglot-mypy (>=2.1.0.post3)", "types-python-dateutil", "types-pytz", "typing_extensions"] +rs = ["sqlglotc (==30.12.0)", "sqlglotrs (==0.13.0)"] [[package]] name = "starlette" @@ -5573,6 +5712,17 @@ files = [ {file = "tomli-2.4.1.tar.gz", hash = "sha256:7c7e1a961a0b2f2472c1ac5b69affa0ae1132c39adcb67aba98568702b9cc23f"}, ] +[[package]] +name = "tomlkit" +version = "0.15.1" +description = "Style preserving TOML library" +optional = false +python-versions = ">=3.9" +files = [ + {file = "tomlkit-0.15.1-py3-none-any.whl", hash = "sha256:177a05aece5a8ca5266fd3c448abb47b8d352f09d477d3ca8332db4d89b24304"}, + {file = "tomlkit-0.15.1.tar.gz", hash = "sha256:e25bbf38843005246210a12982776f27f99cb9be67160e14434d0c0d21ee1e97"}, +] + [[package]] name = "tqdm" version = "4.69.1" @@ -6095,4 +6245,4 @@ type = ["pytest-mypy (>=1.0.1)"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.14" -content-hash = "4ff8d2ad69991d92d83e6543c6527bd6d8104b415e5d8d9d5415183975592cb5" +content-hash = "b412158a0a3e3efc4bf8dcfae086d18661053d2faa940fdece3b56216a46eb3b" diff --git a/pyproject.toml b/pyproject.toml index 984a2f063..2e91a436d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ aiofiles = "^24.1.0" python-jose = "^3.4.0" bcrypt = "^4.3.0" connectorx = "^0.4.2" -sqlglot = ">=25.0.0,<28.0.0" +sqlglot = ">=25.0.0,<31.0.0" polars_simed = ">=0.4.0" polars-gw = "0.1.4" @@ -83,6 +83,7 @@ gitpython = "^3.1.40" duckdb = "^1.5.5" pymssql = "^2.3.2" packaging = ">=24.0" +snowflake-connector-python = "^4.7.1" [tool.poetry.scripts] @@ -131,6 +132,7 @@ griffe = ">=1.0,<2" griffe-pydantic = ">=1.1.6,<1.3" ruff = "^0.8.0" mkdocs-redirects = "^1.2.3" +fakesnow = "0.11.11" [build-system] requires = ["poetry-core"] diff --git a/shared/CLAUDE.md b/shared/CLAUDE.md index f7fd75aba..7cbd64eef 100644 --- a/shared/CLAUDE.md +++ b/shared/CLAUDE.md @@ -17,7 +17,7 @@ It is a Poetry package — `{ include = "shared" }` in root `pyproject.toml`, wi - `models.py` — standalone SQLAlchemy models (`FlowRun`, `FlowSchedule`, `FlowRegistration`, `CatalogTable`, `ScheduleTriggerTable`, `SchedulerLock`) on their own `Base`. - `artifact_storage.py` — `ArtifactStorageBackend` ABC (`prepare_upload`/`prepare_download`/`delete`/`exists`) + `SharedFilesystemStorage` / `S3Storage` (presigned-URL) backends, returning `UploadTarget` / `DownloadSource`. Kernel moves blob bytes via presigned URLs; Core stays metadata-only. - `delta_utils.py` / `delta_models.py` — dependency-light Delta-log helpers (`make_json_safe`, `format_delta_timestamp`, `get_delta_size_bytes`, `validate_catalog_path`, plus `write_delta` / `merge_into_delta`, and the SCD2 primitives `scd2_into_delta` / `scd2_surrogate_keys` / `Scd2Result` — one atomic close+insert MERGE per write, with a frozen `sha256-v1` surrogate-key encoding) + Pydantic `DeltaVersionCommit` / `SourceTableVersion`. -- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver) or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. +- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), `snowflake.py` (native snowflake-connector-python end to end — no connectorx, no SQLAlchemy: Arrow-fetch reads cast through a `cursor.describe()` type-code map for predicted==materialized parity with a row-based fallback when the cursor lacks Arrow (fakesnow), account/warehouse/role arrive via the connection's guarded `extra_params` — `base.is_blocked_extra_param` keys can never override auth — port-less account URIs, qmark `executemany` writes, information_schema browse; behavioral tests run against fakesnow, no Docker), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). `DialectInfo` also serves per-dialect `extra_fields`/`hidden_fields` (from `DbDialect` ClassVars) so a new connection shape renders in the frontend forms with zero frontend changes. Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver), `snowflake.py` (native-driver server dialect with its own connection shape), or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. - `sql_utils.py` — `construct_sql_uri`, `get_sqlalchemy_uri`, `SQLALCHEMY_DRIVER_MAP` (caller passes an already-decrypted password); thin dispatchers over `db_dialects` since the registry landed. - `cloud_storage/` — GCS/S3/ADLS helpers: `storage_options.py` (`build_*_storage_options`), `writers.py` (`write_to_cloud` + per-format parquet/csv/json/delta writers), `directory.py` (first-file listing per backend), `uri.py` (scheme list + `parse_uri`/`uri_parent`/`uri_join`/`canonical_scheme` — `pathlib` corrupts `scheme://`, so URI path maths lives here and `catalog/storage_backend.py` imports the scheme list from it), `browse.py` (one-level listing for the storage-browser UI: `BrowseEntry`/`BrowseResult`, `browse_support`, `list_cloud_uri`; boto3 `Delimiter="/"` for S3, `walk_blobs` for ADLS, `gcsfs.ls` for GCS — a refused bucket list is reported as `root_listing_denied`, not an error, and provider exceptions are translated to `BrowseError` subclasses carrying their own status + `error_code`), `gcs.py`, `utils.py`. - `kafka/` — `consumer.py` (`read_kafka_source`, `infer_topic_schema`, `commit_offsets`, `make_kafka_commit_callback`), `models.py`, `deserializers.py` (`get_deserializer`, JSON deserializer). diff --git a/shared/db_dialects/__init__.py b/shared/db_dialects/__init__.py index 7c8e4007b..c2c8f6f2c 100644 --- a/shared/db_dialects/__init__.py +++ b/shared/db_dialects/__init__.py @@ -18,7 +18,7 @@ from pydantic import BaseModel -from shared.db_dialects.base import POSTGRES_FAMILY, DbDialect +from shared.db_dialects.base import POSTGRES_FAMILY, DbDialect, DialectField, is_blocked_extra_param from shared.db_dialects.builtin import ( GenericDialect, MySQLDialect, @@ -27,6 +27,7 @@ ) from shared.db_dialects.duckdb import DuckDBDialect from shared.db_dialects.mssql import MSSQLDialect +from shared.db_dialects.snowflake import SnowflakeDialect if TYPE_CHECKING: import polars as pl @@ -34,15 +35,19 @@ __all__ = [ "POSTGRES_FAMILY", "DbDialect", + "DialectField", + "DialectFieldInfo", "DialectInfo", "DuckDBDialect", "GenericDialect", "KNOWN_DIALECT_NAMES", "MSSQLDialect", + "SnowflakeDialect", "UnknownDialectError", "dialect_catalog", "get_dialect", "get_dialect_or_generic", + "is_blocked_extra_param", "iter_dialects", "read_sql", ] @@ -58,6 +63,7 @@ class UnknownDialectError(ValueError): SQLiteDialect(), DuckDBDialect(), MSSQLDialect(), + SnowflakeDialect(), ) _REGISTRY: dict[str, DbDialect] = {d.name: d for d in _BUILTIN_DIALECTS} @@ -65,6 +71,14 @@ class UnknownDialectError(ValueError): KNOWN_DIALECT_NAMES: tuple[str, ...] = tuple(_REGISTRY) +class DialectFieldInfo(BaseModel): + """A dialect-specific connection field the frontend renders into the form.""" + + name: str + label: str + required: bool = False + + class DialectInfo(BaseModel): """Catalog entry served to the frontend via GET /db_dialects.""" @@ -74,6 +88,8 @@ class DialectInfo(BaseModel): default_port: int | None supports_ssl: bool available: bool + extra_fields: list[DialectFieldInfo] = [] + hidden_fields: list[str] = [] def get_dialect(name: str) -> DbDialect: @@ -110,6 +126,8 @@ def dialect_catalog() -> list[DialectInfo]: default_port=d.default_port, supports_ssl=d.supports_ssl, available=d.is_available(), + extra_fields=[DialectFieldInfo(name=f.name, label=f.label, required=f.required) for f in d.extra_fields], + hidden_fields=list(d.hidden_fields), ) for d in iter_dialects() ] diff --git a/shared/db_dialects/base.py b/shared/db_dialects/base.py index 8ffb5ab61..05a95f3d6 100644 --- a/shared/db_dialects/base.py +++ b/shared/db_dialects/base.py @@ -20,6 +20,7 @@ class *is* the generic dialect — its method bodies are the historical import logging from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar if TYPE_CHECKING: @@ -31,6 +32,42 @@ class *is* the generic dialect — its method bodies are the historical # sslmode/connect_timeout query params are valid (pymysql rejects unknown params). POSTGRES_FAMILY = {"postgresql", "postgres", "redshift"} +# Connection extra_params (dialect-specific settings like Snowflake's +# account/warehouse/role) must never be able to override credentials, the +# connection target, or transport security. Mirrors the Kafka blocked-config +# guard in shared/kafka/models.py: blocked keys are dropped at point of use; +# core additionally rejects them with a 422 at the API boundary. +_BLOCKED_EXTRA_PARAM_PREFIXES = ("private_key", "ssl") +_BLOCKED_EXTRA_PARAMS = frozenset( + { + "password", + "user", + "username", + "host", + "port", + "database", + "dbname", + "authenticator", + "token", + "insecure_mode", + } +) + + +def is_blocked_extra_param(key: str) -> bool: + """Whether a connection extra_params key could override auth/target settings.""" + lowered = key.lower() + return lowered in _BLOCKED_EXTRA_PARAMS or lowered.startswith(_BLOCKED_EXTRA_PARAM_PREFIXES) + + +@dataclass(frozen=True) +class DialectField: + """A dialect-specific connection form field, carried in the connection's extra_params.""" + + name: str + label: str + required: bool = False + class DbDialect: """One database dialect. Base behavior == the historical generic code paths.""" @@ -43,6 +80,11 @@ class DbDialect: sqlalchemy_driver: ClassVar[str | None] = None sqlglot_name: ClassVar[str] = "postgres" install_hint: ClassVar[str | None] = None + # Dialect-specific connection fields (stored in extra_params) and standard form + # fields the dialect does not use; served to the frontend via dialect_catalog() + # so a new connection shape needs no frontend changes. + extra_fields: ClassVar[tuple[DialectField, ...]] = () + hidden_fields: ClassVar[tuple[str, ...]] = () @property def uri_scheme(self) -> str: @@ -97,7 +139,7 @@ def build_uri( "no SSL parameter was applied and the connection may be unencrypted.", scheme, ) - query_params.update(kwargs) + query_params.update({k: v for k, v in kwargs.items() if v is not None and not is_blocked_extra_param(k)}) if query_params: sep = "&" if "?" in base_uri else "?" diff --git a/shared/db_dialects/snowflake.py b/shared/db_dialects/snowflake.py new file mode 100644 index 000000000..8ff632eb9 --- /dev/null +++ b/shared/db_dialects/snowflake.py @@ -0,0 +1,355 @@ +"""Snowflake dialect: native connector Arrow reads, describe()-based fast schema. + +connectorx has no Snowflake backend and the generic SQLAlchemy paths would need +the separate snowflake-sqlalchemy package, so every operation (read, write, +browse) goes through ``snowflake-connector-python`` directly. Reads use the +connector's Arrow fetch; like duckdb, ``cancel_check`` is accepted for +interface parity but not polled (the fetch happens in one native call). + +Snowflake connections don't fit the host/port shape: the locator is an +*account identifier* and warehouse/role select compute/authorization context. +Those arrive through the connection's guarded ``extra_params`` and land in the +URI as ``snowflake://user:pass@account/database?warehouse=...&role=...``. +``build_uri`` drops any blocked key (``base.is_blocked_extra_param``) so extra +params can never override credentials. + +Fast schema uses ``cursor.describe()`` — the query is compiled server-side but +never executed — plus a dialect-local map keyed on the connector's type codes. +``read`` casts the fetched frame through the same map (Snowflake's Arrow +results pick per-batch physical types, e.g. the smallest int width that fits, +so an uncast schema would be value-dependent), which makes predicted schema +equal materialized schema by construction. Semi-structured types +(VARIANT/OBJECT/ARRAY, GEOGRAPHY/GEOMETRY) materialize as JSON text and map to +String; every FIXED column with scale 0 is normalized to Int64 (Snowflake +reports precision 38 for all integer columns). A column whose type is not in +the map makes the schema hooks return ``None`` (the caller keeps its generic +path) while ``read`` still returns whatever the connector yields for it. + +Writes create the table from the frame's schema and bulk-insert with qmark +binding (the connector converts large ``executemany`` batches into stage-based +loads internally). Column identifiers are written quoted so frames round-trip +with their exact column names; table/schema identifiers stay unquoted when +they are plain, keeping Snowflake's case-insensitive resolution for user SQL. +""" + +from __future__ import annotations + +import logging +import re +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, ClassVar + +from shared.db_dialects.base import DbDialect, DialectField, is_blocked_extra_param + +if TYPE_CHECKING: + import polars as pl + +logger = logging.getLogger(__name__) + +_PLAIN_IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$") + + +def _ident(name: str) -> str: + """Quote only non-plain identifiers so plain names keep case-insensitive resolution.""" + if _PLAIN_IDENT.match(name): + return name + return '"' + name.replace('"', '""') + '"' + + +def _quote_ident(name: str) -> str: + return '"' + name.replace('"', '""') + '"' + + +class SnowflakeDialect(DbDialect): + name: ClassVar[str] = "snowflake" + display_name: ClassVar[str] = "Snowflake" + default_port: ClassVar[int | None] = 443 + sqlglot_name: ClassVar[str] = "snowflake" + install_hint: ClassVar[str | None] = "pip install snowflake-connector-python" + extra_fields: ClassVar[tuple[DialectField, ...]] = ( + DialectField("account", "Account", required=True), + DialectField("warehouse", "Warehouse"), + DialectField("role", "Role"), + ) + hidden_fields: ClassVar[tuple[str, ...]] = ("host", "port", "ssl") + + def is_available(self) -> bool: + try: + import snowflake.connector # noqa: F401 + except ImportError: + return False + return True + + def build_uri( + self, + *, + host: str | None = None, + port: int | None = None, + username: str | None = None, + password: str | None = None, + database: str | None = None, + ssl_enabled: bool = False, + connect_timeout: int | None = None, + **kwargs, + ) -> str: + """Account-shaped URI; ``account`` comes from extra_params (``host`` is accepted as an alias).""" + from urllib.parse import quote_plus + + account = kwargs.pop("account", None) or host + if not account: + raise ValueError("Account is required to create a Snowflake URI") + + credentials = "" + if username: + credentials = quote_plus(username) + if password: + credentials += f":{quote_plus(password)}" + credentials += "@" + + uri = f"{self.uri_scheme}://{credentials}{quote_plus(str(account))}" + if database: + uri += f"/{database}" + params = {k: v for k, v in kwargs.items() if v is not None and not is_blocked_extra_param(k)} + if params: + uri += "?" + "&".join(f"{key}={quote_plus(str(value))}" for key, value in params.items()) + return uri + + @staticmethod + def _connect_kwargs(uri: str) -> dict[str, str]: + from urllib.parse import parse_qsl, unquote, urlparse + + parsed = urlparse(uri) + # netloc split instead of .hostname: account identifiers should keep their case. + kwargs: dict[str, str] = {"account": unquote(parsed.netloc.rsplit("@", 1)[-1])} + if parsed.username: + kwargs["user"] = unquote(parsed.username) + if parsed.password: + kwargs["password"] = unquote(parsed.password) + database = parsed.path.lstrip("/") + if database: + kwargs["database"] = unquote(database) + for key, value in parse_qsl(parsed.query): + if key in ("warehouse", "role", "schema") and value: + kwargs[key] = value + return kwargs + + def _connect(self, uri: str): + import snowflake.connector + + return snowflake.connector.connect(paramstyle="qmark", **self._connect_kwargs(uri)) + + @staticmethod + def _polars_dtype(column) -> Any: + """Map a cursor-describe ResultMetadata entry to the dtype the read path materializes.""" + import polars as pl + from snowflake.connector.constants import FIELD_ID_TO_NAME + + type_name = FIELD_ID_TO_NAME.get(column.type_code, "") + if type_name == "FIXED": + scale = column.scale or 0 + if scale == 0: + return pl.Int64 + return pl.Decimal(column.precision or 38, scale) + simple = { + "REAL": pl.Float64, + "TEXT": pl.String, + "DATE": pl.Date, + "TIME": pl.Time, + "TIMESTAMP": pl.Datetime("us"), + "TIMESTAMP_NTZ": pl.Datetime("us"), + "TIMESTAMP_LTZ": pl.Datetime("us", "UTC"), + "TIMESTAMP_TZ": pl.Datetime("us", "UTC"), + "BOOLEAN": pl.Boolean, + "BINARY": pl.Binary, + "VARIANT": pl.String, + "OBJECT": pl.String, + "ARRAY": pl.String, + "GEOGRAPHY": pl.String, + "GEOMETRY": pl.String, + } + return simple.get(type_name) + + def _describe_columns(self, uri: str, query: str): + """Cursor-describe metadata without executing; None if the query can't be described.""" + try: + con = self._connect(uri) + try: + cur = con.cursor() + try: + return cur.describe(query) + finally: + cur.close() + finally: + con.close() + except Exception as exc: + logger.debug("Snowflake describe failed: %s", exc) + return None + + def read( + self, + query: str, + uri: str, + logger: logging.Logger, + cancel_check: Callable[[], bool] | None = None, + ) -> pl.DataFrame: + import polars as pl + + con = self._connect(uri) + try: + cur = con.cursor() + try: + cur.execute(query) + description = cur.description or [] + if hasattr(cur, "fetch_arrow_all"): + table = cur.fetch_arrow_all() + df = pl.from_arrow(table) if table is not None else None + else: + # No Arrow support on this cursor (e.g. fakesnow): row-based fallback. + names = [col.name for col in description] + df = pl.DataFrame(cur.fetchall(), schema=names, orient="row", strict=False) + finally: + cur.close() + finally: + con.close() + + if df is None: + schema = {col.name: (self._polars_dtype(col) or pl.String) for col in description if col.name} + return pl.DataFrame(schema=schema) + overrides = { + col.name: dtype + for col in description + if col.name and col.name in df.columns and (dtype := self._polars_dtype(col)) is not None + } + return df.cast(overrides) if overrides else df + + def query_schema(self, uri: str, query: str) -> pl.Schema | None: + import polars as pl + + columns = self._describe_columns(uri, query) + if not columns: + return None + dtypes: dict[str, pl.DataType] = {} + for column in columns: + dtype = self._polars_dtype(column) if column.name else None + if dtype is None: + return None + dtypes[column.name] = dtype + return pl.Schema(dtypes) + + def table_schema(self, uri: str, table_name: str, schema_name: str | None) -> pl.Schema | None: + parts = [p for p in ([schema_name] if schema_name else []) + table_name.split(".") if p] + qualified = ".".join(_ident(p) for p in parts) + return self.query_schema(uri, f"SELECT * FROM {qualified}") + + @staticmethod + def _snowflake_type(dtype) -> str: + import polars as pl + + if isinstance(dtype, pl.Decimal): + return f"NUMBER({dtype.precision or 38},{dtype.scale or 0})" + if dtype.is_integer(): + return "BIGINT" + if dtype.is_float(): + return "DOUBLE" + if dtype == pl.Boolean: + return "BOOLEAN" + if dtype == pl.Date: + return "DATE" + if dtype == pl.Time: + return "TIME" + if isinstance(dtype, pl.Datetime): + return "TIMESTAMP_TZ" if dtype.time_zone else "TIMESTAMP_NTZ" + if dtype == pl.Binary: + return "BINARY" + return "VARCHAR" + + @staticmethod + def _table_exists(cur, schema_name: str | None, table_name: str) -> bool: + if schema_name: + cur.execute( + "SELECT 1 FROM information_schema.tables " + "WHERE UPPER(table_name) = UPPER(?) AND UPPER(table_schema) = UPPER(?) LIMIT 1", + (table_name, schema_name), + ) + else: + cur.execute( + "SELECT 1 FROM information_schema.tables " + "WHERE UPPER(table_name) = UPPER(?) AND table_schema = CURRENT_SCHEMA() LIMIT 1", + (table_name,), + ) + return cur.fetchone() is not None + + def write(self, df: pl.DataFrame, *, uri: str, table_name: str, if_exists: str = "append") -> None: + from shared.db_writer import _text_encoders + + schema_name, _, bare_name = table_name.rpartition(".") + con = self._connect(uri) + try: + cur = con.cursor() + try: + if schema_name: + cur.execute(f"CREATE SCHEMA IF NOT EXISTS {_ident(schema_name)}") + qualified = ".".join(_ident(p) for p in ([schema_name] if schema_name else []) + [bare_name]) + if if_exists == "fail" and self._table_exists(cur, schema_name or None, bare_name): + raise ValueError(f"Table '{table_name}' already exists") + column_defs = ", ".join( + f"{_quote_ident(name)} {self._snowflake_type(dtype)}" for name, dtype in df.schema.items() + ) + create = "CREATE OR REPLACE TABLE" if if_exists == "replace" else "CREATE TABLE IF NOT EXISTS" + cur.execute(f"{create} {qualified} ({column_defs})") + + encoders = _text_encoders(df) + encoder_idx = {df.columns.index(name): encode for name, encode in encoders.items()} + column_list = ", ".join(_quote_ident(name) for name in df.columns) + placeholders = ", ".join(["?"] * df.width) + insert = f"INSERT INTO {qualified} ({column_list}) VALUES ({placeholders})" + for batch in df.iter_slices(16_384): + rows = [ + tuple( + encoder_idx[i](value) if i in encoder_idx and value is not None else value + for i, value in enumerate(row) + ) + for row in batch.rows() + ] + if rows: + cur.executemany(insert, rows) + finally: + cur.close() + finally: + con.close() + + def _execute_rows(self, uri: str, sql: str, params: tuple = ()) -> list[tuple]: + con = self._connect(uri) + try: + cur = con.cursor() + try: + cur.execute(sql, params or None) + return cur.fetchall() + finally: + cur.close() + finally: + con.close() + + def list_schemas(self, uri: str) -> list[str] | None: + rows = self._execute_rows( + uri, + "SELECT schema_name FROM information_schema.schemata " + "WHERE schema_name <> 'INFORMATION_SCHEMA' ORDER BY schema_name", + ) + return [row[0] for row in rows] + + def list_tables(self, uri: str, schema_name: str | None) -> list[str] | None: + if schema_name: + rows = self._execute_rows( + uri, + "SELECT table_name FROM information_schema.tables " + "WHERE UPPER(table_schema) = UPPER(?) ORDER BY table_name", + (schema_name,), + ) + return [row[0] for row in rows] + rows = self._execute_rows( + uri, + "SELECT table_schema, table_name FROM information_schema.tables " + "WHERE table_schema <> 'INFORMATION_SCHEMA' ORDER BY table_schema, table_name", + ) + return [f"{schema}.{name}" for schema, name in rows] diff --git a/shared/tests/db_dialects/test_dialect_contract.py b/shared/tests/db_dialects/test_dialect_contract.py index 2a0d0691b..b1c148869 100644 --- a/shared/tests/db_dialects/test_dialect_contract.py +++ b/shared/tests/db_dialects/test_dialect_contract.py @@ -14,6 +14,7 @@ from shared.db_dialects import ( KNOWN_DIALECT_NAMES, + DbDialect, DialectInfo, GenericDialect, UnknownDialectError, @@ -81,7 +82,13 @@ def test_build_uri(dialect): with pytest.raises(ValueError): dialect.build_uri(host=None) uri = dialect.build_uri(host="h", port=dialect.default_port, username="u", password="p", database="d") - assert uri == f"{dialect.uri_scheme}://u:p@h:{dialect.default_port}/d" + if type(dialect).build_uri is DbDialect.build_uri: + assert uri == f"{dialect.uri_scheme}://u:p@h:{dialect.default_port}/d" + else: + # Custom URI shapes (e.g. Snowflake's port-less account locator) assert form, not bytes. + assert uri.startswith(f"{dialect.uri_scheme}://") + assert "u:p@" in uri + assert "/d" in uri @pytest.mark.parametrize("dialect", DIALECTS, ids=_ids) diff --git a/shared/tests/db_dialects/test_snowflake_dialect.py b/shared/tests/db_dialects/test_snowflake_dialect.py new file mode 100644 index 000000000..7c013fb1e --- /dev/null +++ b/shared/tests/db_dialects/test_snowflake_dialect.py @@ -0,0 +1,181 @@ +"""Snowflake dialect unit tests: account-URI shapes, type map, fakesnow round trip. + +The generic contract legs (registry sanity, URI shape, sqlglot-parseable limit +queries, catalog serialization) also run automatically for snowflake via +test_dialect_contract.py. The behavioral legs here run against fakesnow +(duckdb-backed connector emulation) so CI needs no Snowflake account; live +coverage (reads, writes, fast-schema parity against a real account) lives in +flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py, gated +on FLOWFILE_TEST_SNOWFLAKE_* credentials. +""" + +import logging +from types import SimpleNamespace + +import polars as pl +import pytest +import sqlglot + +from shared.db_dialects import get_dialect + +dialect = get_dialect("snowflake") +logger = logging.getLogger(__name__) + + +def _meta(type_code, precision=None, scale=None, name="c"): + return SimpleNamespace(name=name, type_code=type_code, precision=precision, scale=scale) + + +def _type_code(type_name: str) -> int: + from snowflake.connector.constants import FIELD_ID_TO_NAME + + return next(code for code, name in FIELD_ID_TO_NAME.items() if name == type_name) + + +def test_metadata(): + assert dialect.display_name == "Snowflake" + assert dialect.file_based is False + assert dialect.default_port == 443, "catalog metadata only; the account URI carries no port" + assert dialect.sqlalchemy_driver is None + assert dialect.sqlglot_name == "snowflake" + assert dialect.install_hint == "pip install snowflake-connector-python" + assert dialect.is_available() is True, "snowflake-connector-python is a main dependency; must be available" + assert [f.name for f in dialect.extra_fields] == ["account", "warehouse", "role"] + assert dialect.extra_fields[0].required is True + assert dialect.hidden_fields == ("host", "port", "ssl") + + +def test_build_uri_account_shape(): + uri = dialect.build_uri( + username="u", password="p", database="d", account="my-org-acct", warehouse="WH", role="R" + ) + assert uri == "snowflake://u:p@my-org-acct/d?warehouse=WH&role=R" + + +def test_build_uri_requires_account(): + with pytest.raises(ValueError, match="Account is required"): + dialect.build_uri(username="u", password="p", database="d") + + +def test_build_uri_accepts_host_as_account_alias(): + assert dialect.build_uri(host="acct", username="u", password="p", database="d") == "snowflake://u:p@acct/d" + + +def test_build_uri_drops_blocked_extra_params(): + uri = dialect.build_uri( + username="u", + password="p", + database="d", + account="acct", + warehouse="WH", + authenticator="externalbrowser", + private_key_file="/tmp/key.p8", + user="EVIL", + token="t", + ) + assert uri == "snowflake://u:p@acct/d?warehouse=WH" + + +def test_build_uri_quotes_credentials(): + uri = dialect.build_uri(username="u@corp", password="p@ss word", account="acct", database="d") + assert uri == "snowflake://u%40corp:p%40ss+word@acct/d" + + +def test_connect_kwargs_round_trip(): + uri = dialect.build_uri( + username="u@corp", password="p@ss", database="db1", account="Acct-Id", warehouse="WH", role="R" + ) + assert dialect._connect_kwargs(uri) == { + "account": "Acct-Id", + "user": "u@corp", + "password": "p@ss", + "database": "db1", + "warehouse": "WH", + "role": "R", + } + + +def test_limit_query_is_plain_limit_and_parses_as_snowflake(): + limited = dialect.limit_query("SELECT a, b FROM some_table", 5) + assert limited == "SELECT a, b FROM some_table LIMIT 5" + parsed = sqlglot.parse(limited, read=dialect.sqlglot_name) + assert parsed and parsed[0] is not None + + +def test_polars_dtype_maps_snowflake_types(): + fixed = _type_code("FIXED") + assert dialect._polars_dtype(_meta(fixed, 38, 0)) == pl.Int64, "all Snowflake ints report NUMBER(38,0)" + assert dialect._polars_dtype(_meta(fixed, 10, 2)) == pl.Decimal(10, 2) + assert dialect._polars_dtype(_meta(_type_code("REAL"))) == pl.Float64 + assert dialect._polars_dtype(_meta(_type_code("TEXT"))) == pl.String + assert dialect._polars_dtype(_meta(_type_code("DATE"))) == pl.Date + assert dialect._polars_dtype(_meta(_type_code("TIMESTAMP_NTZ"))) == pl.Datetime("us") + assert dialect._polars_dtype(_meta(_type_code("TIMESTAMP_TZ"))) == pl.Datetime("us", "UTC") + assert dialect._polars_dtype(_meta(_type_code("BOOLEAN"))) == pl.Boolean + assert dialect._polars_dtype(_meta(_type_code("VARIANT"))) == pl.String, "semi-structured reads as JSON text" + assert dialect._polars_dtype(_meta(_type_code("ARRAY"))) == pl.String + assert dialect._polars_dtype(_meta(9999)) is None, "unmapped types must disable prediction, not guess" + + +def test_snowflake_type_ddl_map(): + assert dialect._snowflake_type(pl.Int64) == "BIGINT" + assert dialect._snowflake_type(pl.Decimal(10, 2)) == "NUMBER(10,2)" + assert dialect._snowflake_type(pl.Float64) == "DOUBLE" + assert dialect._snowflake_type(pl.Datetime("us")) == "TIMESTAMP_NTZ" + assert dialect._snowflake_type(pl.Datetime("us", "UTC")) == "TIMESTAMP_TZ" + assert dialect._snowflake_type(pl.List(pl.Int64)) == "VARCHAR", "nested columns are JSON-encoded text" + + +@pytest.fixture() +def snow_uri(): + fakesnow = pytest.importorskip("fakesnow") + with fakesnow.patch(): + yield dialect.build_uri( + username="u", password="p", database="db1", account="test", **{"schema": "public"} + ) + + +def test_fakesnow_write_read_roundtrip(snow_uri): + df = pl.DataFrame({"id": [1, 2, 3], "score": [0.5, 1.5, 2.5], "label": ["a", "b", "c"]}) + + dialect.write(df, uri=snow_uri, table_name="t", if_exists="replace") + dialect.write(df, uri=snow_uri, table_name="t", if_exists="append") + with pytest.raises(ValueError, match="Table 't' already exists"): + dialect.write(df, uri=snow_uri, table_name="t", if_exists="fail") + dialect.write(df, uri=snow_uri, table_name="t", if_exists="replace") + + result = dialect.read("SELECT * FROM t", snow_uri, logger) + assert result.height == df.height + assert result.columns == df.columns, "quoted column identifiers must round-trip exact names" + + +def test_fakesnow_fast_schema_parity(snow_uri): + df = pl.DataFrame({"id": [1, 2], "score": [0.5, 1.5], "label": ["a", "b"]}) + dialect.write(df, uri=snow_uri, table_name="t", if_exists="replace") + + result = dialect.read("SELECT * FROM t", snow_uri, logger) + predicted = dialect.query_schema(snow_uri, "SELECT * FROM t") + if predicted is None: + pytest.skip("fakesnow does not support cursor.describe for this query") + assert dict(predicted) == dict(result.schema) + predicted_table = dialect.table_schema(snow_uri, "t", None) + assert predicted_table is not None and dict(predicted_table) == dict(result.schema) + + +def test_fakesnow_schema_qualified_write(snow_uri): + df = pl.DataFrame({"x": [1, 2]}) + dialect.write(df, uri=snow_uri, table_name="reporting.facts", if_exists="replace") + result = dialect.read("SELECT * FROM reporting.facts", snow_uri, logger) + assert result.height == 2 + + +def test_fakesnow_browse(snow_uri): + df = pl.DataFrame({"x": [1]}) + dialect.write(df, uri=snow_uri, table_name="t_browse", if_exists="replace") + + schemas = dialect.list_schemas(snow_uri) + assert schemas is not None and any(s.lower() == "public" for s in schemas) + tables = dialect.list_tables(snow_uri, "public") + assert tables is not None and any(t.lower() == "t_browse" for t in tables) + qualified = dialect.list_tables(snow_uri, None) + assert qualified is not None and any(t.lower() == "public.t_browse" for t in qualified) From 779aaa7a57deb0ace88c2fbbb48887b03fe1e370 Mon Sep 17 00:00:00 2001 From: edwardvaneechoud Date: Wed, 5 Aug 2026 20:24:47 +0200 Subject: [PATCH 2/4] feat: add SQL Server support with command-line utilities and fixtures --- .../python-api/reference/reading-data.md | 18 ++ docs/users/visual-editor/connections.md | 8 + .../031_database_connection_key_pair.py | 45 +++ .../flowfile_core/database/models.py | 3 + .../db_connections.py | 87 +++++- .../flowfile_core/flowfile/flow_graph.py | 101 ++++-- .../external_sources/sql_source/models.py | 51 ++- .../external_sources/sql_source/sql_source.py | 41 ++- .../external_sources/sql_source/utils.py | 11 +- .../flowfile_core/project/importer.py | 26 +- .../flowfile_core/project/manifest_entries.py | 3 + .../flowfile_core/project/projection.py | 20 +- flowfile_core/flowfile_core/routes/routes.py | 20 +- .../flowfile_core/schemas/input_schema.py | 61 +++- .../external_sources/test_snowflake_source.py | 290 +++++++++++++++++- flowfile_core/tests/project/test_roundtrip.py | 72 +++++ .../tests/test_db_dialect_endpoints.py | 95 ++++++ flowfile_core/tests/test_migration.py | 101 ++++++ .../database/connection_manager.py | 46 ++- .../database/connection_manager.pyi | 4 +- .../src/renderer/app/api/dbDialects.test.ts | 12 + .../src/renderer/app/api/dbDialects.ts | 8 + .../DatabaseConnectionSettings.vue | 80 ++++- .../databaseReader/DatabaseReader.vue | 3 +- .../databaseWriter/DatabaseWriter.vue | 3 +- .../renderer/app/composables/useDbDialects.ts | 5 +- .../src/renderer/app/types/node.types.ts | 3 + .../DatabaseConnectionSettings.vue | 106 ++++++- .../app/views/DatabaseView/DatabaseView.vue | 3 + .../renderer/app/views/DatabaseView/api.ts | 7 + .../DatabaseView/databaseConnectionTypes.ts | 8 + .../external_sources/sql_source/models.py | 26 +- .../external_sources/test_dialect_ports.py | 51 +++ shared/CLAUDE.md | 2 +- shared/db_dialects/__init__.py | 2 + shared/db_dialects/base.py | 16 + shared/db_dialects/builtin.py | 3 +- shared/db_dialects/duckdb.py | 3 +- shared/db_dialects/snowflake.py | 70 ++++- shared/sql_utils.py | 10 + .../db_dialects/test_dialect_contract.py | 1 + .../db_dialects/test_snowflake_dialect.py | 127 ++++++++ 42 files changed, 1553 insertions(+), 99 deletions(-) create mode 100644 flowfile_core/flowfile_core/alembic/versions/031_database_connection_key_pair.py diff --git a/docs/users/python-api/reference/reading-data.md b/docs/users/python-api/reference/reading-data.md index d849257ed..f366fb1a2 100644 --- a/docs/users/python-api/reference/reading-data.md +++ b/docs/users/python-api/reference/reading-data.md @@ -419,6 +419,24 @@ ff.create_database_connection( df = ff.read_database("analytics-snowflake", schema_name="PUBLIC", table_name="EVENTS") ``` +For key-pair (JWT) authentication — Snowflake's recommended method for programmatic +access — pass `auth_method="key_pair"` with the private key PEM *text* (never a file +path; read the file yourself). Add `private_key_passphrase` when the PEM is encrypted: + +```python +ff.create_database_connection( + connection_name="analytics-snowflake-kp", + database_type="snowflake", + database="ANALYTICS", + username="svc_user", + auth_method="key_pair", + private_key=open("rsa_key.p8").read(), + extra_params={"account": "myorg-myaccount", "warehouse": "COMPUTE_WH"}, +) +``` + +The key is stored as an encrypted secret, exactly like a password. + Semi-structured columns (`VARIANT`, `OBJECT`, `ARRAY`) are read as JSON text. The tested Snowflake example reads a table and a query through a stored connection (it runs only when Snowflake test credentials are configured): diff --git a/docs/users/visual-editor/connections.md b/docs/users/visual-editor/connections.md index 0e0ff8e40..50cade0c6 100644 --- a/docs/users/visual-editor/connections.md +++ b/docs/users/visual-editor/connections.md @@ -47,6 +47,14 @@ and Cloud Storage Writer nodes without re-entering credentials each time. usual username, password, and database. Connections always use TLS, so there is no SSL toggle. Semi-structured columns (`VARIANT`, `OBJECT`, `ARRAY`) are read as JSON text. + Snowflake also supports **key-pair (JWT) authentication** — Snowflake's recommended + method for programmatic access now that password-only logins are being phased out. + Pick *Key pair (JWT)* in the **Authentication Method** selector, then paste the + private key PEM text into the key field (plus its passphrase when the key is + encrypted). The key is stored as an encrypted secret, exactly like a password, and + is never written back to the form when editing — leave the field blank to keep the + existing key. + ### Creating a Database Connection 1. Open the **Connections** page from the left sidebar and select the **Database** tab diff --git a/flowfile_core/flowfile_core/alembic/versions/031_database_connection_key_pair.py b/flowfile_core/flowfile_core/alembic/versions/031_database_connection_key_pair.py new file mode 100644 index 000000000..20130368d --- /dev/null +++ b/flowfile_core/flowfile_core/alembic/versions/031_database_connection_key_pair.py @@ -0,0 +1,45 @@ +"""Add key-pair auth columns to database_connections. + +auth_method selects the connection's authentication method (NULL means +password, the historical default). private_key_id / private_key_passphrase_id +reference encrypted secrets holding the PEM text and its optional passphrase. +The id columns are plain integers here — the ForeignKey lives in the ORM only, +because SQLite cannot drop a column that carries a FK constraint on downgrade. + +Revision ID: 031 +Revises: 030 +Create Date: 2026-08-05 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import inspect + +revision: str = "031" +down_revision: str | None = "030" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_COLUMNS = ( + ("auth_method", sa.String), + ("private_key_id", sa.Integer), + ("private_key_passphrase_id", sa.Integer), +) + + +def _has_column(table: str, column: str) -> bool: + return column in {c["name"] for c in inspect(op.get_bind()).get_columns(table)} + + +def upgrade() -> None: + for name, column_type in _COLUMNS: + if not _has_column("database_connections", name): + op.add_column("database_connections", sa.Column(name, column_type, nullable=True)) + + +def downgrade() -> None: + for name, _ in _COLUMNS: + if _has_column("database_connections", name): + op.drop_column("database_connections", name) diff --git a/flowfile_core/flowfile_core/database/models.py b/flowfile_core/flowfile_core/database/models.py index 05e558645..2b9078d88 100644 --- a/flowfile_core/flowfile_core/database/models.py +++ b/flowfile_core/flowfile_core/database/models.py @@ -59,7 +59,10 @@ class DatabaseConnection(Base): database = Column(String, default=None) ssl_enabled = Column(Boolean, default=False) extra_params = Column(Text, nullable=True) # JSON dict of dialect-specific params (e.g. snowflake account) + auth_method = Column(String, nullable=True) # NULL == password (the historical default) password_id = Column(Integer, ForeignKey("secrets.id")) + private_key_id = Column(Integer, ForeignKey("secrets.id"), nullable=True) + private_key_passphrase_id = Column(Integer, ForeignKey("secrets.id"), nullable=True) user_id = Column(Integer, ForeignKey("users.id")) diff --git a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py index af5de49f1..7bab39f3c 100644 --- a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py +++ b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py @@ -1,5 +1,6 @@ import json +from pydantic import SecretStr from sqlalchemy.orm import Session from flowfile_core.auth import sharing @@ -50,7 +51,26 @@ def store_database_connection(db: Session, connection: FullDatabaseConnection, u f" Please use a unique connection name or delete the existing connection first." ) + if connection.auth_method == "key_pair" and connection.private_key is None: + raise ValueError("key_pair authentication requires a private key when creating a connection.") + password_id = store_secret(db, SecretInput(name=connection.connection_name, value=connection.password), user_id).id + # `is not None` (not truthiness): a project import may pass an empty placeholder + # value that must still create a linked secret row for the refill UI. + private_key_id = None + if connection.private_key is not None: + private_key_id = store_secret( + db, SecretInput(name=f"{connection.connection_name}_private_key", value=connection.private_key), user_id + ).id + private_key_passphrase_id = None + if connection.private_key_passphrase is not None: + private_key_passphrase_id = store_secret( + db, + SecretInput( + name=f"{connection.connection_name}_private_key_passphrase", value=connection.private_key_passphrase + ), + user_id, + ).id db_connection = DBConnectionModel( connection_name=connection.connection_name, @@ -60,6 +80,9 @@ def store_database_connection(db: Session, connection: FullDatabaseConnection, u database_type=connection.database_type, username=connection.username, password_id=password_id, + private_key_id=private_key_id, + private_key_passphrase_id=private_key_passphrase_id, + auth_method=connection.auth_method, ssl_enabled=connection.ssl_enabled, extra_params=_dump_extra_params(connection.extra_params), user_id=user_id, @@ -92,6 +115,7 @@ def update_database_connection(db: Session, connection: FullDatabaseConnection, db_connection.username = connection.username db_connection.ssl_enabled = connection.ssl_enabled db_connection.extra_params = _dump_extra_params(connection.extra_params) + db_connection.auth_method = connection.auth_method password_value = connection.password.get_secret_value() if password_value: @@ -103,6 +127,40 @@ def update_database_connection(db: Session, connection: FullDatabaseConnection, new_secret = store_secret(db, secret_input, user_id) db_connection.password_id = new_secret.id + incoming_key = connection.private_key.get_secret_value() if connection.private_key else "" + keeps_key_material = connection.auth_method == "key_pair" or bool(incoming_key) + if keeps_key_material: + if connection.auth_method == "key_pair" and not incoming_key and db_connection.private_key_id is None: + raise ValueError("key_pair authentication requires a private key.") + # Rotate key material only when a non-empty value arrived (empty == keep existing, + # mirroring password). _update_cloud_secret is generic despite the name. + db_connection.private_key_id = _update_cloud_secret( + db, + db_connection.private_key_id, + incoming_key, + f"{connection.connection_name}_private_key", + user_id, + ) + db_connection.private_key_passphrase_id = _update_cloud_secret( + db, + db_connection.private_key_passphrase_id, + connection.private_key_passphrase.get_secret_value() if connection.private_key_passphrase else "", + f"{connection.connection_name}_private_key_passphrase", + user_id, + ) + else: + # Switching away from key-pair auth: detach AND delete the key secrets, or a + # rotated-away (possibly compromised) key would keep authenticating silently. + stale_ids = [ + secret_id + for secret_id in (db_connection.private_key_id, db_connection.private_key_passphrase_id) + if secret_id is not None + ] + db_connection.private_key_id = None + db_connection.private_key_passphrase_id = None + if stale_ids: + db.query(Secret).filter(Secret.id.in_(stale_ids)).delete(synchronize_session=False) + db.commit() db.refresh(db_connection) _project_sync_connection("database", connection.connection_name, user_id) @@ -181,6 +239,14 @@ def get_database_connection_schema(db: Session, connection_name: str, user_id: i if not password_secret: raise Exception("Password secret not found") + def _ciphertext(secret_id: int | None) -> str | None: + if secret_id is None: + return None + secret = db.query(Secret).filter(Secret.id == secret_id).first() + return secret.encrypted_value if secret else None + + private_key = _ciphertext(db_connection.private_key_id) + private_key_passphrase = _ciphertext(db_connection.private_key_passphrase_id) return FullDatabaseConnection( connection_name=db_connection.connection_name, host=db_connection.host, @@ -191,6 +257,9 @@ def get_database_connection_schema(db: Session, connection_name: str, user_id: i password=password_secret.encrypted_value, ssl_enabled=db_connection.ssl_enabled, extra_params=parse_extra_params(db_connection.extra_params), + auth_method=db_connection.auth_method, + private_key=SecretStr(private_key) if private_key is not None else None, + private_key_passphrase=SecretStr(private_key_passphrase) if private_key_passphrase is not None else None, ) return None @@ -222,12 +291,21 @@ def delete_database_connection(db: Session, connection_name: str, user_id: int) db_connection = _get_own_database_connection(db, connection_name, user_id) if db_connection: + # Collect secret ids before the delete is staged (cloud-connection ordering). + secret_ids_to_delete = [ + secret_id + for secret_id in ( + db_connection.password_id, + db_connection.private_key_id, + db_connection.private_key_passphrase_id, + ) + if secret_id is not None + ] + sharing.delete_grants_for_resource(db, "database_connection", db_connection.id) db.delete(db_connection) - - password_secret = db.query(Secret).filter(Secret.id == db_connection.password_id).first() - if password_secret: - db.delete(password_secret) + if secret_ids_to_delete: + db.query(Secret).filter(Secret.id.in_(secret_ids_to_delete)).delete(synchronize_session=False) db.commit() _project_sync_connection("database", connection_name, user_id, deleted=True) @@ -248,6 +326,7 @@ def database_connection_interface_from_db_connection( database=db_connection.database, ssl_enabled=db_connection.ssl_enabled, extra_params=parse_extra_params(db_connection.extra_params), + auth_method=db_connection.auth_method, id=db_connection.id, access=access, ) diff --git a/flowfile_core/flowfile_core/flowfile/flow_graph.py b/flowfile_core/flowfile_core/flowfile/flow_graph.py index 5e12e9979..9e9c78824 100644 --- a/flowfile_core/flowfile_core/flowfile/flow_graph.py +++ b/flowfile_core/flowfile_core/flowfile/flow_graph.py @@ -1316,15 +1316,24 @@ def _handle_physical_table_write( return df +class ResolvedDatabaseCredentials(NamedTuple): + """A resolved connection plus its credential ciphertexts (``$ffsec$``, decrypted at point of use).""" + + connection: Any + password: str | None + private_key: str | None + private_key_passphrase: str | None + reference_settings: input_schema.FullDatabaseConnection | None + + def _resolve_database_credentials( database_settings, user_id: int, -) -> tuple: - """Resolve database connection and encrypted password from settings. +) -> ResolvedDatabaseCredentials: + """Resolve database connection and encrypted credentials from settings. - Returns: - (database_connection, encrypted_password, database_reference_settings) - where database_reference_settings is the stored connection (or None for inline). + For key-pair auth the private key is the required credential and the password + becomes optional; ``reference_settings`` is the stored connection (None for inline). """ is_file_based = ( database_settings.connection_mode == "inline" @@ -1333,12 +1342,33 @@ def _resolve_database_credentials( ) if database_settings.connection_mode == "inline" and not is_file_based: database_connection = database_settings.database_connection - encrypted_password = get_encrypted_secret(current_user_id=user_id, secret_name=database_connection.password_ref) - if encrypted_password is None: + use_key_pair = database_connection.auth_method == "key_pair" + encrypted_password = None + if database_connection.password_ref: + encrypted_password = get_encrypted_secret( + current_user_id=user_id, secret_name=database_connection.password_ref + ) + if encrypted_password is None and not use_key_pair: raise HTTPException(status_code=400, detail="Password not found") - return database_connection, encrypted_password, None + encrypted_private_key = None + encrypted_passphrase = None + if use_key_pair: + encrypted_private_key = get_encrypted_secret( + current_user_id=user_id, secret_name=database_connection.private_key_ref + ) + if encrypted_private_key is None: + raise HTTPException(status_code=400, detail="Private key secret not found") + if database_connection.private_key_passphrase_ref: + encrypted_passphrase = get_encrypted_secret( + current_user_id=user_id, secret_name=database_connection.private_key_passphrase_ref + ) + if encrypted_passphrase is None: + raise HTTPException(status_code=400, detail="Private key passphrase secret not found") + return ResolvedDatabaseCredentials( + database_connection, encrypted_password, encrypted_private_key, encrypted_passphrase, None + ) elif is_file_based: - return database_settings.database_connection, None, None + return ResolvedDatabaseCredentials(database_settings.database_connection, None, None, None, None) else: ref_settings = get_local_database_connection(database_settings.database_connection_name, user_id) if ref_settings is None: @@ -1349,8 +1379,13 @@ def _resolve_database_credentials( "or not accessible for this user" ), ) - encrypted_password = ref_settings.password.get_secret_value() - return ref_settings, encrypted_password, ref_settings + return ResolvedDatabaseCredentials( + ref_settings, + ref_settings.password.get_secret_value(), + ref_settings.private_key.get_secret_value() if ref_settings.private_key else None, + ref_settings.private_key_passphrase.get_secret_value() if ref_settings.private_key_passphrase else None, + ref_settings, + ) class _FlowIdentity(NamedTuple): @@ -4570,9 +4605,8 @@ def add_database_writer(self, node_database_writer: input_schema.NodeDatabaseWri database_settings: input_schema.DatabaseWriteSettings = node_database_writer.database_write_settings def _func(df: FlowDataEngine): - database_connection, encrypted_password, database_reference_settings = _resolve_database_credentials( - database_settings, node_database_writer.user_id - ) + creds = _resolve_database_credentials(database_settings, node_database_writer.user_id) + database_connection = creds.connection df.lazy = True table_name = ( database_settings.schema_name + "." + database_settings.table_name @@ -4589,9 +4623,14 @@ def _func(df: FlowDataEngine): port=database_connection.port, database=database_connection.database, username=database_connection.username, - password=decrypt_secret(encrypted_password) if encrypted_password else None, + password=decrypt_secret(creds.password) if creds.password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + auth_method=database_connection.auth_method, + private_key=decrypt_secret(creds.private_key) if creds.private_key else None, + private_key_passphrase=( + decrypt_secret(creds.private_key_passphrase) if creds.private_key_passphrase else None + ), **(database_connection.extra_params or {}), ), table_name=table_name, @@ -4602,12 +4641,14 @@ def _func(df: FlowDataEngine): database_external_write_settings = ( sql_models.DatabaseExternalWriteSettings.create_from_from_node_database_writer( node_database_writer=node_database_writer, - password=encrypted_password, + password=creds.password, table_name=table_name, database_reference_settings=( - database_reference_settings if database_settings.connection_mode == "reference" else None + creds.reference_settings if database_settings.connection_mode == "reference" else None ), lf=df.data_frame, + private_key=creds.private_key, + private_key_passphrase=creds.private_key_passphrase, ) ) external_database_writer = ExternalDatabaseWriter( @@ -4660,7 +4701,8 @@ def _get_creds(): return _creds["v"] def _func(): - database_connection, encrypted_password, database_reference_settings = _get_creds() + creds = _get_creds() + database_connection = creds.connection sql_source = BaseSqlSource( query=None if database_settings.query_mode == "table" else database_settings.query, table_name=database_settings.table_name, @@ -4678,9 +4720,14 @@ def _func(): port=database_connection.port, database=database_connection.database, username=database_connection.username, - password=decrypt_secret(encrypted_password) if encrypted_password else None, + password=decrypt_secret(creds.password) if creds.password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + auth_method=database_connection.auth_method, + private_key=decrypt_secret(creds.private_key) if creds.private_key else None, + private_key_passphrase=( + decrypt_secret(creds.private_key_passphrase) if creds.private_key_passphrase else None + ), **(database_connection.extra_params or {}), ), query=None if database_settings.query_mode == "table" else database_settings.query, @@ -4698,11 +4745,13 @@ def _func(): database_external_read_settings = ( sql_models.DatabaseExternalReadSettings.create_from_from_node_database_reader( node_database_reader=node_database_reader, - password=encrypted_password, + password=creds.password, query=sql_source.query, database_reference_settings=( - database_reference_settings if database_settings.connection_mode == "reference" else None + creds.reference_settings if database_settings.connection_mode == "reference" else None ), + private_key=creds.private_key, + private_key_passphrase=creds.private_key_passphrase, ) ) @@ -4720,7 +4769,8 @@ def schema_callback(): # when fields were never captured (failures here are caught per-node). if node_database_reader.fields: return [FlowfileColumn.from_input(f.name, f.data_type) for f in node_database_reader.fields] - database_connection, encrypted_password, _ = _get_creds() + creds = _get_creds() + database_connection = creds.connection sql_source = SqlSource( connection_string=sql_utils.construct_sql_uri( database_type=database_connection.database_type, @@ -4728,9 +4778,14 @@ def schema_callback(): port=database_connection.port, database=database_connection.database, username=database_connection.username, - password=decrypt_secret(encrypted_password) if encrypted_password else None, + password=decrypt_secret(creds.password) if creds.password else None, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + auth_method=database_connection.auth_method, + private_key=decrypt_secret(creds.private_key) if creds.private_key else None, + private_key_passphrase=( + decrypt_secret(creds.private_key_passphrase) if creds.private_key_passphrase else None + ), **(database_connection.extra_params or {}), ), query=None if database_settings.query_mode == "table" else database_settings.query, diff --git a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/models.py b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/models.py index fa3e0fb15..bf1a8bb2f 100644 --- a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/models.py +++ b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/models.py @@ -29,10 +29,19 @@ def _decode_bytes(v: Any) -> bytes: class ExtDatabaseConnection(DatabaseConnection): - """Database connection configuration with password handling.""" + """Database connection configuration with credential handling. + + password / private_key / private_key_passphrase carry $ffsec$ ciphertext + over the core -> worker wire; the worker decrypts them independently. + """ password: str | None = None ssl_enabled: bool | None = False + private_key: str | None = None + private_key_passphrase: str | None = None + + +_SECRET_FIELDS = ("password", "private_key", "private_key_passphrase") class DatabaseExternalWriteSettings(BaseModel): @@ -53,24 +62,36 @@ def create_from_from_node_database_writer( table_name: str, lf: pl.LazyFrame, database_reference_settings: FullDatabaseConnection = None, + private_key: str | None = None, + private_key_passphrase: str | None = None, ) -> "DatabaseExternalWriteSettings": """ Create DatabaseExternalWriteSettings from NodeDatabaseWriter. Args: node_database_writer (NodeDatabaseWriter): an instance of NodeDatabaseWriter - password (str): the password for the database connection + password (str): the encrypted password for the database connection table_name (str): the table name to be used for writing lf (pl.LazyFrame): the LazyFrame to be written to the database database_reference_settings (FullDatabaseConnection): optional database reference settings + private_key (str): the encrypted private key (key-pair auth) + private_key_passphrase (str): the encrypted private-key passphrase Returns: DatabaseExternalReadSettings: an instance of DatabaseExternalReadSettings """ if node_database_writer.database_write_settings.connection_mode == "inline": database_connection = node_database_writer.database_write_settings.database_connection.model_dump() else: - database_connection = {k: v for k, v in database_reference_settings.model_dump().items() if k != "password"} - - ext_database_connection = ExtDatabaseConnection(**database_connection, password=password) + # Exclude the SecretStr fields from the splat; the ciphertexts are passed explicitly. + database_connection = { + k: v for k, v in database_reference_settings.model_dump().items() if k not in _SECRET_FIELDS + } + + ext_database_connection = ExtDatabaseConnection( + **database_connection, + password=password, + private_key=private_key, + private_key_passphrase=private_key_passphrase, + ) return cls( connection=ext_database_connection, table_name=table_name, @@ -96,23 +117,35 @@ def create_from_from_node_database_reader( password: str, query: str, database_reference_settings: FullDatabaseConnection = None, + private_key: str | None = None, + private_key_passphrase: str | None = None, ) -> "DatabaseExternalReadSettings": """ Create DatabaseExternalReadSettings from NodeDatabaseReader. Args: node_database_reader (NodeDatabaseReader): an instance of NodeDatabaseReader - password (str): the password for the database connection + password (str): the encrypted password for the database connection query (str): the SQL query to be executed database_reference_settings (FullDatabaseConnection): optional database reference settings + private_key (str): the encrypted private key (key-pair auth) + private_key_passphrase (str): the encrypted private-key passphrase Returns: DatabaseExternalReadSettings: an instance of DatabaseExternalReadSettings """ if node_database_reader.database_settings.connection_mode == "inline": database_connection = node_database_reader.database_settings.database_connection.model_dump() else: - database_connection = {k: v for k, v in database_reference_settings.model_dump().items() if k != "password"} - - ext_database_connection = ExtDatabaseConnection(**database_connection, password=password) + # Exclude the SecretStr fields from the splat; the ciphertexts are passed explicitly. + database_connection = { + k: v for k, v in database_reference_settings.model_dump().items() if k not in _SECRET_FIELDS + } + + ext_database_connection = ExtDatabaseConnection( + **database_connection, + password=password, + private_key=private_key, + private_key_passphrase=private_key_passphrase, + ) return cls( connection=ext_database_connection, query=query, diff --git a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py index f516beaf8..198e8ecec 100644 --- a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py +++ b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/sql_source.py @@ -430,20 +430,42 @@ class ResolvedConnection(NamedTuple): def _resolve_connection(database_settings: DatabaseSettings, user_id: int) -> ResolvedConnection: """Resolve DatabaseSettings into a connection URI + dialect, handling inline/reference mode.""" database_connection = database_settings.database_connection + private_key = None + private_key_passphrase = None if database_settings.connection_mode == "inline": if database_connection is None: raise ValueError("Database connection is required in inline mode") is_file_based = get_dialect_or_generic(database_connection.database_type).file_based - if is_file_based: - password = None - else: + use_key_pair = database_connection.auth_method == "key_pair" + password = None + if not is_file_based and database_connection.password_ref: encrypted_secret = get_encrypted_secret( current_user_id=user_id, secret_name=database_connection.password_ref ) - if encrypted_secret is None: - raise ValueError(f"Secret with name {database_connection.password_ref} not found for user {user_id}") - password = decrypt_secret(encrypted_secret) + if encrypted_secret is not None: + password = decrypt_secret(encrypted_secret) + if password is None and not is_file_based and not use_key_pair: + raise ValueError(f"Secret with name {database_connection.password_ref} not found for user {user_id}") + if use_key_pair: + encrypted_key = get_encrypted_secret( + current_user_id=user_id, secret_name=database_connection.private_key_ref + ) + if encrypted_key is None: + raise ValueError( + f"Private key secret with name {database_connection.private_key_ref} not found for user {user_id}" + ) + private_key = decrypt_secret(encrypted_key) + if database_connection.private_key_passphrase_ref: + encrypted_passphrase = get_encrypted_secret( + current_user_id=user_id, secret_name=database_connection.private_key_passphrase_ref + ) + if encrypted_passphrase is None: + raise ValueError( + f"Private key passphrase secret with name " + f"{database_connection.private_key_passphrase_ref} not found for user {user_id}" + ) + private_key_passphrase = decrypt_secret(encrypted_passphrase) else: database_connection = get_local_database_connection(database_settings.database_connection_name, user_id) if database_connection is None: @@ -453,6 +475,10 @@ def _resolve_connection(database_settings: DatabaseSettings, user_id: int) -> Re ) encrypted_secret = database_connection.password.get_secret_value() password = decrypt_secret(encrypted_secret) + if database_connection.private_key: + private_key = decrypt_secret(database_connection.private_key.get_secret_value()) + if database_connection.private_key_passphrase: + private_key_passphrase = decrypt_secret(database_connection.private_key_passphrase.get_secret_value()) uri = construct_sql_uri( database_type=database_connection.database_type, @@ -463,6 +489,9 @@ def _resolve_connection(database_settings: DatabaseSettings, user_id: int) -> Re password=password, ssl_enabled=bool(getattr(database_connection, "ssl_enabled", False)), connect_timeout=10, + auth_method=database_connection.auth_method, + private_key=private_key, + private_key_passphrase=private_key_passphrase, **(database_connection.extra_params or {}), ) return ResolvedConnection(uri=uri, database_type=database_connection.database_type) diff --git a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/utils.py b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/utils.py index 567b0e84a..7962bb9e0 100644 --- a/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/utils.py +++ b/flowfile_core/flowfile_core/flowfile/sources/external_sources/sql_source/utils.py @@ -367,12 +367,15 @@ def construct_sql_uri( url: str | None = None, ssl_enabled: bool = False, connect_timeout: int | None = None, + auth_method: str | None = None, + private_key: SecretStr | None = None, + private_key_passphrase: SecretStr | None = None, **kwargs, ) -> str: """ Constructs a SQL URI string from the provided parameters. - Thin wrapper around shared.sql_utils.construct_sql_uri that unwraps SecretStr passwords. + Thin wrapper around shared.sql_utils.construct_sql_uri that unwraps SecretStr credentials. Args: database_type: Database type (postgresql, mysql, sqlite, etc.) @@ -384,6 +387,9 @@ def construct_sql_uri( url: Complete database URL (overrides other parameters if provided) ssl_enabled: Adds sslmode=require for postgres-family databases connect_timeout: Connection timeout in seconds (postgres-family only) + auth_method: Authentication method ("password"/None, or "key_pair" where supported) + private_key: Private key PEM text as SecretStr (key-pair auth) + private_key_passphrase: Optional passphrase for an encrypted private key **kwargs: Additional connection parameters Returns: @@ -403,6 +409,9 @@ def construct_sql_uri( url=url, ssl_enabled=ssl_enabled, connect_timeout=connect_timeout, + auth_method=auth_method, + private_key=private_key.get_secret_value() if private_key else None, + private_key_passphrase=private_key_passphrase.get_secret_value() if private_key_passphrase else None, **kwargs, ) diff --git a/flowfile_core/flowfile_core/project/importer.py b/flowfile_core/flowfile_core/project/importer.py index 55d209d68..0bfe8e818 100644 --- a/flowfile_core/flowfile_core/project/importer.py +++ b/flowfile_core/flowfile_core/project/importer.py @@ -170,6 +170,15 @@ def _import_db_connection(data: dict, owner_id: int, dotenv: dict, result: Setup name = entry.connection_name secret_name = placeholder_name(entry.password) or name value = _resolve_secret_value(secret_name, owner_id, dotenv, result) + + def _resolve_optional(raw: str | None) -> SecretStr | None: + # Only a placeholder resolves (no `or name` fallback like password): a None + # field means the connection has no such secret at all. + key_name = placeholder_name(raw) + if not key_name: + return None + return SecretStr(_resolve_secret_value(key_name, owner_id, dotenv, result)) + conn = FullDatabaseConnection( connection_name=name, database_type=entry.database_type, @@ -180,12 +189,19 @@ def _import_db_connection(data: dict, owner_id: int, dotenv: dict, result: Setup database=entry.database, ssl_enabled=entry.ssl_enabled, extra_params=entry.extra_params, + auth_method=entry.auth_method, + private_key=_resolve_optional(entry.private_key), + private_key_passphrase=_resolve_optional(entry.private_key_passphrase), ) - with get_db_context() as db: - if _get_own_database_connection(db, name, owner_id): - update_database_connection(db, conn, owner_id) - else: - store_database_connection(db, conn, owner_id) + try: + with get_db_context() as db: + if _get_own_database_connection(db, name, owner_id): + update_database_connection(db, conn, owner_id) + else: + store_database_connection(db, conn, owner_id) + except ValueError: + logger.warning("Project import: skipping invalid database connection %r", name, exc_info=True) + return None result.imported_connections += 1 return name diff --git a/flowfile_core/flowfile_core/project/manifest_entries.py b/flowfile_core/flowfile_core/project/manifest_entries.py index 57ee13a62..bad8cd09d 100644 --- a/flowfile_core/flowfile_core/project/manifest_entries.py +++ b/flowfile_core/flowfile_core/project/manifest_entries.py @@ -23,8 +23,11 @@ class DatabaseConnectionEntry(_Entry): database: str | None = None username: str = "" ssl_enabled: bool = False + auth_method: str | None = None extra_params: dict[str, str] | None = None password: str | None = None + private_key: str | None = None + private_key_passphrase: str | None = None class CloudConnectionEntry(_Entry): diff --git a/flowfile_core/flowfile_core/project/projection.py b/flowfile_core/flowfile_core/project/projection.py index 57c08d9da..6d2ee2292 100644 --- a/flowfile_core/flowfile_core/project/projection.py +++ b/flowfile_core/flowfile_core/project/projection.py @@ -63,7 +63,12 @@ "verify_ssl", ) # Non-secret database-connection fields that round-trip verbatim. -_DB_PLAIN_FIELDS = ("database_type", "host", "port", "database", "username", "ssl_enabled") +_DB_PLAIN_FIELDS = ("database_type", "host", "port", "database", "username", "ssl_enabled", "auth_method") +# Secret-backed database-connection fields: (file field, model FK column), like _CLOUD_SECRETS. +_DB_SECRETS = ( + ("private_key", "private_key_id"), + ("private_key_passphrase", "private_key_passphrase_id"), +) def _secret_name(db: Session, secret_id: int | None) -> str | None: @@ -191,6 +196,9 @@ def _db_connection_dict(db: Session, conn: DatabaseConnection) -> dict: d[f] = getattr(conn, f) d["extra_params"] = parse_extra_params(conn.extra_params) d["password"] = make_placeholder(secret_name) if secret_name else None + for field, fk in _DB_SECRETS: + name = _secret_name(db, getattr(conn, fk)) + d[field] = make_placeholder(name) if name else None return d @@ -283,9 +291,13 @@ def regenerate_secret_manifest(db: Session, root: Path, owner_id: int) -> None: """secrets.yaml lists only standalone secrets; connection secrets are implied by the connection files (and recreated by their store functions on import).""" linked: set[int] = set() - for (sid,) in db.query(DatabaseConnection.password_id).filter(DatabaseConnection.user_id == owner_id): - if sid: - linked.add(sid) + db_secret_fk_columns = ( + DatabaseConnection.password_id, + DatabaseConnection.private_key_id, + DatabaseConnection.private_key_passphrase_id, + ) + for row in db.query(*db_secret_fk_columns).filter(DatabaseConnection.user_id == owner_id): + linked.update(sid for sid in row if sid) for row in db.query(*[getattr(CloudStorageConnection, fk) for fk in _CLOUD_SECRET_FK_COLUMNS]).filter( CloudStorageConnection.user_id == owner_id ): diff --git a/flowfile_core/flowfile_core/routes/routes.py b/flowfile_core/flowfile_core/routes/routes.py index 119fd1284..53baf050a 100644 --- a/flowfile_core/flowfile_core/routes/routes.py +++ b/flowfile_core/flowfile_core/routes/routes.py @@ -760,8 +760,9 @@ def create_db_connection( _require_known_database_type(input_connection.database_type) try: store_database_connection(db, input_connection, current_user.id) - except ValueError: - raise HTTPException(422, "Connection name already exists") from None + except ValueError as e: + # Duplicate name or invalid credential shape (e.g. key_pair without a key). + raise HTTPException(422, str(e)) from None except Exception as e: logger.error(e) raise HTTPException(422, str(e)) from e @@ -791,16 +792,23 @@ def update_db_connection( # JSON against the incoming dict, normalized, so unchanged params don't trip the guard. if (parse_extra_params(db_connection.extra_params) or {}) != (input_connection.extra_params or {}): changed.append("extra_params") + # auth_method needs normalization: a stored NULL and an incoming "password" are the same. + if (db_connection.auth_method or "password") != (input_connection.auth_method or "password"): + changed.append("auth_method") require_credentials_on_target_change( changed, - has_new_credentials=bool(input_connection.password.get_secret_value()), - has_bundled_secrets=db_connection.password_id is not None, + has_new_credentials=bool(input_connection.password.get_secret_value()) + or bool(input_connection.private_key and input_connection.private_key.get_secret_value()), + has_bundled_secrets=db_connection.password_id is not None or db_connection.private_key_id is not None, ) try: # Owner's user_id keeps a rotated password encrypted under the OWNER's key. update_database_connection(db, input_connection, db_connection.user_id) - except ValueError: - raise HTTPException(404, "Database connection not found") from None + except ValueError as e: + if "not found" in str(e): + raise HTTPException(404, "Database connection not found") from None + # Invalid credential shape (e.g. switching to key_pair without a key). + raise HTTPException(422, str(e)) from None except Exception as e: logger.error(e) raise HTTPException(422, str(e)) from e diff --git a/flowfile_core/flowfile_core/schemas/input_schema.py b/flowfile_core/flowfile_core/schemas/input_schema.py index ea4f43791..d2c297819 100644 --- a/flowfile_core/flowfile_core/schemas/input_schema.py +++ b/flowfile_core/flowfile_core/schemas/input_schema.py @@ -929,6 +929,16 @@ def _validate_extra_params(v: dict[str, str] | None) -> dict[str, str] | None: return v +def _validate_auth_method(auth_method: str | None, database_type: str) -> None: + """Reject an auth method the connection's dialect does not support.""" + if auth_method and auth_method != "password": + from shared.db_dialects import get_dialect_or_generic + + dialect = get_dialect_or_generic(database_type) + if auth_method not in dialect.auth_methods: + raise ValueError(f"auth_method '{auth_method}' is not supported by database type '{database_type}'") + + class DatabaseConnection(BaseModel): """Defines the connection parameters for a database.""" @@ -940,6 +950,9 @@ class DatabaseConnection(BaseModel): database: str | None = None url: str | None = None extra_params: dict[str, str] | None = None + auth_method: str | None = None + private_key_ref: SecretRef | None = None + private_key_passphrase_ref: SecretRef | None = None @field_validator("database_type") @classmethod @@ -951,7 +964,7 @@ def known_database_type(cls, v: str) -> str: raise ValueError(f"Unsupported database type '{v}'. Supported types: {', '.join(KNOWN_DIALECT_NAMES)}") return low - @field_validator("password_ref", mode="before") + @field_validator("password_ref", "private_key_ref", "private_key_passphrase_ref", "auth_method", mode="before") @classmethod def empty_string_to_none(cls, v): if v == "": @@ -963,9 +976,25 @@ def empty_string_to_none(cls, v): def guard_extra_params(cls, v): return _validate_extra_params(v) + @model_validator(mode="after") + def check_auth_method(self): + _validate_auth_method(self.auth_method, self.database_type) + # getattr: the wire subclass (ExtDatabaseConnection) carries the key ciphertext + # in a `private_key` field this model does not have. + if self.auth_method == "key_pair" and not self.private_key_ref and not getattr(self, "private_key", None): + raise ValueError("key_pair authentication requires a private key") + return self + class FullDatabaseConnection(BaseModel): - """A complete database connection model including the secret password.""" + """A complete database connection model including the secret password. + + For key-pair auth, ``private_key`` holds the PEM text (and + ``private_key_passphrase`` its optional passphrase). Both mirror + ``password``'s update semantics: empty means "keep the existing secret". + The key is not required at the model level for that reason — creation-time + enforcement lives in ``store_database_connection``. + """ connection_name: str database_type: str = "postgresql" @@ -977,6 +1006,9 @@ class FullDatabaseConnection(BaseModel): ssl_enabled: bool | None = False url: str | None = None extra_params: dict[str, str] | None = None + auth_method: str | None = None + private_key: SecretStr | None = None + private_key_passphrase: SecretStr | None = None @field_validator("database_type") @classmethod @@ -984,14 +1016,36 @@ def normalize_database_type(cls, v: str) -> str: # lowercase only, no vocabulary check: legacy stored types (e.g. redshift) must keep loading return v.lower() + @field_validator("auth_method", mode="before") + @classmethod + def empty_auth_method_to_none(cls, v): + if v == "": + return None + return v + + @field_validator("private_key", "private_key_passphrase", mode="before") + @classmethod + def empty_key_string_to_none(cls, v): + # A raw "" (JSON clients sending the empty form field) means "no key material"; + # an explicit SecretStr("") (the project importer's placeholder refill rows) + # deliberately passes through and creates a linked empty secret. + if v == "": + return None + return v + @field_validator("extra_params") @classmethod def guard_extra_params(cls, v): return _validate_extra_params(v) + @model_validator(mode="after") + def check_auth_method(self): + _validate_auth_method(self.auth_method, self.database_type) + return self + class FullDatabaseConnectionInterface(BaseModel): - """A database connection model intended for UI display, omitting the password.""" + """A database connection model intended for UI display, omitting the password and key material.""" connection_name: str database_type: str = "postgresql" @@ -1002,6 +1056,7 @@ class FullDatabaseConnectionInterface(BaseModel): ssl_enabled: bool | None = False url: str | None = None extra_params: dict[str, str] | None = None + auth_method: str | None = None id: int | None = None access: AccessInfo | None = None diff --git a/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py index 262f8a977..c9f50061d 100644 --- a/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py +++ b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py @@ -18,6 +18,7 @@ from flowfile_core.database.connection import get_db_context from flowfile_core.flowfile.database_connection_manager.db_connections import ( delete_database_connection, + get_all_database_connections_interface, get_database_connection, get_database_connection_schema, parse_extra_params, @@ -25,7 +26,12 @@ update_database_connection, ) from flowfile_core.routes._connection_sharing import require_credentials_on_target_change -from flowfile_core.schemas.input_schema import FullDatabaseConnection +from flowfile_core.schemas.input_schema import ( + DatabaseConnection, + FullDatabaseConnection, + FullDatabaseConnectionInterface, +) +from flowfile_core.secret_manager.secret_manager import decrypt_secret from shared.db_dialects import get_dialect logger = logging.getLogger(__name__) @@ -133,6 +139,246 @@ def test_guard_requires_credentials_on_change(self): require_credentials_on_target_change([], has_new_credentials=False, has_bundled_secrets=True) +def _make_pem() -> str: + """A fresh throwaway RSA key as unencrypted PKCS#8 PEM text.""" + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode("ascii") + + +def _key_pair_connection( + name: str, private_key: str | None, passphrase: str | None = None +) -> FullDatabaseConnection: + return FullDatabaseConnection( + connection_name=name, + database_type="snowflake", + username="svc_user", + password=SecretStr(""), + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="key_pair", + private_key=SecretStr(private_key) if private_key is not None else None, + private_key_passphrase=SecretStr(passphrase) if passphrase is not None else None, + ) + + +class TestKeyPairCrud: + def test_store_and_reload_round_trips_ciphertexts(self): + name = "snowflake_kp_store" + pem = _make_pem() + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, pem, "pass-phrase"), USER_ID) + with get_db_context() as db: + reloaded = get_database_connection_schema(db, name, USER_ID) + assert reloaded is not None + assert reloaded.auth_method == "key_pair" + key_ciphertext = reloaded.private_key.get_secret_value() + assert key_ciphertext.startswith("$ffsec$"), "schema must return ciphertext, never the PEM" + assert pem not in key_ciphertext + assert decrypt_secret(key_ciphertext).get_secret_value() == pem + passphrase_ciphertext = reloaded.private_key_passphrase.get_secret_value() + assert passphrase_ciphertext.startswith("$ffsec$") + assert decrypt_secret(passphrase_ciphertext).get_secret_value() == "pass-phrase" + finally: + _cleanup(name) + + def test_create_requires_private_key(self): + _cleanup("snowflake_kp_missing") + with get_db_context() as db: + with pytest.raises(ValueError, match="requires a private key"): + store_database_connection(db, _key_pair_connection("snowflake_kp_missing", None), USER_ID) + + def test_update_rotates_key_and_empty_keeps_existing(self): + name = "snowflake_kp_rotate" + pem = _make_pem() + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, pem), USER_ID) + + with get_db_context() as db: + update_database_connection(db, _key_pair_connection(name, None), USER_ID) + with get_db_context() as db: + kept = get_database_connection_schema(db, name, USER_ID).private_key.get_secret_value() + assert decrypt_secret(kept).get_secret_value() == pem, "empty key on update must keep the existing secret" + + new_pem = _make_pem() + with get_db_context() as db: + update_database_connection(db, _key_pair_connection(name, new_pem), USER_ID) + with get_db_context() as db: + rotated = get_database_connection_schema(db, name, USER_ID).private_key.get_secret_value() + assert decrypt_secret(rotated).get_secret_value() == new_pem + finally: + _cleanup(name) + + def test_delete_removes_all_connection_secrets(self): + from flowfile_core.database.models import Secret + + name = "snowflake_kp_delete" + _cleanup(name) + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, _make_pem(), "pp"), USER_ID) + with get_db_context() as db: + row = get_database_connection(db, name, USER_ID) + secret_ids = [row.password_id, row.private_key_id, row.private_key_passphrase_id] + assert all(secret_id is not None for secret_id in secret_ids) + with get_db_context() as db: + delete_database_connection(db, name, USER_ID) + with get_db_context() as db: + assert db.query(Secret).filter(Secret.id.in_(secret_ids)).count() == 0 + + def test_switching_to_password_detaches_and_deletes_key_secrets(self): + from flowfile_core.database.models import Secret + + name = "snowflake_kp_flip" + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, _make_pem(), "pp"), USER_ID) + with get_db_context() as db: + row = get_database_connection(db, name, USER_ID) + key_secret_ids = [row.private_key_id, row.private_key_passphrase_id] + assert all(secret_id is not None for secret_id in key_secret_ids) + + flipped = FullDatabaseConnection( + connection_name=name, + database_type="snowflake", + username="svc_user", + password=SecretStr("new-pass"), + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="password", + ) + with get_db_context() as db: + update_database_connection(db, flipped, USER_ID) + with get_db_context() as db: + row = get_database_connection(db, name, USER_ID) + assert row.private_key_id is None, "a rotated-away key must not linger on the row" + assert row.private_key_passphrase_id is None + assert db.query(Secret).filter(Secret.id.in_(key_secret_ids)).count() == 0 + reloaded = get_database_connection_schema(db, name, USER_ID) + assert reloaded.private_key is None, "resolvers must no longer see any key material" + finally: + _cleanup(name) + + def test_update_to_key_pair_without_key_is_rejected(self): + name = "snowflake_kp_bad_flip" + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _snowflake_connection(name), USER_ID) + with get_db_context() as db: + with pytest.raises(ValueError, match="requires a private key"): + update_database_connection(db, _key_pair_connection(name, None), USER_ID) + finally: + _cleanup(name) + + def test_interface_reports_auth_method_without_key_material(self): + name = "snowflake_kp_interface" + pem = _make_pem() + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, pem), USER_ID) + with get_db_context() as db: + interfaces = get_all_database_connections_interface(db, USER_ID) + entry = next(i for i in interfaces if i.connection_name == name) + assert entry.auth_method == "key_pair" + assert pem not in str(entry.model_dump()) + finally: + _cleanup(name) + + +class TestKeyPairModelValidation: + def test_key_pair_rejected_for_non_supporting_dialect(self): + with pytest.raises(ValidationError, match="not supported by database type"): + FullDatabaseConnection( + connection_name="pg_kp", + database_type="postgresql", + username="u", + password=SecretStr(""), + auth_method="key_pair", + private_key=SecretStr("-----BEGIN PRIVATE KEY-----"), + ) + + def test_inline_key_pair_requires_private_key_ref(self): + with pytest.raises(ValidationError, match="requires a private key"): + DatabaseConnection(database_type="snowflake", username="u", auth_method="key_pair") + conn = DatabaseConnection( + database_type="snowflake", username="u", auth_method="key_pair", private_key_ref="kp_secret" + ) + assert conn.private_key_ref == "kp_secret" + + def test_raw_empty_key_strings_normalize_to_none(self): + # JSON clients send "" for the blank form fields; that must not create key secrets. + conn = FullDatabaseConnection( + connection_name="pw_conn", + database_type="postgresql", + username="u", + password=SecretStr("pw"), + private_key="", + private_key_passphrase="", + ) + assert conn.private_key is None + assert conn.private_key_passphrase is None + # The importer's explicit SecretStr("") placeholder rows pass through untouched. + refill = FullDatabaseConnection( + connection_name="kp_refill", + database_type="snowflake", + username="u", + password=SecretStr(""), + extra_params={"account": "acct"}, + auth_method="key_pair", + private_key=SecretStr(""), + ) + assert refill.private_key is not None + + def test_inline_empty_strings_normalize_to_none(self): + conn = DatabaseConnection( + database_type="snowflake", username="u", auth_method="", private_key_ref="", private_key_passphrase_ref="" + ) + assert conn.auth_method is None + assert conn.private_key_ref is None + assert conn.private_key_passphrase_ref is None + + def test_interface_model_never_exposes_key_material(self): + assert "private_key" not in FullDatabaseConnectionInterface.model_fields + assert "private_key_passphrase" not in FullDatabaseConnectionInterface.model_fields + assert "auth_method" in FullDatabaseConnectionInterface.model_fields + + +class TestAuthMethodRepointGuard: + """The PUT route appends auth_method via a normalized comparison (stored NULL == "password").""" + + @staticmethod + def _changed(row_value: str | None, incoming: str | None) -> bool: + return (row_value or "password") != (incoming or "password") + + def test_null_and_password_are_equivalent(self): + assert self._changed(None, "password") is False + assert self._changed("password", None) is False + assert self._changed(None, None) is False + assert self._changed("key_pair", "key_pair") is False + + def test_auth_method_flip_trips_the_guard(self): + from fastapi import HTTPException + + assert self._changed(None, "key_pair") is True + assert self._changed("key_pair", "password") is True + with pytest.raises(HTTPException) as exc: + require_credentials_on_target_change(["auth_method"], has_new_credentials=False, has_bundled_secrets=True) + assert exc.value.status_code == 422 + require_credentials_on_target_change(["auth_method"], has_new_credentials=True, has_bundled_secrets=True) + + def _live_uri() -> str: dialect = get_dialect("snowflake") extra = { @@ -188,3 +434,45 @@ def test_sql_source_schema_prediction(self): source = SqlSource(connection_string=_live_uri(), query="SELECT 1 AS one", database_type="snowflake") columns = source.get_schema() assert [c.name for c in columns] == ["ONE"] + + +_KEY_PAIR_LIVE_ENV = ("ACCOUNT", "USER", "DATABASE", "WAREHOUSE") + + +def _live_key_pair_configured() -> bool: + return bool(os.environ.get("FLOWFILE_TEST_SNOWFLAKE_PRIVATE_KEY_PATH")) and all( + os.environ.get(f"FLOWFILE_TEST_SNOWFLAKE_{name}") for name in _KEY_PAIR_LIVE_ENV + ) + + +snowflake_key_pair_available = pytest.mark.skipif( + not _live_key_pair_configured(), + reason="FLOWFILE_TEST_SNOWFLAKE_PRIVATE_KEY_PATH and account credentials are not configured", +) + + +@snowflake_key_pair_available +class TestSnowflakeKeyPairLive: + def test_key_pair_read_round_trip(self): + dialect = get_dialect("snowflake") + with open(os.environ["FLOWFILE_TEST_SNOWFLAKE_PRIVATE_KEY_PATH"], encoding="utf-8") as f: + pem = f.read() + extra = { + "account": os.environ["FLOWFILE_TEST_SNOWFLAKE_ACCOUNT"], + "warehouse": os.environ["FLOWFILE_TEST_SNOWFLAKE_WAREHOUSE"], + } + if os.environ.get("FLOWFILE_TEST_SNOWFLAKE_ROLE"): + extra["role"] = os.environ["FLOWFILE_TEST_SNOWFLAKE_ROLE"] + if os.environ.get("FLOWFILE_TEST_SNOWFLAKE_SCHEMA"): + extra["schema"] = os.environ["FLOWFILE_TEST_SNOWFLAKE_SCHEMA"] + uri = dialect.build_uri( + username=os.environ["FLOWFILE_TEST_SNOWFLAKE_USER"], + database=os.environ["FLOWFILE_TEST_SNOWFLAKE_DATABASE"], + auth_method="key_pair", + private_key=pem, + private_key_passphrase=os.environ.get("FLOWFILE_TEST_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"), + **extra, + ) + result = dialect.read("SELECT 1 AS one", uri, logger) + assert result.height == 1 + assert result.columns == ["ONE"] diff --git a/flowfile_core/tests/project/test_roundtrip.py b/flowfile_core/tests/project/test_roundtrip.py index bf8d6d657..b64c4d891 100644 --- a/flowfile_core/tests/project/test_roundtrip.py +++ b/flowfile_core/tests/project/test_roundtrip.py @@ -1315,3 +1315,75 @@ def test_database_connection_extra_params_round_trip(tmp_path, monkeypatch): assert restored.extra_params == extra_params finally: _cleanup([conn], []) + + +def test_database_connection_key_pair_round_trip(tmp_path, monkeypatch): + """A key-pair snowflake connection projects auth_method verbatim and the key/passphrase + as ${secret:...} placeholders (never PEM text or ciphertext), the secret manifest keeps + treating the linked key secrets as connection-implied, and import restores a working + connection from FLOWFILE_SECRET_* env vars.""" + import yaml + + project_sync.close_project(OWNER) + conn = "proj_db_snowflake_kp" + pem = "-----BEGIN PRIVATE KEY-----\nkey-material-for-round-trip\n-----END PRIVATE KEY-----\n" + with get_db_context() as db: + if get_database_connection(db, conn, OWNER) is None: + store_database_connection( + db, + input_schema.FullDatabaseConnection( + connection_name=conn, + database_type="snowflake", + username="svc_reader", + password="", + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="key_pair", + private_key=pem, + private_key_passphrase="kp-pass", + ), + OWNER, + ) + root = tmp_path / "project" + try: + project_sync.init_project(str(root), "Snowflake KP RT", OWNER) + conn_text = (root / "connections" / "database" / f"{conn}.yaml").read_text(encoding="utf-8") + assert "auth_method: key_pair" in conn_text + assert f"${{secret:{conn}_private_key}}" in conn_text + assert f"${{secret:{conn}_private_key_passphrase}}" in conn_text + assert "BEGIN PRIVATE KEY" not in conn_text + assert "kp-pass" not in conn_text + assert "$ffsec$" not in conn_text + + manifest_data = yaml.safe_load((root / "secrets.yaml").read_text(encoding="utf-8")) + listed = set(manifest_data.get("required_secrets") or []) + assert f"{conn}_private_key" not in listed, "linked key secrets are connection-implied" + assert f"{conn}_private_key_passphrase" not in listed + + project_sync.close_project(OWNER) + with get_db_context() as db: + delete_database_connection(db, conn, OWNER) + + monkeypatch.setenv("FLOWFILE_SECRET_PROJ_DB_SNOWFLAKE_KP", "") + monkeypatch.setenv("FLOWFILE_SECRET_PROJ_DB_SNOWFLAKE_KP_PRIVATE_KEY", pem) + monkeypatch.setenv("FLOWFILE_SECRET_PROJ_DB_SNOWFLAKE_KP_PRIVATE_KEY_PASSPHRASE", "kp-pass") + from flowfile_core.project.importer import import_project + + import_project(root, OWNER) + from flowfile_core.flowfile.database_connection_manager.db_connections import ( + get_database_connection_schema, + ) + from flowfile_core.secret_manager.secret_manager import decrypt_secret + + with get_db_context() as db: + restored = get_database_connection_schema(db, conn, OWNER) + assert restored is not None + assert restored.auth_method == "key_pair" + assert restored.private_key is not None + assert decrypt_secret(restored.private_key.get_secret_value()).get_secret_value() == pem + assert restored.private_key_passphrase is not None + assert ( + decrypt_secret(restored.private_key_passphrase.get_secret_value()).get_secret_value() == "kp-pass" + ) + finally: + _cleanup([conn], []) diff --git a/flowfile_core/tests/test_db_dialect_endpoints.py b/flowfile_core/tests/test_db_dialect_endpoints.py index 1f654e710..b8d117116 100644 --- a/flowfile_core/tests/test_db_dialect_endpoints.py +++ b/flowfile_core/tests/test_db_dialect_endpoints.py @@ -43,6 +43,7 @@ def test_get_db_dialects(): assert by_name["sqlite"]["default_port"] is None assert by_name["postgresql"]["default_port"] == 5432 assert by_name["postgresql"]["supports_ssl"] is True + assert by_name["postgresql"]["auth_methods"] == ["password"] assert all(d["available"] is True for d in dialects) @@ -223,6 +224,7 @@ def test_snowflake_in_dialect_catalog(): assert [f["name"] for f in entry["extra_fields"]] == ["account", "warehouse", "role"] assert entry["extra_fields"][0]["required"] is True assert entry["hidden_fields"] == ["host", "port", "ssl"] + assert entry["auth_methods"] == ["password", "key_pair"] def test_create_snowflake_connection_with_extra_params(): @@ -257,6 +259,99 @@ def test_create_connection_rejects_blocked_extra_params(): _cleanup_connection("snowflake_evil") +def test_create_connection_rejects_auth_method_extra_param(): + payload = { + "connection_name": "snowflake_evil_auth", + "database_type": "snowflake", + "username": "user", + "password": "pass", + "extra_params": {"account": "acct", "auth_method": "key_pair"}, + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 422, response.text + assert "auth_method" in response.text + _cleanup_connection("snowflake_evil_auth") + + +_TEST_PEM = ( + "-----BEGIN PRIVATE KEY-----\nMIIB-not-a-real-key-but-round-trips-fine\n-----END PRIVATE KEY-----\n" +) + + +def test_snowflake_key_pair_connection_http_round_trip(): + _cleanup_connection("snowflake_kp_http") + payload = { + "connection_name": "snowflake_kp_http", + "database_type": "snowflake", + "username": "svc_user", + "password": "", + "database": "ANALYTICS", + "extra_params": {"account": "myorg-myaccount"}, + "auth_method": "key_pair", + "private_key": _TEST_PEM, + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 200, response.text + try: + listed = client.get("/db_connection_lib").json() + entry = next(c for c in listed if c["connection_name"] == "snowflake_kp_http") + assert entry["auth_method"] == "key_pair" + assert "private_key" not in entry, "the interface payload must never carry key material" + assert _TEST_PEM not in str(listed) + + # The OWNER may flip the auth method freely — the anti-repoint guard applies to + # manage-grantees only (helper-level coverage in test_snowflake_source.py). + rotate = {**payload, "auth_method": "password", "private_key": "", "password": "new-pass"} + response = client.put("/db_connection_lib", json=rotate) + assert response.status_code == 200, response.text + listed = client.get("/db_connection_lib").json() + entry = next(c for c in listed if c["connection_name"] == "snowflake_kp_http") + assert entry["auth_method"] == "password" + + # The flip must fully shed the key material — a stale key must never keep authenticating. + from flowfile_core.flowfile.database_connection_manager.db_connections import ( + get_database_connection_schema, + ) + + with get_db_context() as db: + restored = get_database_connection_schema(db, "snowflake_kp_http", 1) + assert restored.private_key is None + assert restored.private_key_passphrase is None + finally: + _cleanup_connection("snowflake_kp_http") + + +def test_create_key_pair_connection_requires_private_key(): + _cleanup_connection("snowflake_kp_incomplete") + payload = { + "connection_name": "snowflake_kp_incomplete", + "database_type": "snowflake", + "username": "svc_user", + "password": "", + "extra_params": {"account": "acct"}, + "auth_method": "key_pair", + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 422, response.text + assert "private key" in response.text + _cleanup_connection("snowflake_kp_incomplete") + + +def test_create_key_pair_connection_rejected_for_password_only_dialect(): + payload = { + "connection_name": "pg_kp_http", + "database_type": "postgresql", + "username": "u", + "password": "", + "host": "h", + "auth_method": "key_pair", + "private_key": _TEST_PEM, + } + response = client.post("/db_connection_lib", json=payload) + assert response.status_code == 422, response.text + _cleanup_connection("pg_kp_http") + + def test_create_duckdb_connection_without_credentials(): _cleanup_connection("duckdb_conn") payload = { diff --git a/flowfile_core/tests/test_migration.py b/flowfile_core/tests/test_migration.py index 76e487b25..0ab87ef24 100644 --- a/flowfile_core/tests/test_migration.py +++ b/flowfile_core/tests/test_migration.py @@ -913,3 +913,104 @@ def test_upgrade_is_guarded_when_the_column_already_exists(self, tmp_path, monke _run_migration(db_path, monkeypatch) assert "scd2_config" in self._columns(db_path, "catalog_tables") + + +# Migration 031: database_connections key-pair auth columns + + +class TestKeyPairColumnsMigration: + _NEW_COLUMNS = ("auth_method", "private_key_id", "private_key_passphrase_id") + + @staticmethod + def _columns(db_path: Path, table: str) -> set[str]: + engine = create_engine(f"sqlite:///{db_path}") + names = {c["name"] for c in inspect(engine).get_columns(table)} + engine.dispose() + return names + + def test_fresh_install_has_the_columns(self, tmp_path, monkeypatch): + db_path = tmp_path / "catalog.db" + _run_migration(db_path, monkeypatch) + columns = self._columns(db_path, "database_connections") + assert set(self._NEW_COLUMNS) <= columns + + def test_upgrade_from_030_adds_the_columns(self, tmp_path, monkeypatch): + from alembic import command + + from flowfile_core.database.migration import _get_alembic_config + + db_path = tmp_path / "catalog.db" + monkeypatch.setenv("FLOWFILE_DB_PATH", str(db_path)) + command.upgrade(_get_alembic_config(), "030") + assert not set(self._NEW_COLUMNS) & self._columns(db_path, "database_connections") + + _run_migration(db_path, monkeypatch) + assert set(self._NEW_COLUMNS) <= self._columns(db_path, "database_connections") + + def test_existing_rows_read_back_null(self, tmp_path, monkeypatch): + """Pre-031 connections become password connections (NULL auth_method), never half-migrated.""" + from alembic import command + + from flowfile_core.database.migration import _get_alembic_config + + db_path = tmp_path / "catalog.db" + monkeypatch.setenv("FLOWFILE_DB_PATH", str(db_path)) + command.upgrade(_get_alembic_config(), "030") + + engine = create_engine(f"sqlite:///{db_path}") + with engine.connect() as conn: + conn.execute( + text( + "INSERT INTO database_connections (connection_name, database_type, username, user_id) " + "VALUES ('legacy_conn', 'postgresql', 'u', 1)" + ) + ) + conn.commit() + engine.dispose() + + _run_migration(db_path, monkeypatch) + + engine = create_engine(f"sqlite:///{db_path}") + with engine.connect() as conn: + row = conn.execute( + text( + "SELECT auth_method, private_key_id, private_key_passphrase_id " + "FROM database_connections WHERE connection_name = 'legacy_conn'" + ) + ).one() + engine.dispose() + assert tuple(row) == (None, None, None) + + def test_downgrade_then_upgrade_round_trips(self, tmp_path, monkeypatch): + from alembic import command + + from flowfile_core.database.migration import _get_alembic_config + + db_path = tmp_path / "catalog.db" + _run_migration(db_path, monkeypatch) + cfg = _get_alembic_config() + + command.downgrade(cfg, "030") + assert not set(self._NEW_COLUMNS) & self._columns(db_path, "database_connections") + + command.upgrade(cfg, "031") + assert set(self._NEW_COLUMNS) <= self._columns(db_path, "database_connections") + + def test_upgrade_is_guarded_when_a_column_already_exists(self, tmp_path, monkeypatch): + """Dev DBs that added a column out of band must not fail the startup upgrade.""" + from alembic import command + + from flowfile_core.database.migration import _get_alembic_config + + db_path = tmp_path / "catalog.db" + monkeypatch.setenv("FLOWFILE_DB_PATH", str(db_path)) + command.upgrade(_get_alembic_config(), "030") + + engine = create_engine(f"sqlite:///{db_path}") + with engine.connect() as conn: + conn.execute(text("ALTER TABLE database_connections ADD COLUMN auth_method VARCHAR")) + conn.commit() + engine.dispose() + + _run_migration(db_path, monkeypatch) + assert set(self._NEW_COLUMNS) <= self._columns(db_path, "database_connections") diff --git a/flowfile_frame/flowfile_frame/database/connection_manager.py b/flowfile_frame/flowfile_frame/database/connection_manager.py index d93edd1ca..8eda19eac 100644 --- a/flowfile_frame/flowfile_frame/database/connection_manager.py +++ b/flowfile_frame/flowfile_frame/database/connection_manager.py @@ -41,6 +41,9 @@ def create_database_connection( ssl_enabled: bool = False, url: str | None = None, extra_params: dict[str, str] | None = None, + auth_method: str | None = None, + private_key: str | SecretStr | None = None, + private_key_passphrase: str | SecretStr | None = None, ) -> FullDatabaseConnection: """Create and store a new database connection. @@ -58,13 +61,19 @@ def create_database_connection( extra_params: Dialect-specific connection parameters, e.g. for Snowflake ``{"account": "myorg-myaccount", "warehouse": "COMPUTE_WH", "role": "ANALYST"}``. Keys that could override credentials (user, password, host, ...) are rejected. + auth_method: Authentication method; must be supported by the dialect + (``"password"`` everywhere, ``"key_pair"`` for Snowflake JWT auth). + private_key: Private key PEM *text* (never a path) for key-pair auth, e.g. + ``open("rsa_key.p8").read()``. Stored as an encrypted secret. + private_key_passphrase: Optional passphrase when the PEM is encrypted. Returns: FullDatabaseConnection: The created connection object. Raises: - ValueError: If a connection with this name already exists, or the - database_type is not a supported dialect. + ValueError: If a connection with this name already exists, the + database_type is not a supported dialect, or key-pair auth is + requested without a private key. """ if database_type.lower() not in KNOWN_DIALECT_NAMES: raise ValueError( @@ -74,11 +83,17 @@ def create_database_connection( if isinstance(password, str): password = SecretStr(password) + if isinstance(private_key, str): + private_key = SecretStr(private_key) + if isinstance(private_key_passphrase, str): + private_key_passphrase = SecretStr(private_key_passphrase) if get_dialect_or_generic(database_type).file_based: # No credentials for file-based databases; the stored model requires strings username = username or "" password = password if password is not None else SecretStr("") + if auth_method == "key_pair" and password is None: + password = SecretStr("") connection = FullDatabaseConnection( connection_name=connection_name, @@ -91,6 +106,9 @@ def create_database_connection( ssl_enabled=ssl_enabled, url=url, extra_params=extra_params, + auth_method=auth_method, + private_key=private_key, + private_key_passphrase=private_key_passphrase, ) with get_db_context() as db: @@ -111,6 +129,9 @@ def create_database_connection_if_not_exists( ssl_enabled: bool = False, url: str | None = None, extra_params: dict[str, str] | None = None, + auth_method: str | None = None, + private_key: str | SecretStr | None = None, + private_key_passphrase: str | SecretStr | None = None, ) -> FullDatabaseConnection: """Create a database connection if it doesn't already exist. @@ -126,6 +147,9 @@ def create_database_connection_if_not_exists( ssl_enabled: Whether to use SSL for the connection. url: Full database URL (overrides other connection parameters). extra_params: Dialect-specific connection parameters (see create_database_connection). + auth_method: Authentication method (see create_database_connection). + private_key: Private key PEM text for key-pair auth (see create_database_connection). + private_key_passphrase: Optional passphrase when the PEM is encrypted. Returns: FullDatabaseConnection: The existing or newly created connection. @@ -147,6 +171,9 @@ def create_database_connection_if_not_exists( ssl_enabled=ssl_enabled, url=url, extra_params=extra_params, + auth_method=auth_method, + private_key=private_key, + private_key_passphrase=private_key_passphrase, ) @@ -187,6 +214,7 @@ def get_all_available_database_connections() -> list[FullDatabaseConnectionInter database=conn.database, ssl_enabled=conn.ssl_enabled, extra_params=parse_extra_params(conn.extra_params), + auth_method=conn.auth_method, ) for conn in connections ] @@ -207,11 +235,15 @@ def del_database_connection(connection_name: str) -> bool: with get_db_context() as db: connection = get_database_connection(db, connection_name, user_id) if connection: - # Delete the associated password secret - if connection.password_id: - secret = db.query(Secret).filter(Secret.id == connection.password_id).first() - if secret: - db.delete(secret) + # Delete every associated secret (password + key-pair material) + all_secret_ids = ( + connection.password_id, + connection.private_key_id, + connection.private_key_passphrase_id, + ) + secret_ids = [secret_id for secret_id in all_secret_ids if secret_id is not None] + if secret_ids: + db.query(Secret).filter(Secret.id.in_(secret_ids)).delete(synchronize_session=False) db.delete(connection) db.commit() diff --git a/flowfile_frame/flowfile_frame/database/connection_manager.pyi b/flowfile_frame/flowfile_frame/database/connection_manager.pyi index 82dfe0a7f..42bd8897a 100644 --- a/flowfile_frame/flowfile_frame/database/connection_manager.pyi +++ b/flowfile_frame/flowfile_frame/database/connection_manager.pyi @@ -6,8 +6,8 @@ from pydantic import SecretStr from flowfile_core.schemas.input_schema import FullDatabaseConnection, FullDatabaseConnectionInterface def get_current_user_id() -> int: ... -def create_database_connection(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None) -> FullDatabaseConnection: ... -def create_database_connection_if_not_exists(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None) -> FullDatabaseConnection: ... +def create_database_connection(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None, auth_method: str | None=None, private_key: str | SecretStr | None=None, private_key_passphrase: str | SecretStr | None=None) -> FullDatabaseConnection: ... +def create_database_connection_if_not_exists(connection_name: str, *, database_type: str='postgresql', host: str | None=None, port: int | None=None, database: str | None=None, username: str | None=None, password: str | SecretStr | None=None, ssl_enabled: bool=False, url: str | None=None, extra_params: dict[str, str] | None=None, auth_method: str | None=None, private_key: str | SecretStr | None=None, private_key_passphrase: str | SecretStr | None=None) -> FullDatabaseConnection: ... def get_database_connection_by_name(connection_name: str) -> FullDatabaseConnection | None: ... def get_all_available_database_connections() -> list[FullDatabaseConnectionInterface]: ... def del_database_connection(connection_name: str) -> bool: ... diff --git a/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts b/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts index f8acb0d72..fed429c18 100644 --- a/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts +++ b/flowfile_frontend/src/renderer/app/api/dbDialects.test.ts @@ -95,4 +95,16 @@ describe("getDbDialects", () => { expect(snowflake?.extra_fields?.[0].required).toBe(true); expect(snowflake?.hidden_fields).toEqual(["host", "port", "ssl"]); }); + + it("carries auth methods in every fallback entry, key_pair only on snowflake", async () => { + const { FALLBACK_DIALECTS } = await loadModule(); + for (const dialect of FALLBACK_DIALECTS) { + expect(dialect.auth_methods, dialect.name).toContain("password"); + if (dialect.name === "snowflake") { + expect(dialect.auth_methods).toEqual(["password", "key_pair"]); + } else { + expect(dialect.auth_methods).toEqual(["password"]); + } + } + }); }); diff --git a/flowfile_frontend/src/renderer/app/api/dbDialects.ts b/flowfile_frontend/src/renderer/app/api/dbDialects.ts index ec16b6d45..0d724658c 100644 --- a/flowfile_frontend/src/renderer/app/api/dbDialects.ts +++ b/flowfile_frontend/src/renderer/app/api/dbDialects.ts @@ -17,6 +17,8 @@ export interface DbDialectInfo { // form fields the dialect hides; absent on older cores, so keep optional. extra_fields?: DbDialectFieldInfo[]; hidden_fields?: string[]; + // Supported authentication methods (e.g. ["password", "key_pair"]); absent on older cores. + auth_methods?: string[]; } // Rendered when the catalog request fails (offline, older core). Must mirror the @@ -31,6 +33,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ available: true, extra_fields: [], hidden_fields: [], + auth_methods: ["password"], }, { name: "mysql", @@ -41,6 +44,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ available: true, extra_fields: [], hidden_fields: [], + auth_methods: ["password"], }, { name: "sqlite", @@ -51,6 +55,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ available: true, extra_fields: [], hidden_fields: [], + auth_methods: ["password"], }, { name: "duckdb", @@ -61,6 +66,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ available: true, extra_fields: [], hidden_fields: [], + auth_methods: ["password"], }, { name: "mssql", @@ -71,6 +77,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ available: true, extra_fields: [], hidden_fields: [], + auth_methods: ["password"], }, { name: "snowflake", @@ -85,6 +92,7 @@ export const FALLBACK_DIALECTS: DbDialectInfo[] = [ { name: "role", label: "Role", required: false }, ], hidden_fields: ["host", "port", "ssl"], + auth_methods: ["password", "key_pair"], }, ]; diff --git a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue index 9e45d4338..7f5c8c7d4 100644 --- a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue +++ b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue @@ -15,9 +15,7 @@ id="database-type" :value="modelValue.database_type" class="form-control" - @change=" - (e: Event) => updateField('database_type', (e.target as HTMLSelectElement).value) - " + @change="(e: Event) => updateDatabaseType((e.target as HTMLSelectElement).value)" >
-
+
+ + +
+ +
+ + + +
+ +
+ + +
+
(); -const { dialects, isFileBased, extraFields, isFieldHidden } = useDbDialects(); +const { dialects, isFileBased, extraFields, isFieldHidden, authMethods } = useDbDialects(); + +const AUTH_METHOD_LABELS: Record = { + password: "Password", + key_pair: "Key pair (JWT)", +}; const isFileBasedConnection = computed(() => isFileBased(props.modelValue.database_type)); @@ -141,6 +193,12 @@ const dialectExtraFields = computed(() => extraFields(props.modelValue.database_ const isHidden = (field: string) => isFieldHidden(props.modelValue.database_type, field); +const authMethodOptions = computed(() => authMethods(props.modelValue.database_type)); + +const usesKeyPair = computed(() => props.modelValue.auth_method === "key_pair"); + +const authMethodLabel = (method: string): string => AUTH_METHOD_LABELS[method] ?? method; + const emit = defineEmits<{ (e: "update:modelValue", value: DatabaseConnection): void; }>(); @@ -160,6 +218,18 @@ const updateField = ( }); }; +const updateDatabaseType = (value: string) => { + const next: DatabaseConnection = { ...props.modelValue, database_type: value }; + // Auth fields are dialect-specific: a stale key_pair selection on a dialect + // without it would fail backend validation with no visible field to fix. + if (!authMethods(value).includes(next.auth_method || "password")) { + next.auth_method = undefined; + next.private_key_ref = undefined; + next.private_key_passphrase_ref = undefined; + } + emit("update:modelValue", next); +}; + const updateExtraParam = (name: string, value: string) => { const params = { ...(props.modelValue.extra_params || {}) }; if (value) { diff --git a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseReader.vue b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseReader.vue index 2dff7e707..1fba14de8 100644 --- a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseReader.vue +++ b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseReader.vue @@ -371,7 +371,8 @@ const isInlineConnectionComplete = () => { const conn = nodeDatabaseReader.value?.database_settings?.database_connection; if (!conn) return false; if (conn.database_type === "sqlite") return !!conn.database; - return !!(conn.host && conn.port && conn.database && conn.username && conn.password_ref); + const credentialRef = conn.auth_method === "key_pair" ? conn.private_key_ref : conn.password_ref; + return !!(conn.host && conn.port && conn.database && conn.username && credentialRef); }; let inlineDebounceTimer: ReturnType | null = null; diff --git a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseWriter/DatabaseWriter.vue b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseWriter/DatabaseWriter.vue index 2c2163ad6..e1abfc149 100644 --- a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseWriter/DatabaseWriter.vue +++ b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseWriter/DatabaseWriter.vue @@ -317,7 +317,8 @@ const isInlineConnectionComplete = () => { const conn = nodeData.value?.database_write_settings?.database_connection; if (!conn) return false; if (conn.database_type === "sqlite") return !!conn.database; - return !!(conn.host && conn.port && conn.database && conn.username && conn.password_ref); + const credentialRef = conn.auth_method === "key_pair" ? conn.private_key_ref : conn.password_ref; + return !!(conn.host && conn.port && conn.database && conn.username && credentialRef); }; let inlineDebounceTimer: ReturnType | null = null; diff --git a/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts b/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts index a17039166..50d90e829 100644 --- a/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts +++ b/flowfile_frontend/src/renderer/app/composables/useDbDialects.ts @@ -38,5 +38,8 @@ export function useDbDialects() { const isFieldHidden = (databaseType?: string, field?: string): boolean => (findDialect(databaseType)?.hidden_fields ?? []).includes(field || ""); - return { dialects, findDialect, isFileBased, defaultPort, extraFields, isFieldHidden }; + const authMethods = (databaseType?: string): string[] => + findDialect(databaseType)?.auth_methods ?? ["password"]; + + return { dialects, findDialect, isFileBased, defaultPort, extraFields, isFieldHidden, authMethods }; } diff --git a/flowfile_frontend/src/renderer/app/types/node.types.ts b/flowfile_frontend/src/renderer/app/types/node.types.ts index 40d0392df..3ad7f1546 100644 --- a/flowfile_frontend/src/renderer/app/types/node.types.ts +++ b/flowfile_frontend/src/renderer/app/types/node.types.ts @@ -710,6 +710,9 @@ export interface DatabaseConnection extends BaseConnection { username?: string; password_ref?: string; // Unused by file-based databases (sqlite, duckdb) extra_params?: Record | null; // Dialect-specific params (e.g. snowflake account) + auth_method?: string | null; // From the dialect's auth_methods; null/absent == password + private_key_ref?: string; // Secret holding the private key PEM (key-pair auth) + private_key_passphrase_ref?: string; // Secret holding the optional key passphrase } export type ConnectionModeOption = "inline" | "reference"; export type IfExistAction = "append" | "replace" | "fail"; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue index ef25415fd..6bd825e08 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseConnectionSettings.vue @@ -25,6 +25,15 @@
+
+ + +
+
-
+
+
+ + +
+ +
+ +
+ + +
+
+
+ + diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue index 477aa6bab..ec7b6d2cf 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/DatabaseView.vue @@ -226,6 +226,9 @@ const showEditModal = (connection: FullDatabaseConnectionInterface) => { sslEnabled: connection.sslEnabled, url: connection.url || "", extraParams: connection.extraParams || undefined, + authMethod: connection.authMethod || "password", + privateKey: "", // Key material is never returned from the API + privateKeyPassphrase: "", }; dialogVisible.value = true; }; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts index 5a1990954..f72d9ed06 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts @@ -25,6 +25,11 @@ const toPythonFormat = (connection: FullDatabaseConnection): PythonFullDatabaseC ssl_enabled: connection.sslEnabled, url: connection.url, extra_params: connection.extraParams, + auth_method: connection.authMethod, + // undefined (dropped from JSON) rather than "": an empty string would create + // empty key-secret rows server-side; blank means "no key" / "keep existing". + private_key: connection.privateKey || undefined, + private_key_passphrase: connection.privateKeyPassphrase || undefined, }; }; @@ -57,6 +62,7 @@ export const convertConnectionInterfacePytoTs = ( url: pythonConnectionInterface.url, database: pythonConnectionInterface.database, extraParams: pythonConnectionInterface.extra_params, + authMethod: pythonConnectionInterface.auth_method ?? undefined, id: pythonConnectionInterface.id, access: pythonConnectionInterface.access, }; @@ -75,6 +81,7 @@ export const convertConnectionInterfaceTstoPy = ( ssl_enabled: dbConnectionInterface.sslEnabled, url: dbConnectionInterface.url, extra_params: dbConnectionInterface.extraParams, + auth_method: dbConnectionInterface.authMethod, }; }; diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts index 3185eb3a7..933a28332 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts @@ -17,6 +17,9 @@ export interface PythonFullDatabaseConnection { ssl_enabled: boolean; url?: string; extra_params?: Record | null; + auth_method?: string | null; + private_key?: string; + private_key_passphrase?: string; } export interface FullDatabaseConnection { @@ -30,6 +33,9 @@ export interface FullDatabaseConnection { sslEnabled: boolean; url?: string; extraParams?: Record | null; + authMethod?: string; + privateKey?: string; + privateKeyPassphrase?: string; } export interface PythonFullDatabaseConnectionInterface { @@ -42,6 +48,7 @@ export interface PythonFullDatabaseConnectionInterface { ssl_enabled: boolean; url?: string; extra_params?: Record | null; + auth_method?: string | null; id?: number; access?: AccessInfo | null; } @@ -56,6 +63,7 @@ export interface FullDatabaseConnectionInterface { sslEnabled: boolean; url?: string; extraParams?: Record | null; + authMethod?: string; id?: number; access?: AccessInfo | null; } diff --git a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py index e3ea6bf72..77886c307 100644 --- a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py +++ b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py @@ -5,9 +5,11 @@ from flowfile_worker.secrets import decrypt_secret from shared.sql_utils import construct_sql_uri, get_sqlalchemy_uri +_RESERVED_EXTRA_PARAMS = ("auth_method", "private_key", "private_key_passphrase") + class DataBaseConnection(BaseModel): - """Database connection configuration with secure password handling.""" + """Database connection configuration with secure credential handling.""" username: str | None = None password: SecretStr | None = None # Encrypted password @@ -18,10 +20,19 @@ class DataBaseConnection(BaseModel): ssl_enabled: bool | None = False url: str | None = None extra_params: dict[str, str] | None = None # Dialect-specific params (e.g. snowflake account/warehouse) + auth_method: str | None = None # None == password auth + private_key: SecretStr | None = None # Encrypted private key PEM (key-pair auth) + private_key_passphrase: SecretStr | None = None # Encrypted private-key passphrase def get_decrypted_secret(self) -> SecretStr: return decrypt_secret(self.password.get_secret_value()) + @staticmethod + def _decrypt(value: SecretStr | None) -> str | None: + if not value or not value.get_secret_value(): + return None + return decrypt_secret(value.get_secret_value()).get_secret_value() + def create_uri(self) -> str: """ Creates a database URI based on the connection details. @@ -31,20 +42,23 @@ def create_uri(self) -> str: Returns: str: The database URI (base scheme, suitable for connectorx) """ - password_str = None - if self.password: - password_str = decrypt_secret(self.password.get_secret_value()).get_secret_value() + # Belt and braces: the named auth params must never collide with a splatted + # extra_params key (core rejects these at its API boundary already). + extra_params = {k: v for k, v in (self.extra_params or {}).items() if k not in _RESERVED_EXTRA_PARAMS} return construct_sql_uri( database_type=self.database_type, host=self.host, port=self.port, username=self.username, - password=password_str, + password=self._decrypt(self.password), database=self.database, url=self.url, ssl_enabled=bool(self.ssl_enabled), connect_timeout=10, - **(self.extra_params or {}), + auth_method=self.auth_method, + private_key=self._decrypt(self.private_key), + private_key_passphrase=self._decrypt(self.private_key_passphrase), + **extra_params, ) def create_sqlalchemy_uri(self) -> str: diff --git a/flowfile_worker/tests/external_sources/test_dialect_ports.py b/flowfile_worker/tests/external_sources/test_dialect_ports.py index 613554bc8..85e9accc7 100644 --- a/flowfile_worker/tests/external_sources/test_dialect_ports.py +++ b/flowfile_worker/tests/external_sources/test_dialect_ports.py @@ -76,3 +76,54 @@ def test_snowflake_create_uri_carries_extra_params(): extra_params={"account": "myorg-myaccount", "warehouse": "COMPUTE_WH"}, ) assert connection.create_uri() == "snowflake://u@myorg-myaccount/ANALYTICS?warehouse=COMPUTE_WH" + + +def _pem() -> str: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode("ascii") + + +def test_snowflake_create_uri_key_pair_decrypts_and_omits_password(): + # Real crypto under TEST_MODE=1 (conftest), matching test_sql_source.py's pattern. + from flowfile_worker.secrets import encrypt_secret + + pem = _pem() + connection = DataBaseConnection( + database_type="snowflake", + username="svc", + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="key_pair", + private_key=encrypt_secret(pem), + private_key_passphrase=None, + ) + uri = connection.create_uri() + assert uri.startswith("snowflake://svc@myorg-myaccount/ANALYTICS?") + assert "authenticator=snowflake_jwt" in uri + assert "private_key=" in uri + assert pem.splitlines()[1] not in uri, "the PEM body must be b64-wrapped, never verbatim" + + from shared.db_dialects import get_dialect + + kwargs = get_dialect("snowflake")._connect_kwargs(uri) + assert kwargs["authenticator"] == "SNOWFLAKE_JWT" + assert isinstance(kwargs["private_key"], bytes) + assert "password" not in kwargs + + +def test_create_uri_strips_reserved_extra_params(): + connection = DataBaseConnection( + database_type="snowflake", + username="u", + database="D", + extra_params={"account": "acct", "auth_method": "key_pair", "private_key": "evil"}, + ) + # Reserved keys are stripped before the splat, so no TypeError and no auth override. + assert connection.create_uri() == "snowflake://u@acct/D" diff --git a/shared/CLAUDE.md b/shared/CLAUDE.md index 7cbd64eef..021443db0 100644 --- a/shared/CLAUDE.md +++ b/shared/CLAUDE.md @@ -17,7 +17,7 @@ It is a Poetry package — `{ include = "shared" }` in root `pyproject.toml`, wi - `models.py` — standalone SQLAlchemy models (`FlowRun`, `FlowSchedule`, `FlowRegistration`, `CatalogTable`, `ScheduleTriggerTable`, `SchedulerLock`) on their own `Base`. - `artifact_storage.py` — `ArtifactStorageBackend` ABC (`prepare_upload`/`prepare_download`/`delete`/`exists`) + `SharedFilesystemStorage` / `S3Storage` (presigned-URL) backends, returning `UploadTarget` / `DownloadSource`. Kernel moves blob bytes via presigned URLs; Core stays metadata-only. - `delta_utils.py` / `delta_models.py` — dependency-light Delta-log helpers (`make_json_safe`, `format_delta_timestamp`, `get_delta_size_bytes`, `validate_catalog_path`, plus `write_delta` / `merge_into_delta`, and the SCD2 primitives `scd2_into_delta` / `scd2_surrogate_keys` / `Scd2Result` — one atomic close+insert MERGE per write, with a frozen `sha256-v1` surrogate-key encoding) + Pydantic `DeltaVersionCommit` / `SourceTableVersion`. -- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), `snowflake.py` (native snowflake-connector-python end to end — no connectorx, no SQLAlchemy: Arrow-fetch reads cast through a `cursor.describe()` type-code map for predicted==materialized parity with a row-based fallback when the cursor lacks Arrow (fakesnow), account/warehouse/role arrive via the connection's guarded `extra_params` — `base.is_blocked_extra_param` keys can never override auth — port-less account URIs, qmark `executemany` writes, information_schema browse; behavioral tests run against fakesnow, no Docker), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). `DialectInfo` also serves per-dialect `extra_fields`/`hidden_fields` (from `DbDialect` ClassVars) so a new connection shape renders in the frontend forms with zero frontend changes. Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver), `snowflake.py` (native-driver server dialect with its own connection shape), or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. +- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), `snowflake.py` (native snowflake-connector-python end to end — no connectorx, no SQLAlchemy: Arrow-fetch reads cast through a `cursor.describe()` type-code map for predicted==materialized parity with a row-based fallback when the cursor lacks Arrow (fakesnow), account/warehouse/role arrive via the connection's guarded `extra_params` — `base.is_blocked_extra_param` keys can never override auth — port-less account URIs, qmark `executemany` writes, information_schema browse; first-class key-pair (JWT) auth: `build_uri(auth_method="key_pair", private_key=)` omits the password and appends dialect-generated `authenticator=snowflake_jwt&private_key=` **after** the blocked-key filter, `_connect_kwargs` decodes back to unencrypted PKCS#8 DER bytes for the connector (`auth_methods` ClassVar declares support; the base `_check_auth_supported` backstop rejects key material on dialects without it — sqlite/duckdb overrides call it explicitly since they skip super); behavioral tests run against fakesnow, no Docker — fakesnow ignores auth kwargs, so key-pair correctness = DER unit tests + core's live leg), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). `DialectInfo` also serves per-dialect `extra_fields`/`hidden_fields`/`auth_methods` (from `DbDialect` ClassVars) so a new connection shape or auth method renders in the frontend forms with zero frontend changes. Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver), `snowflake.py` (native-driver server dialect with its own connection shape), or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. - `sql_utils.py` — `construct_sql_uri`, `get_sqlalchemy_uri`, `SQLALCHEMY_DRIVER_MAP` (caller passes an already-decrypted password); thin dispatchers over `db_dialects` since the registry landed. - `cloud_storage/` — GCS/S3/ADLS helpers: `storage_options.py` (`build_*_storage_options`), `writers.py` (`write_to_cloud` + per-format parquet/csv/json/delta writers), `directory.py` (first-file listing per backend), `uri.py` (scheme list + `parse_uri`/`uri_parent`/`uri_join`/`canonical_scheme` — `pathlib` corrupts `scheme://`, so URI path maths lives here and `catalog/storage_backend.py` imports the scheme list from it), `browse.py` (one-level listing for the storage-browser UI: `BrowseEntry`/`BrowseResult`, `browse_support`, `list_cloud_uri`; boto3 `Delimiter="/"` for S3, `walk_blobs` for ADLS, `gcsfs.ls` for GCS — a refused bucket list is reported as `root_listing_denied`, not an error, and provider exceptions are translated to `BrowseError` subclasses carrying their own status + `error_code`), `gcs.py`, `utils.py`. - `kafka/` — `consumer.py` (`read_kafka_source`, `infer_topic_schema`, `commit_offsets`, `make_kafka_commit_callback`), `models.py`, `deserializers.py` (`get_deserializer`, JSON deserializer). diff --git a/shared/db_dialects/__init__.py b/shared/db_dialects/__init__.py index c2c8f6f2c..9b9224bc5 100644 --- a/shared/db_dialects/__init__.py +++ b/shared/db_dialects/__init__.py @@ -90,6 +90,7 @@ class DialectInfo(BaseModel): available: bool extra_fields: list[DialectFieldInfo] = [] hidden_fields: list[str] = [] + auth_methods: list[str] = ["password"] def get_dialect(name: str) -> DbDialect: @@ -128,6 +129,7 @@ def dialect_catalog() -> list[DialectInfo]: available=d.is_available(), extra_fields=[DialectFieldInfo(name=f.name, label=f.label, required=f.required) for f in d.extra_fields], hidden_fields=list(d.hidden_fields), + auth_methods=list(d.auth_methods), ) for d in iter_dialects() ] diff --git a/shared/db_dialects/base.py b/shared/db_dialects/base.py index 05a95f3d6..024c78e4e 100644 --- a/shared/db_dialects/base.py +++ b/shared/db_dialects/base.py @@ -47,6 +47,7 @@ class *is* the generic dialect — its method bodies are the historical "port", "database", "dbname", + "auth_method", "authenticator", "token", "insecure_mode", @@ -85,6 +86,10 @@ class DbDialect: # so a new connection shape needs no frontend changes. extra_fields: ClassVar[tuple[DialectField, ...]] = () hidden_fields: ClassVar[tuple[str, ...]] = () + # Authentication methods this dialect's build_uri understands. "password" is + # the implicit default everywhere; dialects opting into more (e.g. Snowflake + # key-pair JWT) extend this and handle the corresponding build_uri params. + auth_methods: ClassVar[tuple[str, ...]] = ("password",) @property def uri_scheme(self) -> str: @@ -94,6 +99,13 @@ def is_available(self) -> bool: """Whether the driver stack for this dialect is importable.""" return True + def _check_auth_supported(self, auth_method: str | None, private_key: str | None) -> None: + """Refuse credentials this dialect cannot honor — silent ignoring is worse than an error.""" + if auth_method not in (None, "", "password") and auth_method not in self.auth_methods: + raise ValueError(f"{self.display_name} does not support auth method {auth_method!r}") + if private_key and "key_pair" not in self.auth_methods: + raise ValueError(f"{self.display_name} does not support private-key (key pair) authentication") + def build_uri( self, *, @@ -104,11 +116,15 @@ def build_uri( database: str | None = None, ssl_enabled: bool = False, connect_timeout: int | None = None, + auth_method: str | None = None, + private_key: str | None = None, + private_key_passphrase: str | None = None, **kwargs, ) -> str: """Build a base (connectorx-style) URI. ``password`` is a plain string.""" from urllib.parse import quote_plus + self._check_auth_supported(auth_method, private_key) if not host: raise ValueError("Host is required to create a URI") diff --git a/shared/db_dialects/builtin.py b/shared/db_dialects/builtin.py index 39b354be1..7f83d257a 100644 --- a/shared/db_dialects/builtin.py +++ b/shared/db_dialects/builtin.py @@ -32,7 +32,8 @@ class SQLiteDialect(DbDialect): file_based = True sqlglot_name = "sqlite" - def build_uri(self, *, host=None, database=None, **kwargs) -> str: + def build_uri(self, *, host=None, database=None, auth_method=None, private_key=None, **kwargs) -> str: + self._check_auth_supported(auth_method, private_key) path = database or host or "./database.db" # Strip sqlite:/// prefix if the full URI was passed as the path if path.startswith("sqlite:///"): diff --git a/shared/db_dialects/duckdb.py b/shared/db_dialects/duckdb.py index d8cdd578e..3065dbf5b 100644 --- a/shared/db_dialects/duckdb.py +++ b/shared/db_dialects/duckdb.py @@ -64,7 +64,8 @@ def is_available(self) -> bool: return False return True - def build_uri(self, *, host=None, database=None, **kwargs) -> str: + def build_uri(self, *, host=None, database=None, auth_method=None, private_key=None, **kwargs) -> str: + self._check_auth_supported(auth_method, private_key) path = database or host or "./database.duckdb" if path.startswith(_URI_PREFIX): path = path[len(_URI_PREFIX) :] diff --git a/shared/db_dialects/snowflake.py b/shared/db_dialects/snowflake.py index 8ff632eb9..c815b5a91 100644 --- a/shared/db_dialects/snowflake.py +++ b/shared/db_dialects/snowflake.py @@ -13,6 +13,15 @@ ``build_uri`` drops any blocked key (``base.is_blocked_extra_param``) so extra params can never override credentials. +Key-pair (JWT) auth is first-class, not an extra param: +``build_uri(auth_method="key_pair", private_key=)`` omits the +password from the URI userinfo and appends dialect-generated +``authenticator=snowflake_jwt&private_key=`` (plus an optional +b64 passphrase) *after* the blocked-key filter, so user extra_params can never +inject them. ``_connect_kwargs`` decodes the PEM back and hands the connector +unencrypted PKCS#8 DER bytes with ``authenticator="SNOWFLAKE_JWT"``; key +material only ever exists in memory, never on disk. + Fast schema uses ``cursor.describe()`` — the query is compiled server-side but never executed — plus a dialect-local map keyed on the connector's type codes. ``read`` casts the fetched frame through the same map (Snowflake's Arrow @@ -72,6 +81,7 @@ class SnowflakeDialect(DbDialect): DialectField("role", "Role"), ) hidden_fields: ClassVar[tuple[str, ...]] = ("host", "port", "ssl") + auth_methods: ClassVar[tuple[str, ...]] = ("password", "key_pair") def is_available(self) -> bool: try: @@ -90,11 +100,22 @@ def build_uri( database: str | None = None, ssl_enabled: bool = False, connect_timeout: int | None = None, + auth_method: str | None = None, + private_key: str | None = None, + private_key_passphrase: str | None = None, **kwargs, ) -> str: """Account-shaped URI; ``account`` comes from extra_params (``host`` is accepted as an alias).""" + import base64 from urllib.parse import quote_plus + self._check_auth_supported(auth_method, private_key) + use_key_pair = auth_method == "key_pair" or bool(private_key and auth_method in (None, "", "password")) + if auth_method == "key_pair" and not private_key: + raise ValueError("A private key is required for Snowflake key-pair authentication") + if use_key_pair: + password = None + account = kwargs.pop("account", None) or host if not account: raise ValueError("Account is required to create a Snowflake URI") @@ -110,17 +131,27 @@ def build_uri( if database: uri += f"/{database}" params = {k: v for k, v in kwargs.items() if v is not None and not is_blocked_extra_param(k)} + if use_key_pair: + # Dialect-generated auth params, added after the blocked-key filter so + # user extra_params can never inject or override them. + params["authenticator"] = "snowflake_jwt" + params["private_key"] = base64.urlsafe_b64encode(private_key.encode("utf-8")).decode("ascii") + if private_key_passphrase: + params["private_key_passphrase"] = base64.urlsafe_b64encode( + private_key_passphrase.encode("utf-8") + ).decode("ascii") if params: uri += "?" + "&".join(f"{key}={quote_plus(str(value))}" for key, value in params.items()) return uri - @staticmethod - def _connect_kwargs(uri: str) -> dict[str, str]: + @classmethod + def _connect_kwargs(cls, uri: str) -> dict[str, Any]: + import base64 from urllib.parse import parse_qsl, unquote, urlparse parsed = urlparse(uri) # netloc split instead of .hostname: account identifiers should keep their case. - kwargs: dict[str, str] = {"account": unquote(parsed.netloc.rsplit("@", 1)[-1])} + kwargs: dict[str, Any] = {"account": unquote(parsed.netloc.rsplit("@", 1)[-1])} if parsed.username: kwargs["user"] = unquote(parsed.username) if parsed.password: @@ -128,11 +159,38 @@ def _connect_kwargs(uri: str) -> dict[str, str]: database = parsed.path.lstrip("/") if database: kwargs["database"] = unquote(database) - for key, value in parse_qsl(parsed.query): - if key in ("warehouse", "role", "schema") and value: - kwargs[key] = value + query = dict(parse_qsl(parsed.query)) + for key in ("warehouse", "role", "schema"): + if query.get(key): + kwargs[key] = query[key] + if query.get("authenticator") == "snowflake_jwt" and query.get("private_key"): + pem = base64.urlsafe_b64decode(query["private_key"]) + passphrase = ( + base64.urlsafe_b64decode(query["private_key_passphrase"]) + if query.get("private_key_passphrase") + else None + ) + kwargs["private_key"] = cls._private_key_der(pem, passphrase) + kwargs["authenticator"] = "SNOWFLAKE_JWT" + kwargs.pop("password", None) return kwargs + @staticmethod + def _private_key_der(pem: bytes, passphrase: bytes | None) -> bytes: + """PEM text (optionally passphrase-encrypted) -> unencrypted PKCS#8 DER bytes. + + snowflake-connector-python accepts DER bytes directly; decrypting here keeps + the passphrase out of the connector call and the key out of any file. + """ + from cryptography.hazmat.primitives import serialization + + key = serialization.load_pem_private_key(pem, password=passphrase) + return key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + def _connect(self, uri: str): import snowflake.connector diff --git a/shared/sql_utils.py b/shared/sql_utils.py index bfbea21a8..f377bfa0f 100644 --- a/shared/sql_utils.py +++ b/shared/sql_utils.py @@ -27,6 +27,9 @@ def construct_sql_uri( url: str | None = None, ssl_enabled: bool = False, connect_timeout: int | None = None, + auth_method: str | None = None, + private_key: str | None = None, + private_key_passphrase: str | None = None, **kwargs, ) -> str: """ @@ -42,6 +45,10 @@ def construct_sql_uri( url: Complete database URL (overrides other parameters if provided) ssl_enabled: Adds sslmode=require for postgres-family databases connect_timeout: Connection timeout in seconds (postgres-family only) + auth_method: Authentication method; must be in the dialect's auth_methods + ("password" or None everywhere, "key_pair" where supported) + private_key: Private key PEM text as a plain string (caller handles decryption) + private_key_passphrase: Optional passphrase for an encrypted private key **kwargs: Additional connection parameters appended as query string Returns: @@ -61,6 +68,9 @@ def construct_sql_uri( database=database, ssl_enabled=ssl_enabled, connect_timeout=connect_timeout, + auth_method=auth_method, + private_key=private_key, + private_key_passphrase=private_key_passphrase, **kwargs, ) diff --git a/shared/tests/db_dialects/test_dialect_contract.py b/shared/tests/db_dialects/test_dialect_contract.py index b1c148869..cddca7787 100644 --- a/shared/tests/db_dialects/test_dialect_contract.py +++ b/shared/tests/db_dialects/test_dialect_contract.py @@ -49,6 +49,7 @@ def test_metadata_sanity(dialect): assert isinstance(dialect.is_available(), bool) if not dialect.is_available(): assert dialect.install_hint + assert "password" in dialect.auth_methods @pytest.mark.parametrize("dialect", DIALECTS, ids=_ids) diff --git a/shared/tests/db_dialects/test_snowflake_dialect.py b/shared/tests/db_dialects/test_snowflake_dialect.py index 7c013fb1e..4cdeeff3f 100644 --- a/shared/tests/db_dialects/test_snowflake_dialect.py +++ b/shared/tests/db_dialects/test_snowflake_dialect.py @@ -43,6 +43,7 @@ def test_metadata(): assert [f.name for f in dialect.extra_fields] == ["account", "warehouse", "role"] assert dialect.extra_fields[0].required is True assert dialect.hidden_fields == ("host", "port", "ssl") + assert dialect.auth_methods == ("password", "key_pair") def test_build_uri_account_shape(): @@ -95,6 +96,107 @@ def test_connect_kwargs_round_trip(): } +def _throwaway_key(passphrase: bytes | None = None) -> tuple[str, bytes]: + """Generate an RSA key; returns (PEM text as build_uri receives it, expected unencrypted PKCS#8 DER).""" + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + encryption = serialization.BestAvailableEncryption(passphrase) if passphrase else serialization.NoEncryption() + pem = key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=encryption, + ) + der = key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + return pem.decode("ascii"), der + + +def test_build_uri_key_pair_shape(): + pem, _ = _throwaway_key() + uri = dialect.build_uri( + username="u", + password="should-not-appear", + database="d", + account="acct", + warehouse="WH", + auth_method="key_pair", + private_key=pem, + ) + assert "should-not-appear" not in uri, "key-pair URIs must not carry the password" + assert uri.startswith("snowflake://u@acct/d?") + assert "warehouse=WH" in uri + assert "authenticator=snowflake_jwt" in uri + assert "private_key=" in uri + assert "private_key_passphrase" not in uri + + +def test_build_uri_key_pair_requires_key(): + with pytest.raises(ValueError, match="private key is required"): + dialect.build_uri(username="u", database="d", account="acct", auth_method="key_pair") + + +def test_build_uri_private_key_implies_key_pair(): + pem, _ = _throwaway_key() + uri = dialect.build_uri(username="u", password="p", database="d", account="acct", private_key=pem) + assert "authenticator=snowflake_jwt" in uri + assert ":p@" not in uri + + +def test_connect_kwargs_key_pair_round_trip(): + pem, der = _throwaway_key() + uri = dialect.build_uri( + username="u", database="db1", account="Acct-Id", warehouse="WH", auth_method="key_pair", private_key=pem + ) + kwargs = dialect._connect_kwargs(uri) + assert kwargs["authenticator"] == "SNOWFLAKE_JWT" + assert kwargs["private_key"] == der, "PEM must round-trip to unencrypted PKCS#8 DER bytes" + assert "password" not in kwargs + assert kwargs["account"] == "Acct-Id" + assert kwargs["user"] == "u" + assert kwargs["warehouse"] == "WH" + + +def test_connect_kwargs_key_pair_passphrase_round_trip(): + passphrase = "s3cret pass+phrase" + pem, der = _throwaway_key(passphrase.encode("utf-8")) + uri = dialect.build_uri( + username="u", + database="db1", + account="acct", + auth_method="key_pair", + private_key=pem, + private_key_passphrase=passphrase, + ) + kwargs = dialect._connect_kwargs(uri) + assert kwargs["private_key"] == der, "encrypted PEM must decrypt to the same unencrypted PKCS#8 DER" + assert kwargs["authenticator"] == "SNOWFLAKE_JWT" + + +def test_other_dialects_reject_key_pair(): + from shared.db_dialects import get_dialect as _get + + postgres = _get("postgresql") + with pytest.raises(ValueError, match="does not support auth method"): + postgres.build_uri(host="h", username="u", password="p", auth_method="key_pair") + with pytest.raises(ValueError, match="key pair"): + postgres.build_uri(host="h", username="u", password="p", private_key="-----BEGIN PRIVATE KEY-----") + for file_based in ("sqlite", "duckdb"): + with pytest.raises(ValueError, match="does not support"): + _get(file_based).build_uri(database="/tmp/x.db", auth_method="key_pair") + + +def test_auth_and_key_material_keys_stay_blocked_as_extra_params(): + from shared.db_dialects import is_blocked_extra_param + + for key in ("auth_method", "authenticator", "token", "private_key", "private_key_file", "private_key_passphrase"): + assert is_blocked_extra_param(key), key + + def test_limit_query_is_plain_limit_and_parses_as_snowflake(): limited = dialect.limit_query("SELECT a, b FROM some_table", 5) assert limited == "SELECT a, b FROM some_table LIMIT 5" @@ -179,3 +281,28 @@ def test_fakesnow_browse(snow_uri): assert tables is not None and any(t.lower() == "t_browse" for t in tables) qualified = dialect.list_tables(snow_uri, None) assert qualified is not None and any(t.lower() == "public.t_browse" for t in qualified) + + +@pytest.fixture() +def snow_uri_key_pair(): + fakesnow = pytest.importorskip("fakesnow") + pem, _ = _throwaway_key() + with fakesnow.patch(): + yield dialect.build_uri( + username="u", + database="db1", + account="test", + auth_method="key_pair", + private_key=pem, + **{"schema": "public"}, + ) + + +def test_fakesnow_key_pair_uri_round_trip(snow_uri_key_pair): + # fakesnow ignores auth kwargs, so this proves the key-pair parse path does not + # break connections — auth correctness itself is covered by the DER unit tests + # and the live leg in test_snowflake_source.py. + df = pl.DataFrame({"x": [1, 2]}) + dialect.write(df, uri=snow_uri_key_pair, table_name="t_kp", if_exists="replace") + result = dialect.read("SELECT * FROM t_kp", snow_uri_key_pair, logger) + assert result.height == 2 From 12b7b7fac2f4dc43043e327c2a59bba0fde86c20 Mon Sep 17 00:00:00 2001 From: edwardvaneechoud Date: Wed, 5 Aug 2026 20:56:07 +0200 Subject: [PATCH 3/4] fix: ensure switching to password auth clears stray key references --- .../db_connections.py | 11 ++--- .../external_sources/test_snowflake_source.py | 40 +++++++++++++++++++ .../DatabaseConnectionSettings.vue | 13 +++++- .../DatabaseConnectionSettings.vue | 6 +++ 4 files changed, 64 insertions(+), 6 deletions(-) diff --git a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py index 7bab39f3c..3ef7b4d2d 100644 --- a/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py +++ b/flowfile_core/flowfile_core/flowfile/database_connection_manager/db_connections.py @@ -128,9 +128,8 @@ def update_database_connection(db: Session, connection: FullDatabaseConnection, db_connection.password_id = new_secret.id incoming_key = connection.private_key.get_secret_value() if connection.private_key else "" - keeps_key_material = connection.auth_method == "key_pair" or bool(incoming_key) - if keeps_key_material: - if connection.auth_method == "key_pair" and not incoming_key and db_connection.private_key_id is None: + if connection.auth_method == "key_pair": + if not incoming_key and db_connection.private_key_id is None: raise ValueError("key_pair authentication requires a private key.") # Rotate key material only when a non-empty value arrived (empty == keep existing, # mirroring password). _update_cloud_secret is generic despite the name. @@ -149,8 +148,10 @@ def update_database_connection(db: Session, connection: FullDatabaseConnection, user_id, ) else: - # Switching away from key-pair auth: detach AND delete the key secrets, or a - # rotated-away (possibly compromised) key would keep authenticating silently. + # Switching away from key-pair auth: detach AND delete the key secrets — and + # ignore any stray incoming key — or a rotated-away (possibly compromised) + # key would keep authenticating silently (build_uri infers key-pair from key + # presence, and the reference resolver forwards whatever the row links). stale_ids = [ secret_id for secret_id in (db_connection.private_key_id, db_connection.private_key_passphrase_id) diff --git a/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py index c9f50061d..40140cd42 100644 --- a/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py +++ b/flowfile_core/tests/flowfile/external_sources/test_snowflake_source.py @@ -269,6 +269,46 @@ def test_switching_to_password_detaches_and_deletes_key_secrets(self): finally: _cleanup(name) + def test_switching_to_password_drops_a_stray_incoming_key(self): + """auth_method="password" with a non-empty private_key in the same payload (e.g. a + client that didn't clear the hidden field): the stray key must be ignored AND the + existing key secrets deleted — build_uri infers key-pair from key presence, so a + persisted stray key would silently out-vote the password.""" + from flowfile_core.database.models import Secret + + name = "snowflake_kp_stray" + _cleanup(name) + try: + with get_db_context() as db: + store_database_connection(db, _key_pair_connection(name, _make_pem(), "pp"), USER_ID) + with get_db_context() as db: + row = get_database_connection(db, name, USER_ID) + old_key_ids = [row.private_key_id, row.private_key_passphrase_id] + + flipped = FullDatabaseConnection( + connection_name=name, + database_type="snowflake", + username="svc_user", + password=SecretStr("new-pass"), + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="password", + private_key=SecretStr(_make_pem()), + private_key_passphrase=SecretStr("stray-pass"), + ) + with get_db_context() as db: + update_database_connection(db, flipped, USER_ID) + with get_db_context() as db: + row = get_database_connection(db, name, USER_ID) + assert row.private_key_id is None + assert row.private_key_passphrase_id is None + assert db.query(Secret).filter(Secret.id.in_(old_key_ids)).count() == 0 + reloaded = get_database_connection_schema(db, name, USER_ID) + assert reloaded.private_key is None, "a stray incoming key must never be persisted" + assert reloaded.private_key_passphrase is None + finally: + _cleanup(name) + def test_update_to_key_pair_without_key_is_rejected(self): name = "snowflake_kp_bad_flip" _cleanup(name) diff --git a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue index 7f5c8c7d4..c0e363332 100644 --- a/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue +++ b/flowfile_frontend/src/renderer/app/components/nodes/node-types/elements/databaseReader/DatabaseConnectionSettings.vue @@ -42,7 +42,7 @@ id="auth-method" :value="modelValue.auth_method || 'password'" class="form-control" - @change="(e: Event) => updateField('auth_method', (e.target as HTMLSelectElement).value)" + @change="(e: Event) => updateAuthMethod((e.target as HTMLSelectElement).value)" >
-
+
+
+ + +
+ +
+ +
+ + +
+
+ +
+ + +
+ +
+ + +
+ +
+ + +
+ +
+ +
+
(); const emit = defineEmits<{ (e: "submit", connection: FullDatabaseConnection): void; (e: "cancel"): void; + (e: "oauthChanged"): void; }>(); const { dialects, isFileBased, defaultPort, extraFields, isFieldHidden, authMethods } = @@ -197,6 +277,7 @@ const { dialects, isFileBased, defaultPort, extraFields, isFieldHidden, authMeth const AUTH_METHOD_LABELS: Record = { password: "Password", key_pair: "Key pair (JWT)", + oauth: "Single sign-on (OAuth)", }; const defaultConnection = (): FullDatabaseConnection => ({ @@ -212,6 +293,11 @@ const defaultConnection = (): FullDatabaseConnection => ({ authMethod: "password", privateKey: "", privateKeyPassphrase: "", + oauthClientId: "", + oauthClientSecret: "", + oauthAuthorizeEndpoint: "", + oauthTokenEndpoint: "", + oauthRedirectUri: "", }); const connection = ref( @@ -261,12 +347,22 @@ watch( } connection.value.privateKey = ""; connection.value.privateKeyPassphrase = ""; + clearOauthFields(); } }, ); const showPassword = ref(false); const showPassphrase = ref(false); +const showClientSecret = ref(false); + +const clearOauthFields = () => { + connection.value.oauthClientId = ""; + connection.value.oauthClientSecret = ""; + connection.value.oauthAuthorizeEndpoint = ""; + connection.value.oauthTokenEndpoint = ""; + connection.value.oauthRedirectUri = ""; +}; const isFileBasedType = computed(() => isFileBased(connection.value.databaseType)); @@ -284,10 +380,14 @@ const authMethodModel = computed({ connection.value.privateKey = ""; connection.value.privateKeyPassphrase = ""; } + if (value !== "oauth") { + clearOauthFields(); + } }, }); const usesKeyPair = computed(() => authMethodModel.value === "key_pair"); +const usesOauth = computed(() => authMethodModel.value === "oauth"); const authMethodLabel = (method: string): string => AUTH_METHOD_LABELS[method] ?? method; @@ -315,7 +415,10 @@ const isValid = computed(() => { } const credentialFilled = usesKeyPair.value ? props.isEditing || !!connection.value.privateKey - : props.isEditing || !!connection.value.password; + : usesOauth.value + ? !!connection.value.oauthClientId && + (props.isEditing || !!connection.value.oauthClientSecret) + : props.isEditing || !!connection.value.password; return ( !!connection.value.connectionName && !!connection.value.username && @@ -340,6 +443,10 @@ const submitForm = () => { diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts index f72d9ed06..e0fa02601 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/api.ts @@ -30,6 +30,11 @@ const toPythonFormat = (connection: FullDatabaseConnection): PythonFullDatabaseC // empty key-secret rows server-side; blank means "no key" / "keep existing". private_key: connection.privateKey || undefined, private_key_passphrase: connection.privateKeyPassphrase || undefined, + oauth_client_id: connection.oauthClientId || undefined, + oauth_client_secret: connection.oauthClientSecret || undefined, + oauth_authorize_endpoint: connection.oauthAuthorizeEndpoint || undefined, + oauth_token_endpoint: connection.oauthTokenEndpoint || undefined, + oauth_redirect_uri: connection.oauthRedirectUri || undefined, }; }; @@ -63,6 +68,11 @@ export const convertConnectionInterfacePytoTs = ( database: pythonConnectionInterface.database, extraParams: pythonConnectionInterface.extra_params, authMethod: pythonConnectionInterface.auth_method ?? undefined, + oauthClientId: pythonConnectionInterface.oauth_client_id ?? undefined, + oauthAuthorizeEndpoint: pythonConnectionInterface.oauth_authorize_endpoint ?? undefined, + oauthTokenEndpoint: pythonConnectionInterface.oauth_token_endpoint ?? undefined, + oauthRedirectUri: pythonConnectionInterface.oauth_redirect_uri ?? undefined, + oauthConnected: pythonConnectionInterface.oauth_connected ?? false, id: pythonConnectionInterface.id, access: pythonConnectionInterface.access, }; @@ -82,9 +92,25 @@ export const convertConnectionInterfaceTstoPy = ( url: dbConnectionInterface.url, extra_params: dbConnectionInterface.extraParams, auth_method: dbConnectionInterface.authMethod, + oauth_client_id: dbConnectionInterface.oauthClientId, + oauth_authorize_endpoint: dbConnectionInterface.oauthAuthorizeEndpoint, + oauth_token_endpoint: dbConnectionInterface.oauthTokenEndpoint, + oauth_redirect_uri: dbConnectionInterface.oauthRedirectUri, + oauth_connected: dbConnectionInterface.oauthConnected, }; }; +/** + * Starts the OAuth sign-in flow for a stored connection; returns the IdP + * authorize URL to open in a popup / the system browser. + */ +export const startDbOauthApi = async (connectionName: string): Promise => { + const response = await axios.get<{ auth_url: string }>(`${API_BASE_URL}/oauth/start`, { + params: { connection_name: connectionName }, + }); + return response.data.auth_url; +}; + /** * Creates a new database connection via the API. * @param connectionData - The database connection configuration to add. diff --git a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts index 933a28332..7f303a6e0 100644 --- a/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts +++ b/flowfile_frontend/src/renderer/app/views/DatabaseView/databaseConnectionTypes.ts @@ -20,6 +20,11 @@ export interface PythonFullDatabaseConnection { auth_method?: string | null; private_key?: string; private_key_passphrase?: string; + oauth_client_id?: string | null; + oauth_client_secret?: string; + oauth_authorize_endpoint?: string | null; + oauth_token_endpoint?: string | null; + oauth_redirect_uri?: string | null; } export interface FullDatabaseConnection { @@ -36,6 +41,11 @@ export interface FullDatabaseConnection { authMethod?: string; privateKey?: string; privateKeyPassphrase?: string; + oauthClientId?: string; + oauthClientSecret?: string; + oauthAuthorizeEndpoint?: string; + oauthTokenEndpoint?: string; + oauthRedirectUri?: string; } export interface PythonFullDatabaseConnectionInterface { @@ -49,6 +59,11 @@ export interface PythonFullDatabaseConnectionInterface { url?: string; extra_params?: Record | null; auth_method?: string | null; + oauth_client_id?: string | null; + oauth_authorize_endpoint?: string | null; + oauth_token_endpoint?: string | null; + oauth_redirect_uri?: string | null; + oauth_connected?: boolean; id?: number; access?: AccessInfo | null; } @@ -64,6 +79,11 @@ export interface FullDatabaseConnectionInterface { url?: string; extraParams?: Record | null; authMethod?: string; + oauthClientId?: string; + oauthAuthorizeEndpoint?: string; + oauthTokenEndpoint?: string; + oauthRedirectUri?: string; + oauthConnected?: boolean; id?: number; access?: AccessInfo | null; } diff --git a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py index 77886c307..c260356b1 100644 --- a/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py +++ b/flowfile_worker/flowfile_worker/external_sources/sql_source/models.py @@ -5,7 +5,7 @@ from flowfile_worker.secrets import decrypt_secret from shared.sql_utils import construct_sql_uri, get_sqlalchemy_uri -_RESERVED_EXTRA_PARAMS = ("auth_method", "private_key", "private_key_passphrase") +_RESERVED_EXTRA_PARAMS = ("auth_method", "private_key", "private_key_passphrase", "oauth_token") class DataBaseConnection(BaseModel): @@ -23,6 +23,7 @@ class DataBaseConnection(BaseModel): auth_method: str | None = None # None == password auth private_key: SecretStr | None = None # Encrypted private key PEM (key-pair auth) private_key_passphrase: SecretStr | None = None # Encrypted private-key passphrase + oauth_token: SecretStr | None = None # Encrypted short-lived OAuth access token (core refreshed it) def get_decrypted_secret(self) -> SecretStr: return decrypt_secret(self.password.get_secret_value()) @@ -58,6 +59,7 @@ def create_uri(self) -> str: auth_method=self.auth_method, private_key=self._decrypt(self.private_key), private_key_passphrase=self._decrypt(self.private_key_passphrase), + oauth_token=self._decrypt(self.oauth_token), **extra_params, ) diff --git a/flowfile_worker/tests/external_sources/test_dialect_ports.py b/flowfile_worker/tests/external_sources/test_dialect_ports.py index 85e9accc7..53fbce274 100644 --- a/flowfile_worker/tests/external_sources/test_dialect_ports.py +++ b/flowfile_worker/tests/external_sources/test_dialect_ports.py @@ -127,3 +127,38 @@ def test_create_uri_strips_reserved_extra_params(): ) # Reserved keys are stripped before the splat, so no TypeError and no auth override. assert connection.create_uri() == "snowflake://u@acct/D" + + +def test_snowflake_create_uri_oauth_decrypts_token_and_omits_password(): + from flowfile_worker.secrets import encrypt_secret + + token = "ver:1-hint:abc.DEF/ghi+jk==" + connection = DataBaseConnection( + database_type="snowflake", + username="svc", + database="ANALYTICS", + extra_params={"account": "myorg-myaccount"}, + auth_method="oauth", + oauth_token=encrypt_secret(token), + ) + uri = connection.create_uri() + assert uri.startswith("snowflake://svc@myorg-myaccount/ANALYTICS?") + assert "authenticator=oauth" in uri + assert token not in uri, "the token must be b64-wrapped, never verbatim" + + from shared.db_dialects import get_dialect + + kwargs = get_dialect("snowflake")._connect_kwargs(uri) + assert kwargs["authenticator"] == "oauth" + assert kwargs["token"] == token + assert "password" not in kwargs + + +def test_create_uri_strips_reserved_oauth_extra_param(): + connection = DataBaseConnection( + database_type="snowflake", + username="u", + database="D", + extra_params={"account": "acct", "oauth_token": "evil"}, + ) + assert connection.create_uri() == "snowflake://u@acct/D" diff --git a/shared/CLAUDE.md b/shared/CLAUDE.md index 021443db0..77f057b26 100644 --- a/shared/CLAUDE.md +++ b/shared/CLAUDE.md @@ -17,7 +17,7 @@ It is a Poetry package — `{ include = "shared" }` in root `pyproject.toml`, wi - `models.py` — standalone SQLAlchemy models (`FlowRun`, `FlowSchedule`, `FlowRegistration`, `CatalogTable`, `ScheduleTriggerTable`, `SchedulerLock`) on their own `Base`. - `artifact_storage.py` — `ArtifactStorageBackend` ABC (`prepare_upload`/`prepare_download`/`delete`/`exists`) + `SharedFilesystemStorage` / `S3Storage` (presigned-URL) backends, returning `UploadTarget` / `DownloadSource`. Kernel moves blob bytes via presigned URLs; Core stays metadata-only. - `delta_utils.py` / `delta_models.py` — dependency-light Delta-log helpers (`make_json_safe`, `format_delta_timestamp`, `get_delta_size_bytes`, `validate_catalog_path`, plus `write_delta` / `merge_into_delta`, and the SCD2 primitives `scd2_into_delta` / `scd2_surrogate_keys` / `Scd2Result` — one atomic close+insert MERGE per write, with a frozen `sha256-v1` surrogate-key encoding) + Pydantic `DeltaVersionCommit` / `SourceTableVersion`. -- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), `snowflake.py` (native snowflake-connector-python end to end — no connectorx, no SQLAlchemy: Arrow-fetch reads cast through a `cursor.describe()` type-code map for predicted==materialized parity with a row-based fallback when the cursor lacks Arrow (fakesnow), account/warehouse/role arrive via the connection's guarded `extra_params` — `base.is_blocked_extra_param` keys can never override auth — port-less account URIs, qmark `executemany` writes, information_schema browse; first-class key-pair (JWT) auth: `build_uri(auth_method="key_pair", private_key=)` omits the password and appends dialect-generated `authenticator=snowflake_jwt&private_key=` **after** the blocked-key filter, `_connect_kwargs` decodes back to unencrypted PKCS#8 DER bytes for the connector (`auth_methods` ClassVar declares support; the base `_check_auth_supported` backstop rejects key material on dialects without it — sqlite/duckdb overrides call it explicitly since they skip super); behavioral tests run against fakesnow, no Docker — fakesnow ignores auth kwargs, so key-pair correctness = DER unit tests + core's live leg), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). `DialectInfo` also serves per-dialect `extra_fields`/`hidden_fields`/`auth_methods` (from `DbDialect` ClassVars) so a new connection shape or auth method renders in the frontend forms with zero frontend changes. Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver), `snowflake.py` (native-driver server dialect with its own connection shape), or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. +- `db_dialects/` — the database-dialect registry: `DbDialect` base (base behavior == the historical generic code paths), `builtin.py` (postgres/mysql/sqlite + the `GenericDialect` compat valve for legacy free-string types), `duckdb.py` (native driver, read_only reads, LIMIT-0 fast schema), `mssql.py` (SQL Server: pymssql-only reads — connectorx's tiberius backend segfaults across fresh reader threads, so it must never enter the hedged race — `TOP n` limits, `sp_describe_first_result_set` fast schema applied as `schema_overrides` on read for predicted==materialized parity, Object-producing types projected to NVARCHAR), `snowflake.py` (native snowflake-connector-python end to end — no connectorx, no SQLAlchemy: Arrow-fetch reads cast through a `cursor.describe()` type-code map for predicted==materialized parity with a row-based fallback when the cursor lacks Arrow (fakesnow), account/warehouse/role arrive via the connection's guarded `extra_params` — `base.is_blocked_extra_param` keys can never override auth — port-less account URIs, qmark `executemany` writes, information_schema browse; first-class key-pair (JWT) auth: `build_uri(auth_method="key_pair", private_key=)` omits the password and appends dialect-generated `authenticator=snowflake_jwt&private_key=` **after** the blocked-key filter, `_connect_kwargs` decodes back to unencrypted PKCS#8 DER bytes for the connector (`auth_methods` ClassVar declares support; the base `_check_auth_supported` backstop rejects key material on dialects without it — sqlite/duckdb overrides call it explicitly since they skip super); behavioral tests run against fakesnow, no Docker — fakesnow ignores auth kwargs, so key-pair correctness = DER unit tests + core's live leg; OAuth (SSO) auth mirrors the same shape: `build_uri(auth_method="oauth", oauth_token=)` appends `authenticator=oauth&token=` after the blocked-key filter and `_connect_kwargs` decodes it back — core mints/refreshes the token (see `shared/snowflake_oauth.py`, the pure httpx token-endpoint client with `TokenResponse` + typed `SnowflakeOAuthError`), this layer only transports it), registry API (`get_dialect`, `get_dialect_or_generic`, `KNOWN_DIALECT_NAMES`, `dialect_catalog` → `GET /db_dialects`, `read_sql`). `DialectInfo` also serves per-dialect `extra_fields`/`hidden_fields`/`auth_methods` (from `DbDialect` ClassVars) so a new connection shape or auth method renders in the frontend forms with zero frontend changes. Heavy driver imports stay function-local; dialect methods receive already-decrypted plain strings. Adding a connector: copy `duckdb.py` (file-based/native-driver), `snowflake.py` (native-driver server dialect with its own connection shape), or override metadata + `limit_query` (connectorx-supported server dialects), register in `_BUILTIN_DIALECTS`, and let `shared/tests/db_dialects/test_dialect_contract.py` run the shared contract over it. - `sql_utils.py` — `construct_sql_uri`, `get_sqlalchemy_uri`, `SQLALCHEMY_DRIVER_MAP` (caller passes an already-decrypted password); thin dispatchers over `db_dialects` since the registry landed. - `cloud_storage/` — GCS/S3/ADLS helpers: `storage_options.py` (`build_*_storage_options`), `writers.py` (`write_to_cloud` + per-format parquet/csv/json/delta writers), `directory.py` (first-file listing per backend), `uri.py` (scheme list + `parse_uri`/`uri_parent`/`uri_join`/`canonical_scheme` — `pathlib` corrupts `scheme://`, so URI path maths lives here and `catalog/storage_backend.py` imports the scheme list from it), `browse.py` (one-level listing for the storage-browser UI: `BrowseEntry`/`BrowseResult`, `browse_support`, `list_cloud_uri`; boto3 `Delimiter="/"` for S3, `walk_blobs` for ADLS, `gcsfs.ls` for GCS — a refused bucket list is reported as `root_listing_denied`, not an error, and provider exceptions are translated to `BrowseError` subclasses carrying their own status + `error_code`), `gcs.py`, `utils.py`. - `kafka/` — `consumer.py` (`read_kafka_source`, `infer_topic_schema`, `commit_offsets`, `make_kafka_commit_callback`), `models.py`, `deserializers.py` (`get_deserializer`, JSON deserializer). diff --git a/shared/db_dialects/base.py b/shared/db_dialects/base.py index 024c78e4e..a21de42d6 100644 --- a/shared/db_dialects/base.py +++ b/shared/db_dialects/base.py @@ -50,6 +50,7 @@ class *is* the generic dialect — its method bodies are the historical "auth_method", "authenticator", "token", + "oauth_token", "insecure_mode", } ) @@ -99,12 +100,16 @@ def is_available(self) -> bool: """Whether the driver stack for this dialect is importable.""" return True - def _check_auth_supported(self, auth_method: str | None, private_key: str | None) -> None: + def _check_auth_supported( + self, auth_method: str | None, private_key: str | None, oauth_token: str | None = None + ) -> None: """Refuse credentials this dialect cannot honor — silent ignoring is worse than an error.""" if auth_method not in (None, "", "password") and auth_method not in self.auth_methods: raise ValueError(f"{self.display_name} does not support auth method {auth_method!r}") if private_key and "key_pair" not in self.auth_methods: raise ValueError(f"{self.display_name} does not support private-key (key pair) authentication") + if oauth_token and "oauth" not in self.auth_methods: + raise ValueError(f"{self.display_name} does not support OAuth token authentication") def build_uri( self, @@ -119,12 +124,13 @@ def build_uri( auth_method: str | None = None, private_key: str | None = None, private_key_passphrase: str | None = None, + oauth_token: str | None = None, **kwargs, ) -> str: """Build a base (connectorx-style) URI. ``password`` is a plain string.""" from urllib.parse import quote_plus - self._check_auth_supported(auth_method, private_key) + self._check_auth_supported(auth_method, private_key, oauth_token) if not host: raise ValueError("Host is required to create a URI") diff --git a/shared/db_dialects/builtin.py b/shared/db_dialects/builtin.py index 7f83d257a..78dbeecb4 100644 --- a/shared/db_dialects/builtin.py +++ b/shared/db_dialects/builtin.py @@ -33,7 +33,7 @@ class SQLiteDialect(DbDialect): sqlglot_name = "sqlite" def build_uri(self, *, host=None, database=None, auth_method=None, private_key=None, **kwargs) -> str: - self._check_auth_supported(auth_method, private_key) + self._check_auth_supported(auth_method, private_key, kwargs.pop("oauth_token", None)) path = database or host or "./database.db" # Strip sqlite:/// prefix if the full URI was passed as the path if path.startswith("sqlite:///"): diff --git a/shared/db_dialects/duckdb.py b/shared/db_dialects/duckdb.py index 3065dbf5b..f72f78134 100644 --- a/shared/db_dialects/duckdb.py +++ b/shared/db_dialects/duckdb.py @@ -65,7 +65,7 @@ def is_available(self) -> bool: return True def build_uri(self, *, host=None, database=None, auth_method=None, private_key=None, **kwargs) -> str: - self._check_auth_supported(auth_method, private_key) + self._check_auth_supported(auth_method, private_key, kwargs.pop("oauth_token", None)) path = database or host or "./database.duckdb" if path.startswith(_URI_PREFIX): path = path[len(_URI_PREFIX) :] diff --git a/shared/db_dialects/snowflake.py b/shared/db_dialects/snowflake.py index c815b5a91..23a0fa5e3 100644 --- a/shared/db_dialects/snowflake.py +++ b/shared/db_dialects/snowflake.py @@ -22,6 +22,13 @@ unencrypted PKCS#8 DER bytes with ``authenticator="SNOWFLAKE_JWT"``; key material only ever exists in memory, never on disk. +OAuth (SSO) auth follows the same shape: ``build_uri(auth_method="oauth", +oauth_token=)`` omits the password and appends dialect-generated +``authenticator=oauth&token=`` after the blocked-key filter; +``_connect_kwargs`` decodes it back and hands the connector +``authenticator="oauth", token=...``. Token refresh happens in core — this +module only transports an already-minted short-lived access token. + Fast schema uses ``cursor.describe()`` — the query is compiled server-side but never executed — plus a dialect-local map keyed on the connector's type codes. ``read`` casts the fetched frame through the same map (Snowflake's Arrow @@ -81,7 +88,7 @@ class SnowflakeDialect(DbDialect): DialectField("role", "Role"), ) hidden_fields: ClassVar[tuple[str, ...]] = ("host", "port", "ssl") - auth_methods: ClassVar[tuple[str, ...]] = ("password", "key_pair") + auth_methods: ClassVar[tuple[str, ...]] = ("password", "key_pair", "oauth") def is_available(self) -> bool: try: @@ -103,17 +110,21 @@ def build_uri( auth_method: str | None = None, private_key: str | None = None, private_key_passphrase: str | None = None, + oauth_token: str | None = None, **kwargs, ) -> str: """Account-shaped URI; ``account`` comes from extra_params (``host`` is accepted as an alias).""" import base64 from urllib.parse import quote_plus - self._check_auth_supported(auth_method, private_key) + self._check_auth_supported(auth_method, private_key, oauth_token) use_key_pair = auth_method == "key_pair" or bool(private_key and auth_method in (None, "", "password")) if auth_method == "key_pair" and not private_key: raise ValueError("A private key is required for Snowflake key-pair authentication") - if use_key_pair: + use_oauth = auth_method == "oauth" + if use_oauth and not oauth_token: + raise ValueError("An access token is required for Snowflake OAuth authentication") + if use_key_pair or use_oauth: password = None account = kwargs.pop("account", None) or host @@ -140,6 +151,11 @@ def build_uri( params["private_key_passphrase"] = base64.urlsafe_b64encode( private_key_passphrase.encode("utf-8") ).decode("ascii") + if use_oauth: + # Dialect-generated auth params, after the blocked-key filter ("token" and + # "authenticator" are blocked keys, so user extra_params can never inject them). + params["authenticator"] = "oauth" + params["token"] = base64.urlsafe_b64encode(oauth_token.encode("utf-8")).decode("ascii") if params: uri += "?" + "&".join(f"{key}={quote_plus(str(value))}" for key, value in params.items()) return uri @@ -173,6 +189,10 @@ def _connect_kwargs(cls, uri: str) -> dict[str, Any]: kwargs["private_key"] = cls._private_key_der(pem, passphrase) kwargs["authenticator"] = "SNOWFLAKE_JWT" kwargs.pop("password", None) + elif query.get("authenticator") == "oauth" and query.get("token"): + kwargs["token"] = base64.urlsafe_b64decode(query["token"]).decode("utf-8") + kwargs["authenticator"] = "oauth" + kwargs.pop("password", None) return kwargs @staticmethod diff --git a/shared/snowflake_oauth.py b/shared/snowflake_oauth.py new file mode 100644 index 000000000..11460b814 --- /dev/null +++ b/shared/snowflake_oauth.py @@ -0,0 +1,130 @@ +"""OAuth token-endpoint client for Snowflake SSO connections. + +Pure HTTP helpers used by flowfile_core to exchange an authorization code and +to refresh access tokens for ``auth_method="oauth"`` database connections. +Works against both Snowflake's built-in OAuth server and an external IdP +(Okta / Entra ID / PingFederate): the caller passes the token endpoint, so +this module needs no knowledge of which flavor is configured. + +Secrets never enter ``shared``: every input is an already-decrypted plain +string and the caller encrypts whatever it persists. Client credentials are +sent as HTTP Basic auth (``client_secret_basic``), which Snowflake requires +and the major IdPs accept. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import httpx + +_TIMEOUT_SECONDS = 20.0 + + +class TokenResponse(NamedTuple): + """The useful subset of an OAuth token-endpoint response.""" + + access_token: str + expires_in: int | None + refresh_token: str | None + + +class SnowflakeOAuthError(Exception): + """A failed token-endpoint call. + + ``error`` carries the OAuth error code when the endpoint returned one + (e.g. ``invalid_grant`` for an expired or revoked refresh token) so the + caller can distinguish "re-authenticate" from transient failures. + """ + + def __init__(self, message: str, *, error: str | None = None, status_code: int | None = None): + super().__init__(message) + self.error = error + self.status_code = status_code + + @property + def requires_reauthentication(self) -> bool: + return self.error in ("invalid_grant", "invalid_token") + + +def derive_snowflake_endpoints(account: str) -> tuple[str, str]: + """Default authorize/token endpoints for Snowflake's built-in OAuth server.""" + base = f"https://{account}.snowflakecomputing.com/oauth" + return f"{base}/authorize", f"{base}/token-request" + + +def _post_token_request(token_endpoint: str, client_id: str, client_secret: str, data: dict[str, str]) -> TokenResponse: + try: + response = httpx.post( + token_endpoint, + data=data, + auth=(client_id, client_secret), + headers={"Accept": "application/json"}, + timeout=_TIMEOUT_SECONDS, + ) + except httpx.HTTPError as e: + raise SnowflakeOAuthError(f"Could not reach OAuth token endpoint: {e}") from e + + if response.status_code >= 400: + error_code = None + detail = response.text[:300] + try: + payload = response.json() + error_code = payload.get("error") + detail = payload.get("error_description") or payload.get("message") or detail + except ValueError: + pass + raise SnowflakeOAuthError( + f"OAuth token request failed ({response.status_code}): {detail}", + error=error_code, + status_code=response.status_code, + ) + + try: + payload = response.json() + except ValueError as e: + raise SnowflakeOAuthError("OAuth token endpoint returned a non-JSON response") from e + access_token = payload.get("access_token") + if not access_token: + raise SnowflakeOAuthError("OAuth token endpoint response is missing access_token") + expires_in = payload.get("expires_in") + return TokenResponse( + access_token=access_token, + expires_in=int(expires_in) if expires_in is not None else None, + refresh_token=payload.get("refresh_token") or None, + ) + + +def exchange_authorization_code( + token_endpoint: str, + client_id: str, + client_secret: str, + code: str, + redirect_uri: str, +) -> TokenResponse: + """Exchange an authorization code for tokens (the interactive callback leg).""" + return _post_token_request( + token_endpoint, + client_id, + client_secret, + {"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri}, + ) + + +def refresh_access_token( + token_endpoint: str, + client_id: str, + client_secret: str, + refresh_token: str, +) -> TokenResponse: + """Mint a short-lived access token from a stored refresh token. + + Some IdPs rotate the refresh token on use; when the response carries one, + the caller must persist it or the next refresh fails. + """ + return _post_token_request( + token_endpoint, + client_id, + client_secret, + {"grant_type": "refresh_token", "refresh_token": refresh_token}, + ) diff --git a/shared/sql_utils.py b/shared/sql_utils.py index f377bfa0f..51f6da0ef 100644 --- a/shared/sql_utils.py +++ b/shared/sql_utils.py @@ -30,6 +30,7 @@ def construct_sql_uri( auth_method: str | None = None, private_key: str | None = None, private_key_passphrase: str | None = None, + oauth_token: str | None = None, **kwargs, ) -> str: """ @@ -49,6 +50,7 @@ def construct_sql_uri( ("password" or None everywhere, "key_pair" where supported) private_key: Private key PEM text as a plain string (caller handles decryption) private_key_passphrase: Optional passphrase for an encrypted private key + oauth_token: OAuth access token as a plain string (caller handles refresh/decryption) **kwargs: Additional connection parameters appended as query string Returns: @@ -71,6 +73,7 @@ def construct_sql_uri( auth_method=auth_method, private_key=private_key, private_key_passphrase=private_key_passphrase, + oauth_token=oauth_token, **kwargs, ) diff --git a/shared/tests/db_dialects/test_snowflake_dialect.py b/shared/tests/db_dialects/test_snowflake_dialect.py index 4cdeeff3f..c306c113e 100644 --- a/shared/tests/db_dialects/test_snowflake_dialect.py +++ b/shared/tests/db_dialects/test_snowflake_dialect.py @@ -43,7 +43,7 @@ def test_metadata(): assert [f.name for f in dialect.extra_fields] == ["account", "warehouse", "role"] assert dialect.extra_fields[0].required is True assert dialect.hidden_fields == ("host", "port", "ssl") - assert dialect.auth_methods == ("password", "key_pair") + assert dialect.auth_methods == ("password", "key_pair", "oauth") def test_build_uri_account_shape(): @@ -193,10 +193,67 @@ def test_other_dialects_reject_key_pair(): def test_auth_and_key_material_keys_stay_blocked_as_extra_params(): from shared.db_dialects import is_blocked_extra_param - for key in ("auth_method", "authenticator", "token", "private_key", "private_key_file", "private_key_passphrase"): + for key in ( + "auth_method", + "authenticator", + "token", + "oauth_token", + "private_key", + "private_key_file", + "private_key_passphrase", + ): assert is_blocked_extra_param(key), key +def test_build_uri_oauth_shape(): + uri = dialect.build_uri( + username="u", + password="should-not-appear", + database="d", + account="acct", + warehouse="WH", + auth_method="oauth", + oauth_token="ver:1-hint:abc.DEF/ghi+jk==", + ) + assert "should-not-appear" not in uri, "OAuth URIs must not carry the password" + assert uri.startswith("snowflake://u@acct/d?") + assert "warehouse=WH" in uri + assert "authenticator=oauth" in uri + assert "token=" in uri + + +def test_build_uri_oauth_requires_token(): + with pytest.raises(ValueError, match="access token is required"): + dialect.build_uri(username="u", database="d", account="acct", auth_method="oauth") + + +def test_connect_kwargs_oauth_round_trip(): + token = "ver:1-hint:5A2b.C/d+e==" + uri = dialect.build_uri( + username="u", database="db1", account="Acct-Id", warehouse="WH", auth_method="oauth", oauth_token=token + ) + kwargs = dialect._connect_kwargs(uri) + assert kwargs["authenticator"] == "oauth" + assert kwargs["token"] == token + assert "password" not in kwargs + assert kwargs["account"] == "Acct-Id" + assert kwargs["user"] == "u" + assert kwargs["warehouse"] == "WH" + + +def test_other_dialects_reject_oauth(): + from shared.db_dialects import get_dialect as _get + + postgres = _get("postgresql") + with pytest.raises(ValueError, match="does not support auth method"): + postgres.build_uri(host="h", username="u", password="p", auth_method="oauth") + with pytest.raises(ValueError, match="OAuth token"): + postgres.build_uri(host="h", username="u", password="p", oauth_token="tok") + for file_based in ("sqlite", "duckdb"): + with pytest.raises(ValueError, match="does not support"): + _get(file_based).build_uri(database="/tmp/x.db", auth_method="oauth") + + def test_limit_query_is_plain_limit_and_parses_as_snowflake(): limited = dialect.limit_query("SELECT a, b FROM some_table", 5) assert limited == "SELECT a, b FROM some_table LIMIT 5" diff --git a/shared/tests/test_snowflake_oauth.py b/shared/tests/test_snowflake_oauth.py new file mode 100644 index 000000000..1b9d04a9f --- /dev/null +++ b/shared/tests/test_snowflake_oauth.py @@ -0,0 +1,125 @@ +import base64 +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.parse import parse_qs + +import pytest + +from shared.snowflake_oauth import ( + SnowflakeOAuthError, + derive_snowflake_endpoints, + exchange_authorization_code, + refresh_access_token, +) + + +class _TokenEndpointHandler(BaseHTTPRequestHandler): + """Mock token endpoint. The test controls behavior via server.responses (a list of + (status, payload) tuples popped per request) and records requests in server.requests.""" + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + body = parse_qs(self.rfile.read(length).decode()) + self.server.requests.append( + { + "path": self.path, + "body": {k: v[0] for k, v in body.items()}, + "authorization": self.headers.get("Authorization"), + } + ) + status, payload = self.server.responses.pop(0) + raw = payload if isinstance(payload, bytes) else json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(raw))) + self.end_headers() + self.wfile.write(raw) + + def log_message(self, *args): + pass + + +@pytest.fixture() +def token_server(): + server = HTTPServer(("127.0.0.1", 0), _TokenEndpointHandler) + server.requests = [] + server.responses = [] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + thread.join(timeout=5) + + +def _endpoint(server) -> str: + return f"http://127.0.0.1:{server.server_address[1]}/oauth/token-request" + + +def test_derive_snowflake_endpoints(): + authorize, token = derive_snowflake_endpoints("myorg-myaccount") + assert authorize == "https://myorg-myaccount.snowflakecomputing.com/oauth/authorize" + assert token == "https://myorg-myaccount.snowflakecomputing.com/oauth/token-request" + + +def test_refresh_access_token_success(token_server): + token_server.responses.append((200, {"access_token": "at-1", "expires_in": "600", "refresh_token": "rt-2"})) + result = refresh_access_token(_endpoint(token_server), "client", "secret", "rt-1") + assert result.access_token == "at-1" + assert result.expires_in == 600 + assert result.refresh_token == "rt-2" + request = token_server.requests[0] + assert request["body"] == {"grant_type": "refresh_token", "refresh_token": "rt-1"} + expected = base64.b64encode(b"client:secret").decode() + assert request["authorization"] == f"Basic {expected}" + + +def test_refresh_without_rotation_returns_none(token_server): + token_server.responses.append((200, {"access_token": "at-1", "expires_in": 599})) + result = refresh_access_token(_endpoint(token_server), "client", "secret", "rt-1") + assert result.refresh_token is None + + +def test_exchange_authorization_code(token_server): + token_server.responses.append((200, {"access_token": "at-1", "refresh_token": "rt-1"})) + result = exchange_authorization_code( + _endpoint(token_server), "client", "secret", "the-code", "http://localhost:63578/cb" + ) + assert result.access_token == "at-1" + assert result.refresh_token == "rt-1" + assert result.expires_in is None + assert token_server.requests[0]["body"] == { + "grant_type": "authorization_code", + "code": "the-code", + "redirect_uri": "http://localhost:63578/cb", + } + + +def test_invalid_grant_flags_reauthentication(token_server): + token_server.responses.append((400, {"error": "invalid_grant", "error_description": "expired"})) + with pytest.raises(SnowflakeOAuthError) as exc_info: + refresh_access_token(_endpoint(token_server), "client", "secret", "rt-old") + assert exc_info.value.error == "invalid_grant" + assert exc_info.value.requires_reauthentication + assert "expired" in str(exc_info.value) + + +def test_server_error_is_not_reauthentication(token_server): + token_server.responses.append((503, b"upstream down")) + with pytest.raises(SnowflakeOAuthError) as exc_info: + refresh_access_token(_endpoint(token_server), "client", "secret", "rt-1") + assert exc_info.value.status_code == 503 + assert not exc_info.value.requires_reauthentication + + +def test_missing_access_token_raises(token_server): + token_server.responses.append((200, {"token_type": "bearer"})) + with pytest.raises(SnowflakeOAuthError, match="missing access_token"): + refresh_access_token(_endpoint(token_server), "client", "secret", "rt-1") + + +def test_unreachable_endpoint_raises(): + with pytest.raises(SnowflakeOAuthError, match="Could not reach"): + refresh_access_token("http://127.0.0.1:1/token", "client", "secret", "rt-1")