From e12d6dd9654dc1ef8b65a0fe09ec9ddf9557732f Mon Sep 17 00:00:00 2001 From: skadel Date: Tue, 28 Jul 2026 00:06:55 +0200 Subject: [PATCH 1/4] fix: harden optional connectors and CLI errors --- README.md | 24 +++++--- back/build_query/query_chain.py | 4 +- back/cli/generate.py | 48 ++++++++++++++- back/tests/test_cli_optional_connectors.py | 69 ++++++++++++++++++++++ back/tests/test_fix_duck_db_sql.py | 3 +- back/tests/test_handle_other_threading.py | 2 +- back/tests/test_llm_errors.py | 41 +++++++++++++ back/utils/examples.py | 68 +++++++++++++++++++-- back/utils/llm_errors.py | 48 +++++++++++---- back/utils/optional_deps.py | 13 ++++ back/utils/snowflake_connector.py | 11 +--- docs/quickstart-dbt.md | 60 +++++++++++++------ docs/quickstart.md | 68 +++++++++++++++------ 13 files changed, 378 insertions(+), 81 deletions(-) create mode 100644 back/tests/test_cli_optional_connectors.py diff --git a/README.md b/README.md index af25133..5cb8f61 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,8 @@ warehouse. ```bash pip install mocksql # CLI and local DuckDB execution pip install mocksql[bigquery] # BigQuery schema import/profiling +pip install mocksql[snowflake] # Snowflake schema import/validation +pip install mocksql[all] # BigQuery + Snowflake + Trino mocksql init ``` @@ -29,9 +31,10 @@ GOOGLE_CLOUD_LOCATION=us-central1 OPENAI_API_KEY=sk-... ``` -For a BigQuery source also set `BQ_TEST_PROJECT` (or rely on its -`VERTEX_PROJECT` fallback) and authenticate with Application Default Credentials -or `GOOGLE_APPLICATION_CREDENTIALS`. +For a BigQuery source, set `BQ_TEST_PROJECT` explicitly and authenticate with +Application Default Credentials or `GOOGLE_APPLICATION_CREDENTIALS`. A +`VERTEX_PROJECT` fallback exists for compatibility, but should not be used when +cost isolation matters. ```bash mocksql generate models/orders.sql @@ -48,7 +51,7 @@ and BigQuery Sandbox/billing details. | BigQuery | Supported | Imports missing schemas with `mocksql[bigquery]`; `--profile` issues real BigQuery queries. | | DuckDB | Cache-only | Local test execution works; prepare `schema_cache` before generation. | | PostgreSQL | Cache-only | Validation is available, but this generation flow does not import Postgres schemas. | -| Snowflake | Supported with an explicit schema refresh | Validation/transpilation work; run `mocksql refresh-schemas --table database.schema.table` before generation. | +| Snowflake | Supported | Imports cache misses automatically with `mocksql[snowflake]`; `refresh-schemas` preloads or refreshes the cache. | | Trino | Partial | Validation and `refresh-schemas` support exist; generation still requires cached schemas. | ## dbt status @@ -58,8 +61,8 @@ from `target/compiled/`. It never treats the dbt manifest as a schema source. - dbt-BigQuery: supported, including BigQuery schema import. - dbt-DuckDB: supported when `schema_cache` has been prepared. -- dbt-Snowflake: supported after explicitly refreshing the referenced schemas - into `schema_cache`; compiled-SQL resolution and validation work. +- dbt-Snowflake: supported; cache misses in compiled SQL are imported + automatically, and `refresh-schemas` is available for an explicit preload. Full setup: [docs/quickstart-dbt.md](docs/quickstart-dbt.md). @@ -80,6 +83,9 @@ With `dialect: snowflake` and `mocksql[snowflake]`, both `mocksql generate` and `mocksql refresh-schemas` read schemas from Snowflake `INFORMATION_SCHEMA`. They require `SNOWFLAKE_ACCOUNT`, `SNOWFLAKE_USER`, `SNOWFLAKE_PASSWORD`, `SNOWFLAKE_WAREHOUSE`, and `SNOWFLAKE_DATABASE`; they never require or call -BigQuery. Snowflake profiling is not available yet: `generate --profile` reports -that limitation and continues without profiling rather than falling back to -BigQuery. +BigQuery. `SNOWFLAKE_DATABASE` is currently required by the CLI connection even +when every SQL relation is fully qualified. `generate` imports cache misses +automatically; use `refresh-schemas --table DATABASE.SCHEMA.TABLE` to preload or +force-refresh schemas. Snowflake profiling is not available yet: +`generate --profile` reports that limitation and continues without profiling +rather than falling back to BigQuery. diff --git a/back/build_query/query_chain.py b/back/build_query/query_chain.py index c29f44e..ad64b65 100644 --- a/back/build_query/query_chain.py +++ b/back/build_query/query_chain.py @@ -393,9 +393,7 @@ async def _handle_other(state: QueryState): result = await chain.ainvoke({"descriptions": descriptions}) except Exception as exc: if is_vertex_permission_error(exc): - # Keep this single-argument call compatible with integrations that - # replace the formatter (including existing API/UI tests). - error_msg = format_vertex_permission_message(get_llm_model()) + error_msg = format_vertex_permission_message(get_llm_model(), exc) return { "messages": [ AIMessage( diff --git a/back/cli/generate.py b/back/cli/generate.py index 47990b6..960b001 100644 --- a/back/cli/generate.py +++ b/back/cli/generate.py @@ -42,6 +42,45 @@ def cache_miss_message(dialect: str) -> str: return "[ERROR] No schema importer is available for this dialect." +def require_source_connector(dialect: str) -> None: + """Fail cleanly before a cache-miss path touches an optional connector.""" + import typer + + try: + if dialect == "bigquery": + from utils.optional_deps import import_bigquery + + import_bigquery() + elif dialect == "snowflake": + from utils.optional_deps import import_snowflake + + import_snowflake() + except ImportError as exc: + typer.echo(f"[ERROR] {exc}", err=True) + raise typer.Exit(1) from None + + +def cli_state_error_message(final_state: dict) -> str: + """Return the safe user-facing graph error instead of an internal error code.""" + error = str(final_state.get("error") or "") + if error != "llm_permission_denied": + return error + + from utils.llm_errors import normalize_llm_content + from utils.msg_types import MsgType + + for message in reversed(final_state.get("messages") or []): + if getattr(message, "additional_kwargs", {}).get("type") == MsgType.ERROR: + content = normalize_llm_content(getattr(message, "content", "")) + if content.strip(): + return content.strip() + + from storage.config import get_llm_model + from utils.llm_errors import format_vertex_permission_message + + return format_vertex_permission_message(get_llm_model()) + + # ── Config ──────────────────────────────────────────────────────────────────── @@ -783,6 +822,7 @@ async def run_generate( from build_query.schema_fetcher import fetch_tables_schema_snowflake from models.env_variables import validate_snowflake_env + require_source_connector(dialect) try: validate_snowflake_env() except RuntimeError as exc: @@ -791,6 +831,7 @@ async def run_generate( schema_rows, failed = await fetch_tables_schema_snowflake(missing) partitions = {} elif dialect == "bigquery": + require_source_connector(dialect) if not billing_project: typer.echo( "[ERROR] BQ_TEST_PROJECT not set. Cannot fetch schemas from BigQuery. " @@ -1057,8 +1098,11 @@ def _inject_existing(input_text: str) -> None: final_state = await graph.ainvoke(state, config={"recursion_limit": 50}) if final_state.get("error"): - err = final_state["error"] - typer.echo(f"[ERROR] {err[:500]}{'…' if len(err) > 500 else ''}") + err = cli_state_error_message(final_state) + if final_state["error"] == "llm_permission_denied": + typer.echo(f"[ERROR] {err}") + else: + typer.echo(f"[ERROR] {err[:500]}{'…' if len(err) > 500 else ''}") raise typer.Exit(1) # Step 6 — write outputs diff --git a/back/tests/test_cli_optional_connectors.py b/back/tests/test_cli_optional_connectors.py new file mode 100644 index 0000000..dc8a896 --- /dev/null +++ b/back/tests/test_cli_optional_connectors.py @@ -0,0 +1,69 @@ +"""Optional source connectors must fail early with an actionable CLI error.""" + +from pathlib import Path + +import pytest +import typer + + +async def _noop(*_args, **_kwargs) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dialect", "sql", "extra"), + [ + ( + "bigquery", + "SELECT order_id FROM `demo-project.analytics.orders`", + "mocksql[bigquery]", + ), + ( + "snowflake", + "SELECT ORDER_ID FROM ANALYTICS.PUBLIC.ORDERS", + "mocksql[snowflake]", + ), + ], +) +async def test_generate_cache_miss_reports_missing_connector_without_traceback( + dialect: str, + sql: str, + extra: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], +) -> None: + from cli import generate + + config = tmp_path / "mocksql.yml" + config.write_text(f"dialect: {dialect}\nmodels_path: ./models\n", encoding="utf-8") + model = tmp_path / "models" / "orders.sql" + model.parent.mkdir() + model.write_text(sql, encoding="utf-8") + + def missing_connector(): + raise ImportError(f"Connecteur absent. Installez l'extra : pip install {extra}") + + monkeypatch.setattr("models.env_variables.validate_required_env", lambda: None) + monkeypatch.setattr("models.database.db_pool.init_pool", _noop) + monkeypatch.setattr("init.init_db.run_migrations", _noop) + monkeypatch.setattr(f"utils.optional_deps.import_{dialect}", missing_connector) + monkeypatch.setattr( + generate, + "fetch_tables_schema", + lambda *_args, **_kwargs: pytest.fail( + "BigQuery must not be called after connector preflight fails" + ), + ) + + with pytest.raises(typer.Exit) as exc: + await generate.run_generate(model, config, tmp_path / ".mocksql" / "tests") + + output = capfd.readouterr() + combined = output.out + output.err + assert exc.value.exit_code == 1 + assert f"pip install {extra}" in combined + assert "Traceback" not in combined + other_extra = "mocksql[snowflake]" if dialect == "bigquery" else "mocksql[bigquery]" + assert other_extra not in combined diff --git a/back/tests/test_fix_duck_db_sql.py b/back/tests/test_fix_duck_db_sql.py index 08a2e85..19a18ac 100644 --- a/back/tests/test_fix_duck_db_sql.py +++ b/back/tests/test_fix_duck_db_sql.py @@ -12,6 +12,7 @@ """ import re +from datetime import date import duckdb import pytest @@ -140,7 +141,7 @@ def test_safe_parse_date_end_to_end(self, con): assert "TRY_STRPTIME" in fixed assert "SAFE.CAST" not in fixed - duckdb_ok(con, fixed) + assert con.execute(fixed).fetchone()[0] == date(2024, 1, 15) def test_parse_datetime_strict_raises_on_malformed_value(self, con): """PARSE_DATETIME strict conserve l'erreur que BigQuery lèverait.""" diff --git a/back/tests/test_handle_other_threading.py b/back/tests/test_handle_other_threading.py index 26f1339..a727306 100644 --- a/back/tests/test_handle_other_threading.py +++ b/back/tests/test_handle_other_threading.py @@ -91,7 +91,7 @@ async def test_handle_other_error_also_threads_under_user_message(monkeypatch): ) monkeypatch.setattr( "utils.llm_errors.format_vertex_permission_message", - lambda _model: "Accès refusé.", + lambda _model, _exc: "Accès refusé.", ) state = { diff --git a/back/tests/test_llm_errors.py b/back/tests/test_llm_errors.py index 64c0fa6..94907d2 100644 --- a/back/tests/test_llm_errors.py +++ b/back/tests/test_llm_errors.py @@ -1,7 +1,9 @@ """Vertex errors must be actionable without exposing credential contents.""" import pytest +from langchain_core.messages import AIMessage +from cli.generate import cli_state_error_message from utils.llm_errors import ( classify_vertex_access_error, format_vertex_permission_message, @@ -47,3 +49,42 @@ def test_vertex_access_errors_are_classified_and_actionable( def test_non_vertex_error_is_not_misclassified() -> None: assert classify_vertex_access_error(RuntimeError("network timeout")) is None assert not is_vertex_permission_error(RuntimeError("network timeout")) + + +def test_vertex_message_covers_all_recovery_paths_without_leaking_exception() -> None: + secret_marker = "token=super-sensitive-value" + message = format_vertex_permission_message( + "gemini-2.5-flash", + RuntimeError(f"PERMISSION_DENIED: {secret_marker}"), + ) + + assert "Appel Vertex AI refusé" in message + assert "VERTEX_PROJECT" in message + assert "gcloud auth application-default login" in message + assert "GOOGLE_APPLICATION_CREDENTIALS" in message + assert "aiplatform.googleapis.com" in message + assert "roles/aiplatform.user" in message + assert "llm.provider: openai" in message + assert "OPENAI_API_KEY" in message + assert secret_marker not in message + + +def test_cli_displays_vertex_message_instead_of_internal_error_code() -> None: + message = format_vertex_permission_message( + "gemini-2.5-flash", + RuntimeError("PERMISSION_DENIED: missing roles/aiplatform.user"), + ) + final_state = { + "error": "llm_permission_denied", + "messages": [ + AIMessage( + content=message, + additional_kwargs={"type": "error"}, + ) + ], + } + + displayed = cli_state_error_message(final_state) + + assert displayed == message + assert displayed != "llm_permission_denied" diff --git a/back/utils/examples.py b/back/utils/examples.py index f24fd86..fa86098 100644 --- a/back/utils/examples.py +++ b/back/utils/examples.py @@ -742,6 +742,67 @@ def execute_queries(queries: list[str], con: duckdb.DuckDBPyConnection): raise errors[0] +def _find_matching_parenthesis(sql: str, opening_index: int) -> int | None: + """Return the matching ``)`` while ignoring parentheses inside SQL strings.""" + depth = 0 + quote: str | None = None + index = opening_index + while index < len(sql): + char = sql[index] + if quote: + if char == quote: + if index + 1 < len(sql) and sql[index + 1] == quote: + index += 2 + continue + quote = None + elif char in {"'", '"', "`"}: + quote = char + elif char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + return index + index += 1 + return None + + +def _replace_safe_strptime_casts(sql: str) -> str: + """Translate sqlglot's nested ``SAFE.CAST(STRPTIME(...))`` safely. + + sqlglot 30.14 can emit concatenated arguments containing nested expressions, + which the previous regex could not match. A balanced-parenthesis scan keeps + the STRPTIME arguments intact and preserves the requested result type. + """ + pattern = re.compile(r"SAFE\.CAST\s*\(\s*STRPTIME\s*\(", re.IGNORECASE) + cursor = 0 + while match := pattern.search(sql, cursor): + outer_open = sql.find("(", match.start(), match.end()) + inner_open = match.end() - 1 + inner_close = _find_matching_parenthesis(sql, inner_open) + outer_close = ( + _find_matching_parenthesis(sql, outer_open) if outer_open >= 0 else None + ) + if inner_close is None or outer_close is None or inner_close >= outer_close: + cursor = match.end() + continue + + cast_suffix = sql[inner_close + 1 : outer_close] + cast_match = re.fullmatch( + r"\s+AS\s+(.+?)\s*", cast_suffix, flags=re.IGNORECASE | re.DOTALL + ) + if not cast_match: + cursor = match.end() + continue + + arguments = sql[inner_open + 1 : inner_close] + cast_type = cast_match.group(1) + replacement = f"TRY_CAST(TRY_STRPTIME({arguments}) AS {cast_type})" + sql = sql[: match.start()] + replacement + sql[outer_close + 1 :] + cursor = match.start() + len(replacement) + return sql + + def fix_duck_db_sql(duckdb_sql: str, source_dialect: str = "bigquery") -> str: """ Applique des corrections et des traductions sémantiques pour les requêtes DuckDB @@ -838,12 +899,7 @@ def _fix_date_trunc_week(match): # sqlglot 30+ : SAFE.CAST(STRPTIME(col, '%fmt') AS DATE) # ou SAFE.STRPTIME(col, '%fmt'). DuckDB attend une variante tolérante. - s = re.sub( - r"SAFE\.CAST\s*\(\s*STRPTIME\s*\(\s*([^,]+?)\s*,\s*'([^']+)'\s*\)\s*AS\s+\w+\s*\)", - r"TRY_STRPTIME(\1, '\2')", - s, - flags=re.IGNORECASE, - ) + s = _replace_safe_strptime_casts(s) s = re.sub( r"SAFE\.PARSE_DATE\s*\(\s*'([^']+)'\s*,\s*([^)]+)\)", diff --git a/back/utils/llm_errors.py b/back/utils/llm_errors.py index 83453df..0c98254 100644 --- a/back/utils/llm_errors.py +++ b/back/utils/llm_errors.py @@ -46,14 +46,26 @@ def is_vertex_permission_error(exc: Exception) -> bool: def classify_vertex_access_error(exc: Exception) -> str | None: """Return a stable, credential-safe category for common Vertex failures.""" error = str(exc).upper() - if "DEFAULTCREDENTIALSERROR" in error or "APPLICATION DEFAULT CREDENTIALS" in error: + if ( + "DEFAULTCREDENTIALSERROR" in error + or "APPLICATION DEFAULT CREDENTIALS" in error + or "COULD NOT AUTOMATICALLY DETERMINE CREDENTIALS" in error + ): return "adc_missing" + if "VERTEX_PROJECT" in error or "PROJECT ID" in error and "MISSING" in error: + return "project_missing" if ( "VERTEX AI API HAS NOT BEEN USED" in error + or "SERVICE_DISABLED" in error or "AIPLATFORM.GOOGLEAPIS.COM" in error ): return "api_disabled" - if "PERMISSION_DENIED" in error or "BILLING_DISABLED" in error: + if ( + "PERMISSION_DENIED" in error + or "BILLING_DISABLED" in error + or "FORBIDDEN" in error + or "403" in error + ): if "AIPLATFORM.USER" in error: return "iam_role_missing" if ( @@ -70,16 +82,28 @@ def format_vertex_permission_message( model_name: str, exc: Exception | None = None ) -> str: category = classify_vertex_access_error(exc) if exc else None - guidance = { - "adc_missing": "Configurez les Application Default Credentials avec `gcloud auth application-default login` ou GOOGLE_APPLICATION_CREDENTIALS.", - "api_disabled": "Activez l'API Vertex AI (`aiplatform.googleapis.com`) pour le projet Vertex.", - "iam_role_missing": "Accordez le rôle `roles/aiplatform.user` au compte qui exécute MockSQL.", - "model_access_denied": "Vérifiez que Gemini et ce modèle sont autorisés pour le projet, la région et l'organisation.", + probable_causes = { + "adc_missing": "credentials ADC absents/expirés ou service account mal configuré", + "project_missing": "`VERTEX_PROJECT` absent ou incorrect", + "api_disabled": "API Vertex AI désactivée sur le projet", + "iam_role_missing": "rôle IAM `roles/aiplatform.user` manquant", + "model_access_denied": "modèle Gemini indisponible pour ce projet, cette région ou cette organisation", + "permission_denied": "credentials, projet, API ou rôle IAM incorrects", } - if category in guidance: - return f"Erreur d'accès Vertex AI ({category}) pour « {model_name} ».\n• {guidance[category]}" + probable = probable_causes.get( + category, "credentials, projet, API ou rôle IAM incorrects" + ) return ( - f"Erreur d'accès au modèle LLM (PERMISSION_DENIED).\n" - f"• Vérifiez que le modèle « {model_name} » est accessible dans votre organisation.\n" - f"• Vérifiez vos permissions IAM sur Vertex AI (rôle minimum requis : AI Platform Developer)." + f"Appel Vertex AI refusé pour « {model_name} »" + f"{f' ({category})' if category else ''}.\n" + f"Cause probable : {probable}.\n" + "Actions :\n" + "• Vérifiez `VERTEX_PROJECT` et la région `GOOGLE_CLOUD_LOCATION`.\n" + "• Configurez les Application Default Credentials (ADC) avec " + "`gcloud auth application-default login`, ou un service account via " + "`GOOGLE_APPLICATION_CREDENTIALS`.\n" + "• Activez l'API Vertex AI : `gcloud services enable aiplatform.googleapis.com " + '--project="$VERTEX_PROJECT"`.\n' + "• Accordez `roles/aiplatform.user` au compte qui exécute MockSQL.\n" + "Alternative : définissez `llm.provider: openai` et `OPENAI_API_KEY`." ) diff --git a/back/utils/optional_deps.py b/back/utils/optional_deps.py index bee9a62..6b97fe4 100644 --- a/back/utils/optional_deps.py +++ b/back/utils/optional_deps.py @@ -23,6 +23,19 @@ def import_bigquery(): ) from e +def import_snowflake(): + """Retourne le module `snowflake.connector` ou lève un message clair.""" + try: + import snowflake.connector + + return snowflake.connector + except ImportError as e: + raise ImportError( + "Le connecteur Snowflake n'est pas installé. " + "Installez l'extra correspondant : pip install mocksql[snowflake]" + ) from e + + def import_trino(): """Retourne le module `trino` ou lève un message clair.""" try: diff --git a/back/utils/snowflake_connector.py b/back/utils/snowflake_connector.py index 33ba832..b6980a1 100644 --- a/back/utils/snowflake_connector.py +++ b/back/utils/snowflake_connector.py @@ -15,20 +15,13 @@ SNOWFLAKE_USER, SNOWFLAKE_WAREHOUSE, ) +from utils.optional_deps import import_snowflake _sf_conn: snowflake.connector.SnowflakeConnection | None = None def _import_snowflake(): - try: - import snowflake.connector - - return snowflake.connector - except ImportError as e: - raise ImportError( - "Le connecteur Snowflake n'est pas installé. " - "Installez l'extra correspondant : pip install mocksql[snowflake]" - ) from e + return import_snowflake() def get_sf_connection() -> snowflake.connector.SnowflakeConnection: diff --git a/docs/quickstart-dbt.md b/docs/quickstart-dbt.md index 2c7dcc6..1cd66eb 100644 --- a/docs/quickstart-dbt.md +++ b/docs/quickstart-dbt.md @@ -3,7 +3,7 @@ MockSQL reads a dbt model's **compiled SQL**. It does not compile dbt itself and does not derive schemas from `manifest.json`: the manifest identifies the model; `target/compiled/` supplies the rendered SQL; schemas come from MockSQL's schema -cache or, for BigQuery, from BigQuery. +cache or the automatic BigQuery/Snowflake cache-miss importer. ## Support matrix @@ -11,7 +11,7 @@ cache or, for BigQuery, from BigQuery. |---|---|---| | dbt-BigQuery | Supported | Automatic BigQuery import for cache misses, or `schema_cache` | | dbt-DuckDB | Supported with a prepared cache | `schema_cache` only; no DuckDB schema-import command exists yet | -| dbt-Snowflake | Supported with an explicit schema refresh | Refresh into `schema_cache`, then generate | +| dbt-Snowflake | Supported | Automatic Snowflake import for cache misses, or `schema_cache` | All generated cases are executed locally in DuckDB. The warehouse is never used to execute the synthetic test data. @@ -65,9 +65,10 @@ pip install mocksql[bigquery] mocksql generate models/marts/sales.sql --config mocksql.yml ``` -Set a BigQuery job project (`BQ_TEST_PROJECT`, or `VERTEX_PROJECT` as its -fallback) and Google application credentials. On a cache miss, MockSQL fetches -the referenced table schema and saves it in `.mocksql/schema_cache.json`. +Set an explicit BigQuery job project (`BQ_TEST_PROJECT`) and Google application +credentials. A `VERTEX_PROJECT` fallback exists, but do not use it when cost +isolation matters. On a cache miss, MockSQL fetches the referenced table schema +and saves it in `.mocksql/schema_cache.json`. ### dbt-DuckDB @@ -78,25 +79,46 @@ in the CLI generation path. ### dbt-Snowflake -Install `mocksql[snowflake]`, configure the Snowflake connection variables, and -refresh each required relation into the cache before generation: +Install `mocksql[snowflake]` and configure: + +```dotenv +SNOWFLAKE_ACCOUNT=org-account +SNOWFLAKE_USER=mocksql +SNOWFLAKE_PASSWORD=... +SNOWFLAKE_WAREHOUSE=COMPUTE_WH +SNOWFLAKE_DATABASE=ANALYTICS +# SNOWFLAKE_SCHEMA=PUBLIC # optional +# SNOWFLAKE_ROLE=ANALYST # optional +``` + +`SNOWFLAKE_DATABASE` is required by the current CLI connection validation even +when compiled SQL uses fully qualified relations. Then generate directly: ```bash -mocksql refresh-schemas --table DATABASE.SCHEMA.PARENT_MODEL mocksql generate models/marts/sales.sql --config mocksql.yml ``` -The dbt connector still supplies only compiled SQL; `refresh-schemas` is the -schema source. This manual step is required because `generate` auto-imports -cache misses only for BigQuery. +On a cache miss, `generate` imports the compiled SQL's Snowflake relations +automatically. `refresh-schemas` remains recommended when CI should preload the +cache or when a warehouse schema changed: + +```bash +mocksql refresh-schemas --table DATABASE.SCHEMA.PARENT_MODEL +``` ## BigQuery Sandbox and billing -BigQuery dry-runs validate SQL and estimate bytes processed; they do not scan -table data. Schema metadata reads are also distinct from profiling. The sandbox -can therefore be enough to compile/dry-run, read metadata, and run queries -within the Sandbox free-tier quotas and feature limits. MockSQL profiling is a -real query over the source tables: it consumes that quota and needs a -billing-enabled BigQuery project once those limits or required capabilities are -exceeded. Set `BQ_TEST_PROJECT` explicitly for any BigQuery job; dry-runs are -estimates and are not a guarantee that later profiling is free. +BigQuery dry-runs validate SQL and estimate bytes without executing or charging +for the query. Plain table metadata reads do not scan source rows; however, +MockSQL can issue a real `INFORMATION_SCHEMA.PARTITIONS` query for partition +discovery, and BigQuery applies a 10 MB minimum on-demand processing amount to +each `INFORMATION_SCHEMA` query. `--profile` runs real source-table queries. + +Use an explicit isolated `BQ_TEST_PROJECT`; do not rely on `VERTEX_PROJECT` when +cost isolation matters. A Sandbox project has no billing account and currently +includes 10 GiB active storage plus 1 TiB query processing per month, subject to +Sandbox limits. Omit `--profile` and keep the schema cache warm when the goal is +zero source-query spend. ADC can be set up with +`gcloud auth application-default login`; use `roles/bigquery.metadataViewer` +(or `roles/bigquery.dataViewer` for profiling) plus +`roles/bigquery.jobUser` on the job project. diff --git a/docs/quickstart.md b/docs/quickstart.md index 23c2e80..4ea4df0 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -10,10 +10,12 @@ they do not run the generated synthetic data. - `pip install mocksql` - An LLM credential: Vertex AI/Gemini or OpenAI - `mocksql[bigquery]` only when MockSQL must import or profile BigQuery tables +- `mocksql[snowflake]` when MockSQL must validate or import Snowflake schemas ```bash pip install mocksql pip install mocksql[bigquery] # optional BigQuery connector +pip install mocksql[snowflake] # optional Snowflake connector mocksql --help ``` @@ -63,17 +65,31 @@ is BigQuery. | `bigquery` | BigQuery dry-run | Imports cache misses from BigQuery with `mocksql[bigquery]` | | `postgres` | Postgres validation | Use a prepared `schema_cache`; no Postgres import in this flow | | `duckdb` | Local DuckDB validation | Use a prepared `schema_cache`; no DuckDB import command in this flow | -| `snowflake` | Snowflake `EXPLAIN` validation | Refresh schemas explicitly, then generate from `schema_cache` | +| `snowflake` | Snowflake `EXPLAIN` validation | Imports cache misses automatically with `mocksql[snowflake]` | | `trino` | Trino validation | Use a prepared cache for generation; `refresh-schemas` has Trino support | `mocksql generate` needs schemas. It reads them from `schema_cache` first, then -automatically imports only BigQuery cache misses. It never guesses column types -from generated rows. `mocksql refresh-schemas` refreshes BigQuery schemas by -by default; Snowflake and Trino each have explicit branches. It is not a DuckDB -schema importer. +automatically imports BigQuery or Snowflake cache misses through the connector +selected by `dialect`. It never guesses column types from generated rows. +`mocksql refresh-schemas` explicitly preloads or refreshes BigQuery, Snowflake, +or Trino schemas. It is not a DuckDB schema importer. -For Snowflake, install `mocksql[snowflake]`, set the Snowflake environment -variables, then populate the cache explicitly, for example: +For Snowflake, install `mocksql[snowflake]` and set every required connection +variable: + +```dotenv +SNOWFLAKE_ACCOUNT=org-account +SNOWFLAKE_USER=mocksql +SNOWFLAKE_PASSWORD=... +SNOWFLAKE_WAREHOUSE=COMPUTE_WH +SNOWFLAKE_DATABASE=ANALYTICS +# SNOWFLAKE_SCHEMA=PUBLIC # optional +# SNOWFLAKE_ROLE=ANALYST # optional +``` + +`SNOWFLAKE_DATABASE` is currently required when the CLI opens the connection, +even if every table in the SQL is fully qualified. A normal generation imports +missing schemas automatically. Preloading is optional but useful in CI: ```bash mocksql refresh-schemas --table DATABASE.SCHEMA.ORDERS @@ -85,24 +101,38 @@ mocksql generate models/orders.sql For BigQuery schema import, configure an execution project and credentials: ```dotenv -BQ_TEST_PROJECT=my-billing-project # falls back to VERTEX_PROJECT +BQ_TEST_PROJECT=my-isolated-project # explicit project for BigQuery jobs # GOOGLE_APPLICATION_CREDENTIALS=/absolute/path/service-account.json ``` Application Default Credentials (`gcloud auth application-default login`) are -also supported. Typical permissions are `roles/bigquery.dataViewer` for metadata -and `roles/bigquery.user` to create BigQuery jobs; Gemini additionally needs -`roles/aiplatform.user`. +also supported. Grant `roles/bigquery.metadataViewer` for schema metadata (or +`roles/bigquery.dataViewer` when profiling must read table data) and +`roles/bigquery.jobUser` on `BQ_TEST_PROJECT` to create dry-run/query jobs. +Gemini additionally needs `roles/aiplatform.user`; OpenAI does not. BigQuery dry-runs validate a query and return an estimated `total_bytes_processed`. -They do not read table data and MockSQL uses them to validate/estimate work. A -BigQuery Sandbox can run dry-runs, read metadata, and run real queries within -its free-tier quotas and feature limits; it does not require a billing account. -`mocksql generate --profile` is a real query over source tables, so it consumes -that quota and needs a billing-enabled project once Sandbox/free-tier limits or -required capabilities are exceeded. Dry-run estimates are not charges, nor a -promise that a later real profiling query is free. `profile_budget_tb` limits -which estimated profiling queries are run; it does not make a query free. +They do not execute the query, use query slots, or incur a charge. Plain table +metadata reads do not scan source data. One nuance: for a day-partitioned table, +MockSQL queries `INFORMATION_SCHEMA.PARTITIONS` to discover representative +partitions. That is a real metadata query job; BigQuery applies a 10 MB minimum +on-demand processing amount to each `INFORMATION_SCHEMA` query. + +A [BigQuery Sandbox](https://cloud.google.com/bigquery/docs/sandbox) has no +billing account and provides the free-tier limits (currently 10 GiB active +storage and 1 TiB of query data processed per month), plus Sandbox feature +restrictions. `mocksql generate --profile` executes real queries over source +tables and consumes that allowance; on a billing-enabled project it can incur +charges after the applicable free tier. Dry-run estimates are not charges, nor +a promise that a later real query is free. `profile_budget_tb` filters profiling +queries by their dry-run estimate, but does not make them free. + +To avoid accidental billing, use an explicit isolated Sandbox project in +`BQ_TEST_PROJECT`, omit `--profile`, and keep `schema_cache` warm. Do not rely on +the `VERTEX_PROJECT` fallback when cost isolation matters. See Google's +[dry-run documentation](https://cloud.google.com/bigquery/docs/running-queries#dry-run), +[INFORMATION_SCHEMA pricing](https://cloud.google.com/bigquery/docs/information-schema-intro#pricing), +and [ADC setup](https://cloud.google.com/docs/authentication/provide-credentials-adc). ## Generate and replay From 3f4f703bf8665f319255d20c72a9acff3d28e5cf Mon Sep 17 00:00:00 2001 From: skadel Date: Tue, 28 Jul 2026 00:24:48 +0200 Subject: [PATCH 2/4] fix: allow fully qualified Snowflake references --- back/build_query/validator.py | 9 ++++ back/cli/generate.py | 2 +- back/cli/main.py | 11 ++-- back/models/env_variables.py | 14 +++-- .../tests/test_cli_snowflake_schema_import.py | 54 +++++++++++++++++++ back/utils/snowflake_connector.py | 5 +- 6 files changed, 83 insertions(+), 12 deletions(-) diff --git a/back/build_query/validator.py b/back/build_query/validator.py index 2057fa7..30c59ee 100644 --- a/back/build_query/validator.py +++ b/back/build_query/validator.py @@ -290,7 +290,16 @@ async def compile_query(sql_code, project, dialect): await dk.close() elif dialect == "snowflake": + from models.env_variables import validate_snowflake_env from utils.snowflake_connector import run_sf_query + from utils.sql_code import extract_real_table_refs + + tables = extract_real_table_refs(sql_code, "snowflake") + refs = [ + ".".join(part for part in (table.catalog, table.db, table.name) if part) + for table in tables + ] + validate_snowflake_env(refs) async with atimed("validate: dry-run Snowflake"): await asyncio.to_thread(run_sf_query, sql_code, True) diff --git a/back/cli/generate.py b/back/cli/generate.py index 960b001..7afe1f6 100644 --- a/back/cli/generate.py +++ b/back/cli/generate.py @@ -824,7 +824,7 @@ async def run_generate( require_source_connector(dialect) try: - validate_snowflake_env() + validate_snowflake_env(missing) except RuntimeError as exc: typer.echo(f"[ERROR] {exc}", err=True) raise typer.Exit(1) diff --git a/back/cli/main.py b/back/cli/main.py index 8c3aee4..4625d71 100644 --- a/back/cli/main.py +++ b/back/cli/main.py @@ -1011,12 +1011,6 @@ async def _run() -> None: from build_query.schema_fetcher import fetch_tables_schema_snowflake from models.env_variables import validate_snowflake_env - try: - validate_snowflake_env() - except RuntimeError as exc: - typer.echo(f"[ERROR] {exc}", err=True) - raise typer.Exit(1) - if tables: refs = list(tables) elif from_tests: @@ -1040,6 +1034,11 @@ async def _run() -> None: "(or schema.table), or run `mocksql generate` first." ) raise typer.Exit() + try: + validate_snowflake_env(refs) + except RuntimeError as exc: + typer.echo(f"[ERROR] {exc}", err=True) + raise typer.Exit(1) typer.echo(f"Re-importing {len(refs)} table(s) from Snowflake...") schema_rows, failed = await fetch_tables_schema_snowflake(refs) partitions = {} diff --git a/back/models/env_variables.py b/back/models/env_variables.py index 9697768..4c880ec 100644 --- a/back/models/env_variables.py +++ b/back/models/env_variables.py @@ -87,15 +87,23 @@ def validate_required_env() -> None: SNOWFLAKE_ROLE = os.getenv("SNOWFLAKE_ROLE", "") -def validate_snowflake_env() -> None: - """Fail early with the exact Snowflake connection settings to provide.""" +def validate_snowflake_env(refs: list[str] | None = None) -> None: + """Fail early with the exact Snowflake connection settings to provide. + + A default database is only needed when at least one relation is not fully + qualified as ``database.schema.table``. Connection-only callers may omit + ``refs`` and validate the credentials shared by every Snowflake request. + """ required = { "SNOWFLAKE_ACCOUNT": "identifiant de compte (ex. ORG-ACCOUNT)", "SNOWFLAKE_USER": "utilisateur Snowflake", "SNOWFLAKE_PASSWORD": "mot de passe ou secret d'authentification", "SNOWFLAKE_WAREHOUSE": "warehouse utilisé pour lire INFORMATION_SCHEMA", - "SNOWFLAKE_DATABASE": "base de données source", } + if refs and any(len(ref.split(".")) < 3 for ref in refs): + required["SNOWFLAKE_DATABASE"] = ( + "base par défaut requise pour les relations non pleinement qualifiées" + ) missing = [ f" • {name} — {description}" for name, description in required.items() diff --git a/back/tests/test_cli_snowflake_schema_import.py b/back/tests/test_cli_snowflake_schema_import.py index fe19be6..b5c132c 100644 --- a/back/tests/test_cli_snowflake_schema_import.py +++ b/back/tests/test_cli_snowflake_schema_import.py @@ -51,6 +51,58 @@ def test_snowflake_configuration_error_lists_missing_settings( assert "pip install mocksql[snowflake]" in str(exc.value) +def test_snowflake_database_is_optional_for_fully_qualified_refs( + monkeypatch: pytest.MonkeyPatch, +): + from models.env_variables import validate_snowflake_env + + _snowflake_env(monkeypatch) + monkeypatch.delenv("SNOWFLAKE_DATABASE") + + validate_snowflake_env(["ANALYTICS.PUBLIC.ORDERS"]) + + +def test_snowflake_database_is_required_for_two_part_refs( + monkeypatch: pytest.MonkeyPatch, +): + from models.env_variables import validate_snowflake_env + + _snowflake_env(monkeypatch) + monkeypatch.delenv("SNOWFLAKE_DATABASE") + + with pytest.raises(RuntimeError, match="SNOWFLAKE_DATABASE"): + validate_snowflake_env(["PUBLIC.ORDERS"]) + + +def test_snowflake_connection_omits_empty_database_and_schema( + monkeypatch: pytest.MonkeyPatch, +): + import utils.snowflake_connector as connector + + _snowflake_env(monkeypatch) + monkeypatch.delenv("SNOWFLAKE_DATABASE") + monkeypatch.setattr(connector, "SNOWFLAKE_DATABASE", "") + monkeypatch.setattr(connector, "_sf_conn", None) + captured = {} + + class Connection: + def is_closed(self): + return False + + class FakeConnector: + @staticmethod + def connect(**kwargs): + captured.update(kwargs) + return Connection() + + monkeypatch.setattr(connector, "_import_snowflake", lambda: FakeConnector) + + connector.get_sf_connection() + + assert "database" not in captured + assert "schema" not in captured + + @pytest.mark.asyncio async def test_generate_fetches_missing_snowflake_schema_without_bigquery( tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -59,6 +111,7 @@ async def test_generate_fetches_missing_snowflake_schema_without_bigquery( from cli import generate _snowflake_env(monkeypatch) + monkeypatch.delenv("SNOWFLAKE_DATABASE") config = tmp_path / "mocksql.yml" config.write_text("dialect: snowflake\nmodels_path: ./models\n", encoding="utf-8") model = tmp_path / "models" / "orders.sql" @@ -108,6 +161,7 @@ def test_refresh_schemas_fetches_snowflake_without_bigquery( from cli.generate import save_schema_cache _snowflake_env(monkeypatch) + monkeypatch.delenv("SNOWFLAKE_DATABASE") config = tmp_path / "mocksql.yml" config.write_text("dialect: snowflake\n", encoding="utf-8") save_schema_cache( diff --git a/back/utils/snowflake_connector.py b/back/utils/snowflake_connector.py index b6980a1..6b94f2d 100644 --- a/back/utils/snowflake_connector.py +++ b/back/utils/snowflake_connector.py @@ -37,9 +37,10 @@ def get_sf_connection() -> snowflake.connector.SnowflakeConnection: "user": SNOWFLAKE_USER, "password": SNOWFLAKE_PASSWORD, "warehouse": SNOWFLAKE_WAREHOUSE, - "database": SNOWFLAKE_DATABASE, - "schema": SNOWFLAKE_SCHEMA_NAME, } + if SNOWFLAKE_DATABASE: + kwargs["database"] = SNOWFLAKE_DATABASE + kwargs["schema"] = SNOWFLAKE_SCHEMA_NAME # Le rôle est requis sur certains comptes (ex. comptes partagés type Spider2 # qui imposent role=PARTICIPANT). Optionnel : omis si non défini. if SNOWFLAKE_ROLE: From e2a852fe131719d782e81d4f72e1a6b5b5cf1d7b Mon Sep 17 00:00:00 2001 From: skadel Date: Tue, 28 Jul 2026 08:25:26 +0200 Subject: [PATCH 3/4] fix: map bare Snowflake arrays to JSON --- back/tests/test_snowflake_compat.py | 7 +++++++ back/utils/examples.py | 15 +++++++++------ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/back/tests/test_snowflake_compat.py b/back/tests/test_snowflake_compat.py index 3edd737..49c6c35 100644 --- a/back/tests/test_snowflake_compat.py +++ b/back/tests/test_snowflake_compat.py @@ -713,6 +713,13 @@ def test_variant_ddl_maps_to_json(): assert _get_ddl_type("topics", cols) == "JSON" +def test_snowflake_bare_array_ddl_maps_to_json(): + """Snowflake ARRAY is semi-structured; bare ARRAY must not render as invalid `[]`.""" + assert _resolve_duck_type("ARRAY") == "JSON" + cols = [{"name": "indicators", "type": "ARRAY", "mode": "NULLABLE"}] + assert _get_ddl_type("indicators", cols) == "JSON" + + def test_variant_column_created_as_json(): con = duckdb.connect(":memory:") create_test_tables( diff --git a/back/utils/examples.py b/back/utils/examples.py index fa86098..5d61c4b 100644 --- a/back/utils/examples.py +++ b/back/utils/examples.py @@ -508,13 +508,16 @@ def _resolve_duck_type(bq_ddl_type: str) -> str: Le type d'entrée est toujours en syntaxe BigQuery (STRING / STRUCT<> / ARRAY<>), donc on parse comme bigquery quel que soit le dialect source. """ - # Snowflake semi-structuré (VARIANT/OBJECT) → JSON DuckDB. Laissé tel quel, `VARIANT` - # est un type opaque : l'accès bracket/`->>` rend NULL en silence et un INSERT de - # string nu passe sans broncher → résultats faux muets (sf_bq444). En JSON, l'accès + # Snowflake semi-structuré (VARIANT/OBJECT/ARRAY) → JSON DuckDB. Laissé tel quel, + # un ARRAY sans type d'élément est rendu `[]` par sqlglot, ce qui produit un DDL + # DuckDB invalide. Les ARRAY BigQuery typés restent exprimés sous forme ARRAY<...>. + # Pour VARIANT, le type opaque ferait aussi que l'accès bracket/`->>` rende NULL + # en silence et qu'un INSERT de string nu passe sans broncher → résultats faux + # muets (sf_bq444). En JSON, l'accès # `->`/`->>` est 0-based (aligné avec la réécriture bracket de _fix_snowflake_idioms) # et un INSERT de string non-JSON échoue tôt (`Conversion Error: Malformed JSON`), # routé vers la boucle bad_data par `_is_duckdb_data_error`. - if bq_ddl_type.strip().upper() in ("VARIANT", "OBJECT"): + if bq_ddl_type.strip().upper() in ("VARIANT", "OBJECT", "ARRAY"): return "JSON" try: dummy = sqlglot.parse_one( @@ -561,11 +564,11 @@ def _get_ddl_type(col_name: str, filtered_columns: list) -> str: # que la colonne DuckDB soit réellement JSON (accès 0-based + fail-fast INSERT). return ( "JSON" - if bq_ddl_type.strip().upper() in ("VARIANT", "OBJECT") + if bq_ddl_type.strip().upper() in ("VARIANT", "OBJECT", "ARRAY") else bq_ddl_type ) base = col["type"].upper() - if base in ("VARIANT", "OBJECT"): + if base in ("VARIANT", "OBJECT", "ARRAY"): base = "JSON" mode = col.get("mode", "NULLABLE").upper() if base in ("RECORD", "STRUCT"): From e6472f0ab3cf8b13c137250c9fac3ad0eaa5712d Mon Sep 17 00:00:00 2001 From: skadel Date: Tue, 28 Jul 2026 08:26:19 +0200 Subject: [PATCH 4/4] test: isolate Snowflake schema import connector --- back/tests/test_cli_snowflake_schema_import.py | 1 + 1 file changed, 1 insertion(+) diff --git a/back/tests/test_cli_snowflake_schema_import.py b/back/tests/test_cli_snowflake_schema_import.py index b5c132c..873cd21 100644 --- a/back/tests/test_cli_snowflake_schema_import.py +++ b/back/tests/test_cli_snowflake_schema_import.py @@ -132,6 +132,7 @@ class StopAfterImport(Exception): pass monkeypatch.setattr("models.env_variables.validate_required_env", lambda: None) + monkeypatch.setattr(generate, "require_source_connector", lambda _dialect: None) monkeypatch.setattr( "build_query.schema_fetcher.fetch_tables_schema_snowflake", fake_sf_fetch )