diff --git a/README.md b/README.md index cb06fff..575308f 100644 --- a/README.md +++ b/README.md @@ -141,7 +141,8 @@ else: ### `query(sql, connection=None)` -Execute SQL and return results. +Render SQL with the Bruin CLI (Jinja templates and environment schema prefixes), +then execute the rendered SQL with the existing Python database client and return results. ```python from bruin import query @@ -175,6 +176,54 @@ df_pg = query("SELECT * FROM users", connection="my_postgres") Every query is automatically annotated with `@bruin.config` metadata for observability and cost tracking. +#### CLI rendering requirements + +Both `query()` and `get_connection(...).query()` render before executing. The selected +connection name is passed to the renderer; execution still uses that connection's +existing Python client and injected credentials. Result handling is unchanged: the +driver supplies the DataFrame or `None` (DuckDB also returns DataFrames for DDL/DML). +The CLI only returns SQL; it never executes the query or transports result data. + +```python +df = query("SELECT * FROM analytics.events WHERE dt >= '{{ start_date }}'") +conn = get_connection("my_bigquery") +df = conn.query("SELECT '{{ start_date }}' AS run_start") +``` + +**Pending CLI integration:** the CLI change is being developed separately. The +following provisional contract must be verified against that implementation before +release; these flags and response format are not claimed to exist in a released CLI: + +```text +bruin render --stdin --connection NAME --raw-query --output json +stdin: the query SQL +stdout: {"query": "rendered SQL"} +``` + +The stdin mode must render only the supplied query, including Jinja and schema +prefixes, without asset materialization or asset hooks. The SDK does not create a +temporary SQL asset. The process inherits the current working directory and a copy +of the full environment, including `BRUIN_CONFIG_FILE`, connection credentials, +`BRUIN_THIS`, asset/pipeline names, execution time, and run ID. Only the child's +`BRUIN_CONNECTION` is overridden with the selected connection name; the parent's +environment and Python connection remain unchanged. + +`BRUIN_ENVIRONMENT`, when present, is also passed as `--environment`. +`BRUIN_START_DATETIME` / `BRUIN_END_DATETIME` (falling back to the corresponding +`*_DATE` values) are passed as `--start-date` / `--end-date`. +`BRUIN_FULL_REFRESH=1` adds `--full-refresh`, and typed `context.vars` are passed +as repeated `--var KEY=JSON` arguments. The CLI must honor the inherited run +context in stdin mode. With no run context, CLI defaults apply. + +The SDK invokes `bruin` from `PATH`, or the executable specified by `BRUIN_PATH`, +without a shell and sends SQL through stdin. A missing/unusable binary, nonzero +render exit, invalid/empty JSON query, or 60-second render timeout raises +`QueryError` before touching the database client. No original-SQL fallback is used. +Compatibility is based on the command contract, not a guessed minimum version: +older CLIs without this stdin mode fail with an upgrade message. The CLI release +providing this contract must be identified during the pending integration check. +Direct native client calls such as `conn.client.query(...)` bypass SDK rendering. + --- ### `get_connection(name)` @@ -292,7 +341,8 @@ from bruin.exceptions import ( ConnectionNotFoundError, # Connection name not found or env var missing ConnectionParseError, # Invalid JSON in connection env var ConnectionTypeError, # Unsupported or generic connection type - QueryError, # SQL execution failed + CLIError, # Bruin CLI process failed + QueryError, # SQL rendering or execution failed ) ``` diff --git a/src/bruin/_cli.py b/src/bruin/_cli.py new file mode 100644 index 0000000..7f663bb --- /dev/null +++ b/src/bruin/_cli.py @@ -0,0 +1,45 @@ +"""Shared process boundary for Bruin CLI commands.""" + +import os +import subprocess + +from bruin.exceptions import CLIError + + +def _run_cli( + args: list[str], *, input: str | None = None, env: dict[str, str] | None = None +) -> str: + """Run CLI arguments without a shell and return stdout, or raise CLIError. + + When supplied, env is the complete child environment, as in subprocess.run. + Otherwise the current environment is inherited. Neither is modified. + """ + if env is None: + env = os.environ.copy() + binary = env.get("BRUIN_PATH") or "bruin" + try: + result = subprocess.run( + [binary, *args], + input=input, + text=True, + encoding="utf-8", + capture_output=True, + env=env, + shell=False, + timeout=60, + check=False, + ) + except FileNotFoundError as exc: + raise CLIError( + f"Bruin CLI binary '{binary}' not found. Install Bruin and put it on PATH, " + "or set BRUIN_PATH to its executable." + ) from exc + except subprocess.TimeoutExpired as exc: + raise CLIError("Bruin CLI command timed out after 60s.") from exc + except OSError as exc: + raise CLIError(f"Cannot start Bruin CLI '{binary}': {exc}") from exc + + if result.returncode: + detail = result.stderr.strip() or result.stdout.strip() or "no error output" + raise CLIError(f"Bruin CLI command failed (exit {result.returncode}): {detail}") + return result.stdout diff --git a/src/bruin/_connection.py b/src/bruin/_connection.py index b07e7fd..5c039df 100644 --- a/src/bruin/_connection.py +++ b/src/bruin/_connection.py @@ -41,7 +41,7 @@ def client(self): return self._client def query(self, sql: str) -> "pd.DataFrame | None": # noqa: F821 - """Execute *sql* on this connection. + """Render *sql* with the Bruin CLI, then execute on this connection's client. Returns a pandas DataFrame for data-returning statements, or None for DDL/DML. diff --git a/src/bruin/_query.py b/src/bruin/_query.py index 4895f45..c0781c0 100644 --- a/src/bruin/_query.py +++ b/src/bruin/_query.py @@ -3,6 +3,7 @@ import time from bruin._connection import GCPConnection, get_connection +from bruin._render import _render_sql from bruin.exceptions import ConnectionNotFoundError, ConnectionTypeError, QueryError logger = logging.getLogger("bruin") @@ -63,7 +64,7 @@ def _run_query(conn, sql: str): if conn.type == "generic": raise ConnectionTypeError(f"Cannot run queries against generic connection '{conn.name}'.") - annotated = _annotate_sql(sql) + annotated = _annotate_sql(_render_sql(sql, conn.name)) logger.debug( "Executing query on '%s' (%s, %d chars)", @@ -101,7 +102,8 @@ def query(sql: str, connection: str | None = None) -> "pd.DataFrame | None": # Parameters ---------- sql : str - The SQL statement to execute. + The SQL statement to render with the Bruin CLI (Jinja and schema prefix) + before executing with the Python database client. connection : str, optional Connection name. When *None*, falls back to the asset's default connection (``BRUIN_CONNECTION`` env var). diff --git a/src/bruin/_render.py b/src/bruin/_render.py new file mode 100644 index 0000000..f44e3b4 --- /dev/null +++ b/src/bruin/_render.py @@ -0,0 +1,58 @@ +"""Render query SQL with the Bruin CLI; database execution stays in Python.""" + +import json +import os + +from bruin._cli import _run_cli +from bruin.exceptions import CLIError, QueryError + + +def _render_sql(sql: str, connection: str) -> str: + env = os.environ.copy() + env["BRUIN_CONNECTION"] = connection + # shortcut: provisional stdin contract; verify against the CLI change before release. + command = [ + "render", + "--stdin", + "--connection", + connection, + "--raw-query", + "--output", + "json", + ] + if env.get("BRUIN_ENVIRONMENT"): + command.extend(["--environment", env["BRUIN_ENVIRONMENT"]]) + for flag, keys in ( + ("--start-date", ("BRUIN_START_DATETIME", "BRUIN_START_DATE")), + ("--end-date", ("BRUIN_END_DATETIME", "BRUIN_END_DATE")), + ): + value = next((env[key] for key in keys if env.get(key)), None) + if value: + command.extend([flag, value]) + if env.get("BRUIN_FULL_REFRESH") == "1": + command.append("--full-refresh") + # Use the same typed variables exposed by bruin.context. + from bruin._context import context + + for key, value in context.vars.items(): + command.extend(["--var", f"{key}={json.dumps(value)}"]) + + try: + output = _run_cli(command, input=sql, env=env) + except CLIError as exc: + raise QueryError( + f"SQL rendering failed on connection '{connection}': {exc}\nA CLI supporting " + "'render --stdin --connection ... --raw-query --output json' " + "is required; upgrade the CLI if these options are unsupported." + ) from exc + try: + payload = json.loads(output) + except (ValueError, TypeError) as exc: + raise QueryError( + "Invalid Bruin render output: expected a JSON object containing a 'query' string. " + "Check that the CLI supports stdin query rendering." + ) from exc + rendered = payload.get("query") if isinstance(payload, dict) else None + if not isinstance(rendered, str) or not rendered.strip(): + raise QueryError("Invalid Bruin render output: expected a non-empty 'query' string.") + return rendered diff --git a/src/bruin/exceptions.py b/src/bruin/exceptions.py index 20652d7..5e15b92 100644 --- a/src/bruin/exceptions.py +++ b/src/bruin/exceptions.py @@ -15,4 +15,8 @@ class ConnectionTypeError(BruinError): class QueryError(BruinError): - """Raised when a SQL query fails to execute.""" + """Raised when SQL rendering or database execution fails.""" + + +class CLIError(BruinError): + """Raised when a Bruin CLI command cannot start or fails.""" diff --git a/tests/conftest.py b/tests/conftest.py index 0b84693..7270d10 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,12 @@ import pytest +@pytest.fixture +def render_passthrough(monkeypatch): + """Keep driver tests independent of the CLI; test_render covers the boundary.""" + monkeypatch.setattr("bruin._query._render_sql", lambda sql, connection: sql) + + @pytest.fixture def bq_connection_json(): """Realistic BigQuery connection payload (service_account_json is double-serialized).""" diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 622fe8b..2218989 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -19,6 +19,12 @@ import pytest + +@pytest.fixture(autouse=True) +def _driver_only_rendering(render_passthrough): + """These tests exercise real database drivers, not CLI rendering.""" + + # --------------------------------------------------------------------------- # Auto-apply the integration marker to every test collected in this directory # --------------------------------------------------------------------------- diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..ec19af2 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,56 @@ +import os +import subprocess +from unittest.mock import MagicMock + +import pytest + +from bruin._cli import _run_cli +from bruin.exceptions import CLIError + + +@pytest.mark.parametrize("binary", [None, "", "/path with spaces/bruin"]) +def test_generic_command_inherits_environment(monkeypatch, binary): + monkeypatch.delenv("BRUIN_PATH", raising=False) + if binary is not None: + monkeypatch.setenv("BRUIN_PATH", binary) + original_env = dict(os.environ) + run = MagicMock(return_value=subprocess.CompletedProcess([], 0, "version output", "diagnostic")) + monkeypatch.setattr("bruin._cli.subprocess.run", run) + + assert _run_cli(["version"]) == "version output" + + run.assert_called_once_with( + [binary or "bruin", "version"], + input=None, + text=True, + encoding="utf-8", + capture_output=True, + env=original_env, + shell=False, + timeout=60, + check=False, + ) + assert dict(os.environ) == original_env + + +def test_explicit_environment_is_complete_and_not_modified(monkeypatch): + monkeypatch.setenv("BRUIN_PATH", "/parent/bruin") + env = {"BRUIN_PATH": "/child/bruin", "CUSTOM": "child"} + original_env = env.copy() + run = MagicMock(return_value=subprocess.CompletedProcess([], 0, "output", "")) + monkeypatch.setattr("bruin._cli.subprocess.run", run) + + _run_cli(["validate", "pipeline"], input="input", env=env) + + assert run.call_args.args[0] == ["/child/bruin", "validate", "pipeline"] + assert run.call_args.kwargs["input"] == "input" + assert run.call_args.kwargs["env"] == original_env + assert env == original_env + + +def test_command_failure_raises_cli_error(monkeypatch): + run = MagicMock(return_value=subprocess.CompletedProcess([], 2, "", "invalid pipeline")) + monkeypatch.setattr("bruin._cli.subprocess.run", run) + + with pytest.raises(CLIError, match=r"exit 2.*invalid pipeline"): + _run_cli(["validate", "pipeline"]) diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py index 790aa98..96d4d9d 100644 --- a/tests/test_exceptions.py +++ b/tests/test_exceptions.py @@ -88,7 +88,9 @@ def test_unsupported_type_client_access(self, monkeypatch): class TestQueryError: - def test_wraps_underlying_exception(self, monkeypatch, snowflake_connection_json): + def test_wraps_underlying_exception( + self, monkeypatch, snowflake_connection_json, render_passthrough + ): monkeypatch.setenv( "BRUIN_CONNECTION_TYPES", json.dumps({"sf": "snowflake"}), diff --git a/tests/test_logging.py b/tests/test_logging.py index ae59020..c768b01 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -6,6 +6,8 @@ from bruin import get_connection, query +pytestmark = pytest.mark.usefixtures("render_passthrough") + @pytest.fixture def pg_env(monkeypatch, postgres_connection_json): diff --git a/tests/test_query.py b/tests/test_query.py index dc90190..5be5cc0 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -7,6 +7,8 @@ from bruin import query from bruin.exceptions import ConnectionNotFoundError, ConnectionTypeError, QueryError +pytestmark = pytest.mark.usefixtures("render_passthrough") + @pytest.fixture def _setup_bq(monkeypatch, bq_connection_json): diff --git a/tests/test_render.py b/tests/test_render.py new file mode 100644 index 0000000..94d1a56 --- /dev/null +++ b/tests/test_render.py @@ -0,0 +1,210 @@ +"""SDK boundary tests for the provisional CLI contract, not live CLI validation.""" + +import json +import os +import subprocess +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from bruin import get_connection, query +from bruin._render import _render_sql +from bruin.exceptions import QueryError + + +@pytest.fixture(autouse=True) +def clean_context(monkeypatch): + for key in os.environ: + if key.startswith("BRUIN_"): + monkeypatch.delenv(key) + + +@pytest.fixture +def cli(monkeypatch): + run = MagicMock( + return_value=subprocess.CompletedProcess( + [], 0, stdout=json.dumps({"query": "SELECT 42 AS value"}), stderr="" + ) + ) + monkeypatch.setattr("bruin._cli.subprocess.run", run) + return run + + +def test_render_transfers_environment_and_context(monkeypatch, cli): + context = { + "BRUIN_PATH": "/path with spaces/bruin", + "BRUIN_CONNECTION": "asset_default", + "BRUIN_ENVIRONMENT": "staging", + "BRUIN_CONFIG_FILE": "/project/.bruin.yml", + "BRUIN_CONNECTION_TYPES": '{"selected":"sqlite"}', + "BRUIN_START_DATE": "2026-10-01", + "BRUIN_START_DATETIME": "2026-10-01T12:30:00", + "BRUIN_END_DATE": "2026-10-02", + "BRUIN_FULL_REFRESH": "1", + "BRUIN_VARS": '{"count":"2","region":"EU","enabled":true,"ids":[1,2]}', + "BRUIN_VARS_SCHEMA": '{"count":{"type":"integer"}}', + "BRUIN_THIS": "analytics.output", + "BRUIN_ASSET": "analytics.output", + "BRUIN_PIPELINE": "analytics", + "BRUIN_EXECUTION_DATETIME": "2026-10-02T15:00:00", + "BRUIN_RUN_ID": "original-run", + "selected": '{"path":":memory:"}', + } + for key, value in context.items(): + monkeypatch.setenv(key, value) + original_env = dict(os.environ) + sql = "SELECT '{{ start_date }}', '$(echo unsafe)'" + + assert _render_sql(sql, "selected") == "SELECT 42 AS value" + + cli.assert_called_once_with( + [ + "/path with spaces/bruin", + "render", + "--stdin", + "--connection", + "selected", + "--raw-query", + "--output", + "json", + "--environment", + "staging", + "--start-date", + "2026-10-01T12:30:00", + "--end-date", + "2026-10-02", + "--full-refresh", + "--var", + "count=2", + "--var", + 'region="EU"', + "--var", + "enabled=true", + "--var", + "ids=[1, 2]", + ], + input=sql, + text=True, + encoding="utf-8", + capture_output=True, + env={**original_env, "BRUIN_CONNECTION": "selected"}, + shell=False, + timeout=60, + check=False, + ) + assert dict(os.environ) == original_env + + +def test_render_without_context_uses_cli_defaults(cli): + _render_sql("SELECT 1", "selected") + assert cli.call_args.args[0] == [ + "bruin", + "render", + "--stdin", + "--connection", + "selected", + "--raw-query", + "--output", + "json", + ] + + +@pytest.mark.parametrize("api", ["default", "explicit", "bound"]) +@pytest.mark.parametrize( + "rendered,expected", [("SELECT 42 AS value", 42), ("CREATE TABLE t(x)", None)] +) +def test_rendered_sql_reaches_existing_client(monkeypatch, cli, api, rendered, expected): + monkeypatch.setenv("BRUIN_CONNECTION_TYPES", '{"selected":"sqlite"}') + monkeypatch.setenv("selected", '{"path":":memory:"}') + monkeypatch.setenv("BRUIN_CONNECTION", "selected" if api == "default" else "other") + conn = get_connection("selected") + client = MagicMock() + cursor = client.cursor.return_value + cursor.description = [("value",)] if expected is not None else None + cursor.fetchall.return_value = [(42,)] + conn._client = client + monkeypatch.setattr( + "bruin._query.get_connection", lambda name: conn if name == conn.name else None + ) + cli.return_value.stdout = json.dumps({"query": rendered}) + + if api == "bound": + result = conn.query("{{ unrendered }}") + elif api == "explicit": + result = query("{{ unrendered }}", connection="selected") + else: + result = query("{{ unrendered }}") + + executed = cursor.execute.call_args.args[0] + assert executed.startswith("-- @bruin.config:") + assert executed.split("\n", 1)[1] == rendered + assert conn._client is client + assert cli.call_args.kwargs["env"]["BRUIN_CONNECTION"] == "selected" + if expected is None: + assert result is None + client.commit.assert_called_once() + else: + pd.testing.assert_frame_equal(result, pd.DataFrame({"value": [expected]})) + client.commit.assert_not_called() + + +@pytest.mark.parametrize( + "response,message", + [ + (subprocess.CompletedProcess([], 1, "", "Jinja failed"), "Jinja failed"), + (subprocess.CompletedProcess([], 2, "unknown flag: --stdin", ""), "upgrade the CLI"), + (subprocess.CompletedProcess([], 1, "", ""), "no error output"), + (subprocess.CompletedProcess([], 0, "plain SQL", ""), "expected a JSON object"), + (subprocess.CompletedProcess([], 0, '{"query":null}', ""), "non-empty 'query'"), + (subprocess.CompletedProcess([], 0, '{"query":" "}', ""), "non-empty 'query'"), + (subprocess.CompletedProcess([], 0, '{"query":[]}', ""), "non-empty 'query'"), + (subprocess.CompletedProcess([], 0, "[]", ""), "non-empty 'query'"), + (subprocess.CompletedProcess([], 0, '{"error":"failed"}', ""), "non-empty 'query'"), + (FileNotFoundError(), "not found.*BRUIN_PATH"), + (PermissionError("not executable"), "Cannot start Bruin CLI"), + (subprocess.TimeoutExpired("bruin", 60), "timed out"), + ], +) +def test_render_failure_never_initializes_client(monkeypatch, cli, response, message): + monkeypatch.setenv("BRUIN_CONNECTION_TYPES", '{"selected":"sqlite"}') + monkeypatch.setenv("selected", '{"path":":memory:"}') + if isinstance(response, Exception): + cli.side_effect = response + else: + cli.return_value = response + factory = MagicMock() + monkeypatch.setattr("bruin._connection._create_client", factory) + + with pytest.raises(QueryError, match=message): + query("SELECT original_sql", "selected") + + factory.assert_not_called() + cli.assert_called_once() + + +def test_real_subprocess_stdin_to_python_database(monkeypatch, tmp_path): + # This executable is a contract stub, deliberately not a live Bruin CLI test. + import sys + + binary = tmp_path / "render stub" + binary.write_text( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "assert sys.argv[1:] == ['render', '--stdin', '--connection', 'selected', " + "'--raw-query', '--output', 'json', '--environment', 'staging']\n" + "assert sys.stdin.read() == \"SELECT '{{ vars.n }}' AS value; -- $(exit 99)\"\n" + "assert os.environ['BRUIN_CONNECTION'] == 'selected'\n" + "assert os.environ['BRUIN_RUN_ID'] == 'test-run'\n" + "print(json.dumps({'query': 'SELECT 42 AS value'}))\n" + ) + binary.chmod(0o755) + monkeypatch.setenv("BRUIN_PATH", str(binary)) + monkeypatch.setenv("BRUIN_ENVIRONMENT", "staging") + monkeypatch.setenv("BRUIN_RUN_ID", "test-run") + monkeypatch.setenv("BRUIN_CONNECTION_TYPES", '{"selected":"sqlite"}') + monkeypatch.setenv("selected", '{"path":":memory:"}') + + result = query("SELECT '{{ vars.n }}' AS value; -- $(exit 99)", "selected") + + pd.testing.assert_frame_equal(result, pd.DataFrame({"value": [42]}))