Skip to content

Commit 4b03750

Browse files
sfc-gh-jcieslaksfc-gh-jcieslak
authored
Test and document session parameters support (#644)
* add test * wip * Test query tag usage * Add TOC and docs for session parameters * Update readme * Remove unit test * Adjust pyproject * Adjust pyproject * Adjust pyproject * Fix execute commits on connection (some versions don't have commit on connection object) * Add description
1 parent b95de49 commit 4b03750

File tree

3 files changed

+165
-1
lines changed

3 files changed

+165
-1
lines changed

‎DESCRIPTION.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ Source code is also available at:
2222
- Add `pytest-xdist` parallel test support via per-worker schema provisioning hooks.
2323
- Bump `pandas` lower bound in `sa14` test environment from `<2.1` to `>=2.1.1,<2.2` to ensure pre-built wheels are available for Python 3.12
2424
- Fix SQLAlchemy version parsing (SNOW-3066571)
25+
- Document support for session parameters (like [QUERY_TAG](https://docs.snowflake.com/en/sql-reference/parameters#query-tag)), references: [#644](https://github.com/snowflakedb/snowflake-sqlalchemy/issues/495)
2526

2627
# Release Notes
2728

‎README.md‎

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,47 @@
88

99
Snowflake SQLAlchemy runs on the top of the Snowflake Connector for Python as a [dialect](http://docs.sqlalchemy.org/en/latest/dialects/) to bridge a Snowflake database and SQLAlchemy applications.
1010

11-
11+
Table of contents:
12+
<!-- TOC -->
13+
* [Snowflake SQLAlchemy](#snowflake-sqlalchemy)
14+
* [Prerequisites](#prerequisites)
15+
* [Snowflake Connector for Python](#snowflake-connector-for-python)
16+
* [Data Analytics and Web Application Frameworks (Optional)](#data-analytics-and-web-application-frameworks-optional)
17+
* [Installing Snowflake SQLAlchemy](#installing-snowflake-sqlalchemy)
18+
* [Verifying Your Installation](#verifying-your-installation)
19+
* [Parameters and Behavior](#parameters-and-behavior)
20+
* [Connection Parameters](#connection-parameters)
21+
* [Escaping Special Characters such as `%, @` signs in Passwords](#escaping-special-characters-such-as---signs-in-passwords)
22+
* [Using a proxy server](#using-a-proxy-server)
23+
* [Using session parameters](#using-session-parameters)
24+
* [Opening and Closing Connection](#opening-and-closing-connection)
25+
* [Auto-increment Behavior](#auto-increment-behavior)
26+
* [Object Name Case Handling](#object-name-case-handling)
27+
* [Index Support](#index-support)
28+
* [Single Column Index](#single-column-index)
29+
* [Multi-Column Index](#multi-column-index)
30+
* [Numpy Data Type Support](#numpy-data-type-support)
31+
* [DECFLOAT Data Type Support](#decfloat-data-type-support)
32+
* [DECFLOAT Precision](#decfloat-precision)
33+
* [VECTOR Data Type Support](#vector-data-type-support)
34+
* [Cache Column Metadata](#cache-column-metadata)
35+
* [VARIANT, ARRAY and OBJECT Support](#variant-array-and-object-support)
36+
* [Structured Data Types Support](#structured-data-types-support)
37+
* [MAP](#map)
38+
* [OBJECT](#object)
39+
* [ARRAY](#array)
40+
* [CLUSTER BY Support](#cluster-by-support)
41+
* [Alembic Support](#alembic-support)
42+
* [Key Pair Authentication Support](#key-pair-authentication-support)
43+
* [Merge Command Support](#merge-command-support)
44+
* [CopyIntoStorage Support](#copyintostorage-support)
45+
* [Iceberg Table with Snowflake Catalog support](#iceberg-table-with-snowflake-catalog-support)
46+
* [Hybrid Table support](#hybrid-table-support)
47+
* [Dynamic Tables support](#dynamic-tables-support)
48+
* [Notes](#notes)
49+
* [Verifying Package Signatures](#verifying-package-signatures)
50+
* [Support](#support)
51+
<!-- TOC -->
1252

1353
## Prerequisites
1454

@@ -184,6 +224,44 @@ engine = create_engine(URL(
184224

185225
Use the supported environment variables, `HTTPS_PROXY`, `HTTP_PROXY` and `NO_PROXY` to configure a proxy server.
186226

227+
#### Using session parameters
228+
229+
Snowflake [session parameters](https://docs.snowflake.com/en/sql-reference/parameters#session-parameters) (such as [`QUERY_TAG`](https://docs.snowflake.com/en/sql-reference/parameters#query-tag)) cannot be set directly through the `URL` helper.
230+
Instead, pass them via the `connect_args` parameter of `create_engine`, using the `session_parameters` dict — the same way you would [through the Python connector](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect#setting-session-parameters):
231+
232+
```python
233+
from snowflake.sqlalchemy import URL
234+
from sqlalchemy import create_engine
235+
236+
engine = create_engine(
237+
URL(
238+
# CONNECTION_PARAMETERS
239+
),
240+
connect_args={
241+
"session_parameters": {
242+
"QUERY_TAG": "SOME_QUERY_TAGS",
243+
}
244+
},
245+
)
246+
```
247+
248+
Session parameters set this way apply to all queries executed within the session.
249+
To change a session parameter for specific queries mid-session, use `ALTER SESSION`:
250+
251+
```python
252+
from sqlalchemy import text
253+
254+
with engine.connect() as conn:
255+
conn.execute(text("ALTER SESSION SET QUERY_TAG = 'batch_job_1'"))
256+
conn.execute(text("...")) # Uses 'batch_job_1'
257+
258+
conn.execute(text("ALTER SESSION SET QUERY_TAG = 'batch_job_2'"))
259+
conn.execute(text("...")) # Uses 'batch_job_2'
260+
261+
conn.execute(text("ALTER SESSION UNSET QUERY_TAG"))
262+
conn.execute(text("...")) # No tag
263+
```
264+
187265
### Opening and Closing Connection
188266

189267
Open a connection by executing `engine.connect()`; avoid using `engine.execute()`. Make certain to close the connection by executing `connection.close()` before

‎tests/test_core.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,91 @@ def test_boolean_query_argument_parsing():
197197
engine.dispose()
198198

199199

200+
def test_query_tag_appears_in_query_history():
201+
"""
202+
Tests that query_tag actually appears in Snowflake's query history.
203+
"""
204+
test_query_tag = "sqlalchemy_history_test_tag"
205+
engine = create_engine(
206+
URL(
207+
**CONNECTION_PARAMETERS,
208+
),
209+
connect_args={
210+
"session_parameters": {
211+
"QUERY_TAG": test_query_tag,
212+
}
213+
},
214+
)
215+
216+
try:
217+
with engine.connect() as conn:
218+
# Execute a simple query
219+
with conn.begin():
220+
conn.execute(text("SELECT 1; -- query tag test"))
221+
222+
# Query the query history to verify the tag
223+
result = conn.execute(
224+
text(
225+
"""
226+
SELECT QUERY_TAG
227+
FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION())
228+
WHERE QUERY_TEXT = 'SELECT 1; -- query tag test'
229+
ORDER BY START_TIME DESC LIMIT 1
230+
"""
231+
)
232+
)
233+
row = result.fetchone()
234+
assert row is not None, "Query should be found in history"
235+
assert row[0] == test_query_tag, f"Query tag should be '{test_query_tag}'"
236+
finally:
237+
engine.dispose()
238+
239+
240+
def test_query_tag_per_query():
241+
"""
242+
Tests setting query_tag dynamically per query or per set of SQL actions
243+
using ALTER SESSION SET QUERY_TAG.
244+
"""
245+
engine = create_engine(URL(**CONNECTION_PARAMETERS))
246+
try:
247+
with engine.connect() as conn:
248+
# Set query_tag for first set of operations
249+
with conn.begin():
250+
conn.execute(text("ALTER SESSION SET QUERY_TAG = 'batch_job_1'"))
251+
conn.execute(text("SELECT 1; -- batch 1 query"))
252+
253+
# Change query_tag for second set of operations
254+
with conn.begin():
255+
conn.execute(text("ALTER SESSION SET QUERY_TAG = 'batch_job_2'"))
256+
conn.execute(text("SELECT 2; -- batch 2 query"))
257+
258+
# Reset query_tag (unset it)
259+
with conn.begin():
260+
conn.execute(text("ALTER SESSION UNSET QUERY_TAG"))
261+
conn.execute(text("SELECT 3; -- no tag query"))
262+
263+
# Verify the query tags in history
264+
result = conn.execute(
265+
text(
266+
"""
267+
SELECT QUERY_TEXT, QUERY_TAG
268+
FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION())
269+
WHERE QUERY_TEXT LIKE 'SELECT %; -- batch%'
270+
OR QUERY_TEXT LIKE 'SELECT %; -- no tag%'
271+
ORDER BY START_TIME DESC
272+
"""
273+
)
274+
)
275+
rows = result.fetchall()
276+
query_tags = {row[0]: row[1] for row in rows}
277+
278+
assert query_tags.get("SELECT 1; -- batch 1 query") == "batch_job_1"
279+
assert query_tags.get("SELECT 2; -- batch 2 query") == "batch_job_2"
280+
assert query_tags.get("SELECT 3; -- no tag query") in (None, "")
281+
finally:
282+
engine.dispose()
283+
284+
200285
def test_create_dialect():
201286
"""
202287
Tests getting only dialect object through create_engine

0 commit comments

Comments
 (0)