fix: harden optional connectors and Snowflake references

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.

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(

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)

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,14 +822,16 @@ 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()
+                validate_snowflake_env(missing)
             except RuntimeError as exc:
                 typer.echo(f"[ERROR] {exc}", err=True)
                 raise typer.Exit(1)
             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

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 = {}

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()

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

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"
@@ -79,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
     )
@@ -108,6 +162,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(

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."""

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 = {