Skip to content

Commit 6b30815

Browse files
test(e2e): live kernel-backend (use_sea=True) suite
Moves the previously-ad-hoc /tmp/connector_smoke.py into the repo as a real pytest module under tests/e2e/ — same convention as the rest of the e2e suite. Uses the existing session-scoped `connection_details` fixture from the top-level conftest so it shares the credential surface with every other live test. 11 tests cover: - connect() with use_sea=True opens a session. - SELECT 1: rows + description shape (column name + dbapi type slug). - SELECT * FROM range(10000): multi-batch drain. - fetchmany() pacing across the buffer boundary. - fetchall_arrow() returns a pyarrow Table. - All four metadata methods (catalogs / schemas / tables / columns). - session_configuration={'ANSI_MODE': 'false'} round-trips. - Bad SQL surfaces as DatabaseError with `code='SqlError'` and `sql_state='42P01'` attached as exception attributes. Module-level skips: - `databricks_sql_kernel` not importable → whole module skipped via pytest.importorskip (the wheel hasn't been installed). - Live creds missing → fixture-level skip with a pointed message. Run: `pytest tests/e2e/test_kernel_backend.py -v`. All 11 pass against dogfood in ~20s. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
1 parent 2572362 commit 6b30815

1 file changed

Lines changed: 186 additions & 0 deletions

File tree

tests/e2e/test_kernel_backend.py

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
"""E2E tests for ``use_sea=True`` (routes through the Rust kernel
2+
via the PyO3 ``databricks_sql_kernel`` module).
3+
4+
PAT auth only. Anything else surfaces as ``NotSupportedError``
5+
from the auth bridge — covered as a unit test, not exercised here.
6+
7+
Skipped automatically when:
8+
- The standard ``DATABRICKS_SERVER_HOSTNAME`` / ``HTTP_PATH`` /
9+
``TOKEN`` creds aren't set (existing connector convention).
10+
- ``databricks_sql_kernel`` isn't importable (the wheel hasn't
11+
been installed; run ``pip install
12+
'databricks-sql-connector[kernel]'`` or, for local dev,
13+
``cd databricks-sql-kernel/pyo3 && maturin develop --release``
14+
into this venv).
15+
16+
Run from the connector repo root:
17+
18+
set -a && source ~/.databricks/pecotesting-creds && set +a
19+
.venv/bin/pytest tests/e2e/test_kernel_backend.py -v
20+
"""
21+
22+
from __future__ import annotations
23+
24+
import pytest
25+
26+
import databricks.sql as sql
27+
from databricks.sql.exc import DatabaseError
28+
29+
30+
# Skip the whole module unless the kernel wheel is importable.
31+
pytest.importorskip(
32+
"databricks_sql_kernel",
33+
reason="use_sea=True requires the databricks-sql-kernel package",
34+
)
35+
36+
37+
@pytest.fixture(scope="module")
38+
def kernel_conn_params(connection_details):
39+
"""Live-cred check + connection params for use_sea=True.
40+
41+
Skips the module if any cred is missing rather than letting
42+
every test fail with a confusing connect-time error.
43+
"""
44+
host = connection_details.get("host")
45+
http_path = connection_details.get("http_path")
46+
token = connection_details.get("access_token")
47+
if not (host and http_path and token):
48+
pytest.skip(
49+
"DATABRICKS_SERVER_HOSTNAME / DATABRICKS_HTTP_PATH / "
50+
"DATABRICKS_TOKEN not set"
51+
)
52+
return {
53+
"server_hostname": host,
54+
"http_path": http_path,
55+
"access_token": token,
56+
"use_sea": True,
57+
}
58+
59+
60+
@pytest.fixture
61+
def conn(kernel_conn_params):
62+
"""One-shot connection per test (the simple_test pattern the
63+
existing e2e suite uses for cursor-level tests)."""
64+
c = sql.connect(**kernel_conn_params)
65+
try:
66+
yield c
67+
finally:
68+
c.close()
69+
70+
71+
def test_connect_with_use_sea_opens_a_session(conn):
72+
assert conn.open, "connection should report open after connect()"
73+
74+
75+
def test_select_one(conn):
76+
with conn.cursor() as cur:
77+
cur.execute("SELECT 1 AS n")
78+
assert cur.description[0][0] == "n"
79+
# description type slug matches what Thrift produces
80+
assert cur.description[0][1] == "int"
81+
rows = cur.fetchall()
82+
assert len(rows) == 1
83+
assert rows[0][0] == 1
84+
85+
86+
def test_drain_large_range_to_arrow(conn):
87+
"""SELECT * FROM range(10000) drains as a pyarrow Table with
88+
10000 rows. Exercises the CloudFetch / multi-batch path on the
89+
kernel side."""
90+
with conn.cursor() as cur:
91+
cur.execute("SELECT * FROM range(10000)")
92+
rows = cur.fetchall()
93+
assert len(rows) == 10000
94+
95+
96+
def test_fetchmany_pacing(conn):
97+
"""fetchmany honours the requested size and stops cleanly at
98+
end-of-stream — covers the buffer-slicing logic in
99+
KernelResultSet."""
100+
with conn.cursor() as cur:
101+
cur.execute("SELECT * FROM range(50)")
102+
r1 = cur.fetchmany(10)
103+
r2 = cur.fetchmany(20)
104+
r3 = cur.fetchmany(100) # capped at remaining
105+
assert (len(r1), len(r2), len(r3)) == (10, 20, 20)
106+
107+
108+
def test_fetchall_arrow(conn):
109+
with conn.cursor() as cur:
110+
cur.execute("SELECT 1 AS a, 'hi' AS b")
111+
table = cur.fetchall_arrow()
112+
assert table.num_rows == 1
113+
assert table.column_names == ["a", "b"]
114+
115+
116+
# ── Metadata ──────────────────────────────────────────────────────
117+
118+
119+
def test_metadata_catalogs(conn):
120+
with conn.cursor() as cur:
121+
cur.catalogs()
122+
rows = cur.fetchall()
123+
assert len(rows) > 0
124+
125+
126+
def test_metadata_schemas(conn):
127+
with conn.cursor() as cur:
128+
cur.schemas(catalog_name="main")
129+
rows = cur.fetchall()
130+
assert len(rows) > 0
131+
132+
133+
def test_metadata_tables(conn):
134+
with conn.cursor() as cur:
135+
cur.tables(catalog_name="system", schema_name="information_schema")
136+
rows = cur.fetchall()
137+
assert len(rows) > 0
138+
139+
140+
def test_metadata_columns(conn):
141+
with conn.cursor() as cur:
142+
cur.columns(
143+
catalog_name="system",
144+
schema_name="information_schema",
145+
table_name="tables",
146+
)
147+
rows = cur.fetchall()
148+
assert len(rows) > 0
149+
150+
151+
# ── Session configuration ─────────────────────────────────────────
152+
153+
154+
def test_session_configuration_round_trips(kernel_conn_params):
155+
"""`session_configuration` flows through to the kernel's
156+
`session_conf` and is honoured by the server.
157+
158+
`ANSI_MODE` is the safe choice — it's on the SEA allow-list and
159+
isn't workspace-policy-clamped (unlike `STATEMENT_TIMEOUT`) or
160+
rejected by the warehouse (unlike `TIMEZONE` on dogfood)."""
161+
params = dict(kernel_conn_params)
162+
params["session_configuration"] = {"ANSI_MODE": "false"}
163+
with sql.connect(**params) as c:
164+
with c.cursor() as cur:
165+
cur.execute("SET ANSI_MODE")
166+
rows = cur.fetchall()
167+
kv = {r[0]: r[1] for r in rows}
168+
assert kv.get("ANSI_MODE") == "false", f"got {rows!r}"
169+
170+
171+
# ── Error mapping ─────────────────────────────────────────────────
172+
173+
174+
def test_bad_sql_surfaces_as_databaseerror(conn):
175+
"""Bad SQL should surface as a PEP 249 ``DatabaseError`` with
176+
the kernel's structured fields (`code`, `sql_state`, `query_id`)
177+
attached as attributes — the connector backend re-raises the
178+
kernel's ``SqlError`` to ``DatabaseError`` while preserving the
179+
server-reported state."""
180+
with conn.cursor() as cur:
181+
with pytest.raises(DatabaseError) as exc_info:
182+
cur.execute("SELECT * FROM definitely_not_a_table_xyz_kernel_e2e")
183+
err = exc_info.value
184+
# Structured fields copied off the kernel exception:
185+
assert getattr(err, "code", None) == "SqlError"
186+
assert getattr(err, "sql_state", None) == "42P01"

0 commit comments

Comments
 (0)