Add connection-scoped query resolution

sql2json/sql2json.py

@@ -103,6 +103,47 @@ def list_queries(config_path: Optional[str] = None) -> list:
     return list(config.get("queries", {}).keys())
 
 
+def _get_connection_queries_dict(config: dict, connections: dict) -> dict:
+    """Return validated per-connection query mappings from config."""
+    connection_queries = config.get("connection_queries", {})
+
+    if connection_queries is None:
+        return {}
+
+    if not isinstance(connection_queries, dict):
+        raise ValueError("connection_queries must be an object")
+
+    for connection_name, queries in connection_queries.items():
+        if connection_name not in connections:
+            raise ValueError(
+                f"connection_queries references unknown connection '{connection_name}'"
+            )
+
+        if not isinstance(queries, dict):
+            raise ValueError(f"connection_queries.{connection_name} must be an object")
+
+        for query_name, raw_query in queries.items():
+            if not isinstance(raw_query, str):
+                raise ValueError(
+                    f"connection_queries.{connection_name}.{query_name} must be a string"
+                )
+
+    return connection_queries
+
+
+def _resolve_query_string(
+    connection_name: str, query_name: str, connections: dict, config: dict
+) -> str:
+    connection_queries = _get_connection_queries_dict(config, connections)
+    scoped_queries = connection_queries.get(connection_name, {})
+
+    if connection_name in connections and query_name in scoped_queries:
+        return scoped_queries[query_name]
+
+    config_queries = config.get("queries", {})
+    return config_queries.get(query_name, query_name)
+
+
 def run_query_by_name(
     conection_name: str = "default", query_name: str = "default", **kwargs
 ) -> list:
@@ -117,13 +158,13 @@ def run_query_by_name(
     config = load_config_file(config_path)
 
     config_dbs = _get_connections_dict(config)
-    config_queries = config.get("queries", {})
 
     # If conection_name does not exist, try to use as connection string
     conection_string = config_dbs.get(conection_name, conection_name)
 
-    # If query_name does not exist, try to use as inline SQL
-    raw_query_string = config_queries.get(query_name, query_name)
+    raw_query_string = _resolve_query_string(
+        conection_name, query_name, config_dbs, config
+    )
 
     if raw_query_string.startswith("@"):
         raw_query_string = load_query_from_file(raw_query_string[1:])

tests/test_cli.py

@@ -191,6 +191,25 @@ def test_library_api_still_raises(self, tmp_path):
         with pytest.raises(Exception):
             run_query2json("default", "SELECT * FROM nonexistent_table", config=cfg)
 
+    def test_malformed_connection_queries_stderr_is_json(self, tmp_path):
+        cfg = str(tmp_path / "config.json")
+        _write_config(
+            cfg,
+            {
+                "connections": {"default": "sqlite:///:memory:"},
+                "connection_queries": {"default": "SELECT 1"},
+                "queries": {},
+            },
+        )
+
+        result = run_cli("--name", "default", "--query", "anything", "--config", cfg)
+
+        assert result.returncode != 0
+        assert result.stdout == ""
+        error = json.loads(result.stderr)
+        assert error["type"] == "ValueError"
+        assert error["error"] == "connection_queries.default must be an object"
+
 
 class TestTimezone:
     def test_utc_exits_zero(self, tmp_path):

tests/test_sql2json.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 from importlib.metadata import version
 
 import pytest
@@ -90,3 +91,90 @@ def test_run_query_raw_sql_parameters():
     assert x == json_result["a"]
     assert y == json_result["b"]
     assert x + y == json_result["xy"]
+
+
+def _write_config(path, data):
+    with open(path, "w") as f:
+        json.dump(data, f)
+
+
+def test_run_query_by_name_uses_connection_scoped_query(tmp_path):
+    cfg = tmp_path / "config.json"
+    _write_config(
+        cfg,
+        {
+            "connections": {"default": "sqlite:///:memory:"},
+            "connection_queries": {"default": {"answer": "SELECT 246 AS ticket"}},
+            "queries": {},
+        },
+    )
+
+    json_results = run_query_by_name("default", "answer", config=str(cfg))
+
+    assert json_results == [{"ticket": 246}]
+
+
+def test_run_query_by_name_falls_back_to_global_query(tmp_path):
+    cfg = tmp_path / "config.json"
+    _write_config(
+        cfg,
+        {
+            "connections": {"default": "sqlite:///:memory:"},
+            "connection_queries": {"default": {"scoped": "SELECT 1 AS scoped"}},
+            "queries": {"global": "SELECT 2 AS global_value"},
+        },
+    )
+
+    json_results = run_query_by_name("default", "global", config=str(cfg))
+
+    assert json_results == [{"global_value": 2}]
+
+
+def test_run_query_by_name_scoped_query_takes_precedence_over_global(tmp_path):
+    cfg = tmp_path / "config.json"
+    _write_config(
+        cfg,
+        {
+            "connections": {"default": "sqlite:///:memory:"},
+            "connection_queries": {"default": {"same": "SELECT 'scoped' AS source"}},
+            "queries": {"same": "SELECT 'global' AS source"},
+        },
+    )
+
+    json_results = run_query_by_name("default", "same", config=str(cfg))
+
+    assert json_results == [{"source": "scoped"}]
+
+
+def test_run_query_by_name_rejects_malformed_connection_queries(tmp_path):
+    cfg = tmp_path / "config.json"
+    _write_config(
+        cfg,
+        {
+            "connections": {"default": "sqlite:///:memory:"},
+            "connection_queries": {"default": "SELECT 1"},
+            "queries": {},
+        },
+    )
+
+    with pytest.raises(
+        ValueError, match="connection_queries.default must be an object"
+    ):
+        run_query_by_name("default", "anything", config=str(cfg))
+
+
+def test_run_query_by_name_rejects_connection_queries_for_unknown_connection(tmp_path):
+    cfg = tmp_path / "config.json"
+    _write_config(
+        cfg,
+        {
+            "connections": {"default": "sqlite:///:memory:"},
+            "connection_queries": {"missing": {"answer": "SELECT 1"}},
+            "queries": {},
+        },
+    )
+
+    with pytest.raises(
+        ValueError, match="connection_queries references unknown connection 'missing'"
+    ):
+        run_query_by_name("default", "anything", config=str(cfg))