From 05de740ccd2e8251ace5664c8d62d72990315100 Mon Sep 17 00:00:00 2001 From: Manush Murali Date: Fri, 23 Oct 2026 11:56:13 -0700 Subject: [PATCH 1/2] Name the Snowflake sessions the client opens Snowflake attributes a session's queries by the application it names at login, and the client named none, so its sessions read as the driver's default. Each session it opens now passes application="nvidia/kumo-sdm", unless the caller names one through application or SF_PARTNER. The driver reads its default connection only when given nothing but a connections file. Any other argument, application included, makes it skip that connection, so in that case the default connection is named explicitly; otherwise the account and user would be dropped. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Manush Murali --- AGENTS.md | 3 + CHANGELOG.md | 4 ++ docs/about/architecture.md | 1 + .../src/sdm_connectors/backends/snowflake.py | 39 ++++++++++-- .../tests/test_backends.py | 62 ++++++++++++++++++- 5 files changed, 102 insertions(+), 7 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 15e798218..ff38884a6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -28,6 +28,9 @@ A rename that sweeps the tree will change these, and nothing will fail: `SAMPLING_QUERY_TAG` sits beside the second. They are what a warehouse sees, so changing them splits Databricks partner telemetry and query history at that version. +- `APPLICATION` in `sdm_connectors/backends/snowflake.py`, `nvidia/kumo-sdm`. + It is the application name Snowflake attributes this client's sessions to, + so changing it splits that attribution at that version. - `_term('sdm-relational-engine')` in `rfm/graph_fingerprint.py`. It is hashed. Changing it gives identical graphs different names, which the module says only `FINGERPRINT_VERSION` may do. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f8734742..4bfcaee38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ these packages. All five distributions share one version and one entry. ### Added +- A Snowflake session the client opens names itself to Snowflake as + `nvidia/kumo-sdm`, so Snowflake can attribute its queries. An `application` + you pass, or `SF_PARTNER`, is kept instead. A connection or Snowpark session + you open yourself is used as it is. - The relational engine publishes wheels for macOS (Apple silicon and Intel) and 64-bit Windows, and for CPython 3.14, alongside the `manylinux_2_28` x86-64 and aarch64 wheels it already built, so `structured-data-models-client[relational]` diff --git a/docs/about/architecture.md b/docs/about/architecture.md index 80cdd88b9..0579c4a2f 100644 --- a/docs/about/architecture.md +++ b/docs/about/architecture.md @@ -282,6 +282,7 @@ It is the type the engine already accepts, so there is no second connection type On Databricks, `sdm_connectors.databricks.warehouse()` and `lakebase()` open one as the caller's `WorkspaceClient` and hand back that same native connection, for the caller to close. Every connection and API request structured-data-models-client opens there names it as `nvidia_structured-data-models-client/`, as Databricks requires of a partner product, and a warehouse session it opens tags its statements `structured_data_models_client_purpose: sampling`. Each package keeps a private copy of the identifier, so the connectors never import the client and the serving extra brings no connectors; a test holds the two copies to the same value. +A Snowflake session it opens names itself to Snowflake as `nvidia/kumo-sdm`, the connector's `application`, unless an `application` argument or `SF_PARTNER` names another. A connection the caller opens and passes in is used as it is. A notebook's `spark` is wrapped in the few calls the samplers make: a cursor that runs `spark.sql` with positional parameters and returns Arrow. A schema over a source does not pickle. diff --git a/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py b/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py index 6f464de60..0b07f837c 100644 --- a/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py +++ b/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py @@ -5,6 +5,7 @@ from __future__ import annotations import inspect +import os from typing import TYPE_CHECKING, Any, TypeAlias from sdm_connectors.sql import ( @@ -28,6 +29,12 @@ Connection: TypeAlias = snowflake_connector.SnowflakeConnection +#: The application every Snowflake session this client opens names itself as, +#: which Snowflake attributes the session's queries to. +APPLICATION = 'nvidia/kumo-sdm' + +_FILE_ARGS = frozenset({'connection_name', 'connections_file_path'}) + _AUTH_ARGS = frozenset( { 'account', @@ -103,6 +110,10 @@ def connect( the escape hatch can never quietly replace a validated credential. What it does supply still counts towards the borrow decision, so credentials passed that way are honoured rather than silently ignored. + + A connection it opens names itself to Snowflake as ``nvidia/kumo-sdm``, + unless ``application`` or the driver's ``SF_PARTNER`` names another. A + borrowed Snowpark session is used as it is. """ check_connect_args('snowflake', kwargs, _CONNECT_ARGS) kwargs = merge_driver_options('snowflake', kwargs, driver_options) @@ -122,19 +133,35 @@ def connect( return borrowed if not kwargs: return _connect_without_arguments() - return snowflake_connector.connect(**kwargs) + return snowflake_connector.connect(**_named(kwargs)) + + +def _named(kwargs: dict[str, Any]) -> dict[str, Any]: + r"""Return ``kwargs`` with :data:`APPLICATION`, unless the caller names + an application as ``application`` or ``SF_PARTNER``. + + The driver reads its default connection only when given nothing but a + connections file, so that connection is named explicitly to keep it. + """ + named = dict(kwargs) + if named.get('connection_name') is None and set(named) <= _FILE_ARGS: + manager = snowflake_connector.config_manager.CONFIG_MANAGER + named['connection_name'] = manager['default_connection_name'] + if 'application' not in named and not os.environ.get('SF_PARTNER'): + named['application'] = APPLICATION + return named def _connect_without_arguments() -> Connection: r"""Connect with no arguments, explaining what this client expects on failure. - With no arguments the driver falls back to its own default-connection file, - whose absence it reports as "Default connection with name 'default' cannot - be found", a feature this package never mentions. The fallback still works - where it is configured; only the failure is re-stated in the client's terms. + With no arguments the driver falls back to its own default connection, + whose absence it reports in terms of a connections file this package + never mentions. The fallback still works where it is configured; only + the failure is re-stated in the client's terms. """ try: - connection = snowflake_connector.connect() + connection = snowflake_connector.connect(**_named({})) except Exception as error: raise ConnectorError( f'snowflake connector requires connection arguments (at least ' diff --git a/packages/structured-data-models-connectors/tests/test_backends.py b/packages/structured-data-models-connectors/tests/test_backends.py index eaf9e241a..d0276e520 100644 --- a/packages/structured-data-models-connectors/tests/test_backends.py +++ b/packages/structured-data-models-connectors/tests/test_backends.py @@ -126,7 +126,12 @@ def test_snowflake_driver_options_credentials_suppress_borrowing(monkeypatch): ) assert result is opened - assert seen == {'account': 'a', 'user': 'u', 'password': 'p'} + assert seen == { + 'account': 'a', + 'user': 'u', + 'password': 'p', + 'application': 'nvidia/kumo-sdm', + } def test_snowflake_borrows_session_when_no_auth_arguments(monkeypatch): @@ -687,3 +692,58 @@ def test_a_databricks_default_is_taken_from_the_environment_only_when_absent( assert {k: v for k, v in sent.items() if k in ('catalog', 'schema')} == ( taken ) + + +_NAMED = {'application': 'nvidia/kumo-sdm'} +_DEFAULT = {'connection_name': 'default'} +_CREDENTIALS = {'account': 'a', 'user': 'u'} + + +@pytest.mark.parametrize( + ('arguments', 'partner', 'sent'), + [ + ({}, None, {**_DEFAULT, **_NAMED}), + (_CREDENTIALS, None, {**_CREDENTIALS, **_NAMED}), + ({'connection_name': 'c'}, None, {'connection_name': 'c', **_NAMED}), + ( + {'connections_file_path': 'f.toml'}, + None, + {'connections_file_path': 'f.toml', **_DEFAULT, **_NAMED}, + ), + ( + {**_CREDENTIALS, 'application': 'theirs'}, + None, + {**_CREDENTIALS, 'application': 'theirs'}, + ), + ({}, 'their_partner', _DEFAULT), + (_CREDENTIALS, 'their_partner', _CREDENTIALS), + ], +) +def test_snowflake_names_the_sessions_it_opens( + monkeypatch, arguments, partner, sent +): + r"""The driver reads its default connection only when given nothing but a + connections file, so that connection is named alongside the application.""" + pytest.importorskip('snowflake.connector') + from sdm_connectors.backends import snowflake as backend + + seen = {} + monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: None) + monkeypatch.setattr( + backend.snowflake_connector.config_manager, + 'CONFIG_MANAGER', + {'default_connection_name': 'default'}, + ) + monkeypatch.setattr( + backend.snowflake_connector, + 'connect', + lambda **kwargs: seen.update(kwargs) or _Conn(), + ) + if partner is None: + monkeypatch.delenv('SF_PARTNER', raising=False) + else: + monkeypatch.setenv('SF_PARTNER', partner) + + backend.connect(**arguments) + + assert seen == sent From c31f3eb2938857712e756a73240c794889505462 Mon Sep 17 00:00:00 2001 From: Manush Murali Date: Fri, 23 Oct 2026 19:02:11 -0700 Subject: [PATCH 2/2] Name a Snowflake session only where nothing else names it The application was added to the caller's arguments before the driver saw them, which meant re-implementing the driver's choice of default connection and overriding everything the driver resolves itself: - an `application` in the `connections.toml` profile, because the arguments win over the profile; - the driver's own tags for Streamlit, Jupyter and Snowflake notebooks, which it adds only when no `application` is given; - the default connection when a caller named another connections file, which was read from `~/.snowflake` before the driver switched files. `_NamedConnection` subclasses the driver's connection and adds `nvidia/kumo-sdm` in `connect`, which the driver first calls with what it resolved from the arguments, the profile and its detection. A name any of them gives is kept, and the default connection, the connections file and the no-argument error are the driver's own again. The test runs the driver's real resolution over a real connections.toml and stops at login. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Manush Murali --- CHANGELOG.md | 8 +- docs/about/architecture.md | 2 +- .../src/sdm_connectors/backends/snowflake.py | 46 +++---- .../tests/test_backends.py | 130 +++++++++++------- 4 files changed, 110 insertions(+), 76 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ea4d236d1..d6bd4906f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,9 +10,11 @@ these packages. All five distributions share one version and one entry. ### Added - A Snowflake session the client opens names itself to Snowflake as - `nvidia/kumo-sdm`, so Snowflake can attribute its queries. An `application` - you pass, or `SF_PARTNER`, is kept instead. A connection or Snowpark session - you open yourself is used as it is. + `nvidia/kumo-sdm`, so Snowflake can attribute its queries. An application + named any other way is kept: an `application` argument, one in the + `connections.toml` profile, `SF_PARTNER`, or the connector's own detection + of Streamlit and notebooks. A connection or Snowpark session you open + yourself is used as it is. - The relational engine publishes wheels for macOS (Apple silicon and Intel) and 64-bit Windows, and for CPython 3.14, alongside the `manylinux_2_28` x86-64 and aarch64 wheels it already built, so `structured-data-models-client[relational]` diff --git a/docs/about/architecture.md b/docs/about/architecture.md index 0579c4a2f..cfd1d24a8 100644 --- a/docs/about/architecture.md +++ b/docs/about/architecture.md @@ -282,7 +282,7 @@ It is the type the engine already accepts, so there is no second connection type On Databricks, `sdm_connectors.databricks.warehouse()` and `lakebase()` open one as the caller's `WorkspaceClient` and hand back that same native connection, for the caller to close. Every connection and API request structured-data-models-client opens there names it as `nvidia_structured-data-models-client/`, as Databricks requires of a partner product, and a warehouse session it opens tags its statements `structured_data_models_client_purpose: sampling`. Each package keeps a private copy of the identifier, so the connectors never import the client and the serving extra brings no connectors; a test holds the two copies to the same value. -A Snowflake session it opens names itself to Snowflake as `nvidia/kumo-sdm`, the connector's `application`, unless an `application` argument or `SF_PARTNER` names another. +A Snowflake session it opens names itself to Snowflake as `nvidia/kumo-sdm`, the connector's `application`, unless the arguments, the `connections.toml` profile or the connector's own detection (`SF_PARTNER`, Streamlit, a notebook) already name one. A connection the caller opens and passes in is used as it is. A notebook's `spark` is wrapped in the few calls the samplers make: a cursor that runs `spark.sql` with positional parameters and returns Arrow. A schema over a source does not pickle. diff --git a/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py b/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py index 0b07f837c..990bd1dce 100644 --- a/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py +++ b/packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py @@ -5,7 +5,6 @@ from __future__ import annotations import inspect -import os from typing import TYPE_CHECKING, Any, TypeAlias from sdm_connectors.sql import ( @@ -33,8 +32,6 @@ #: which Snowflake attributes the session's queries to. APPLICATION = 'nvidia/kumo-sdm' -_FILE_ARGS = frozenset({'connection_name', 'connections_file_path'}) - _AUTH_ARGS = frozenset( { 'account', @@ -111,9 +108,9 @@ def connect( does supply still counts towards the borrow decision, so credentials passed that way are honoured rather than silently ignored. - A connection it opens names itself to Snowflake as ``nvidia/kumo-sdm``, - unless ``application`` or the driver's ``SF_PARTNER`` names another. A - borrowed Snowpark session is used as it is. + A connection it opens names itself to Snowflake as ``nvidia/kumo-sdm`` + when nothing else names it; see :class:`_NamedConnection`. A borrowed + Snowpark session is used as it is. """ check_connect_args('snowflake', kwargs, _CONNECT_ARGS) kwargs = merge_driver_options('snowflake', kwargs, driver_options) @@ -133,35 +130,36 @@ def connect( return borrowed if not kwargs: return _connect_without_arguments() - return snowflake_connector.connect(**_named(kwargs)) + return _NamedConnection(**kwargs) -def _named(kwargs: dict[str, Any]) -> dict[str, Any]: - r"""Return ``kwargs`` with :data:`APPLICATION`, unless the caller names - an application as ``application`` or ``SF_PARTNER``. +class _NamedConnection(Connection): + r"""A connection that names itself :data:`APPLICATION` if nothing else does. - The driver reads its default connection only when given nothing but a - connections file, so that connection is named explicitly to keep it. + The driver resolves the arguments, a ``connections.toml`` profile and its + own detection of ``SF_PARTNER``, Streamlit and notebooks before it first + calls ``connect``, so an ``application`` any of them gives is kept. """ - named = dict(kwargs) - if named.get('connection_name') is None and set(named) <= _FILE_ARGS: - manager = snowflake_connector.config_manager.CONFIG_MANAGER - named['connection_name'] = manager['default_connection_name'] - if 'application' not in named and not os.environ.get('SF_PARTNER'): - named['application'] = APPLICATION - return named + + _application_defaulted = False + + def connect(self, **kwargs: Any) -> None: + if not self._application_defaulted: + self._application_defaulted = True + kwargs.setdefault('application', APPLICATION) + super().connect(**kwargs) def _connect_without_arguments() -> Connection: r"""Connect with no arguments, explaining what this client expects on failure. - With no arguments the driver falls back to its own default connection, - whose absence it reports in terms of a connections file this package - never mentions. The fallback still works where it is configured; only - the failure is re-stated in the client's terms. + With no arguments the driver falls back to its own default-connection file, + whose absence it reports as "Default connection with name 'default' cannot + be found", a feature this package never mentions. The fallback still works + where it is configured; only the failure is re-stated in the client's terms. """ try: - connection = snowflake_connector.connect(**_named({})) + connection = _NamedConnection() except Exception as error: raise ConnectorError( f'snowflake connector requires connection arguments (at least ' diff --git a/packages/structured-data-models-connectors/tests/test_backends.py b/packages/structured-data-models-connectors/tests/test_backends.py index d0276e520..5b437be5b 100644 --- a/packages/structured-data-models-connectors/tests/test_backends.py +++ b/packages/structured-data-models-connectors/tests/test_backends.py @@ -5,6 +5,8 @@ from __future__ import annotations import importlib +import sys +import types import pytest @@ -116,8 +118,8 @@ def test_snowflake_driver_options_credentials_suppress_borrowing(monkeypatch): opened, seen = _Conn(), {} monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: _Conn()) monkeypatch.setattr( - backend.snowflake_connector, - 'connect', + backend, + '_NamedConnection', lambda **kwargs: seen.update(kwargs) or opened, ) @@ -126,12 +128,7 @@ def test_snowflake_driver_options_credentials_suppress_borrowing(monkeypatch): ) assert result is opened - assert seen == { - 'account': 'a', - 'user': 'u', - 'password': 'p', - 'application': 'nvidia/kumo-sdm', - } + assert seen == {'account': 'a', 'user': 'u', 'password': 'p'} def test_snowflake_borrows_session_when_no_auth_arguments(monkeypatch): @@ -168,9 +165,7 @@ def test_snowflake_credentials_still_open_their_own_connection(monkeypatch): opened = _Conn() monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: _Conn()) - monkeypatch.setattr( - backend.snowflake_connector, 'connect', lambda **kwargs: opened - ) + monkeypatch.setattr(backend, '_NamedConnection', lambda **kwargs: opened) assert backend.connect(account='a', user='u', password='p') is opened @@ -187,7 +182,7 @@ def _fail(**kwargs): ) monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: None) - monkeypatch.setattr(backend.snowflake_connector, 'connect', _fail) + monkeypatch.setattr(backend, '_NamedConnection', _fail) with pytest.raises(ConnectorError) as excinfo: backend.connect() @@ -694,56 +689,95 @@ def test_a_databricks_default_is_taken_from_the_environment_only_when_absent( ) -_NAMED = {'application': 'nvidia/kumo-sdm'} -_DEFAULT = {'connection_name': 'default'} -_CREDENTIALS = {'account': 'a', 'user': 'u'} +_PROFILES = """\ +[default] +account = "from_default" +user = "u" +password = "p" + +[plain] +account = "from_plain" +user = "u" +password = "p" + +[named] +account = "from_named" +user = "u" +password = "p" +application = "from_profile" +""" + +_CREDENTIALS = {'account': 'a', 'user': 'u', 'password': 'p'} +_THE_FILE = 'the connections.toml written below' + + +class _Resolved(Exception): + pass @pytest.mark.parametrize( - ('arguments', 'partner', 'sent'), + ('arguments', 'detected', 'account', 'application'), [ - ({}, None, {**_DEFAULT, **_NAMED}), - (_CREDENTIALS, None, {**_CREDENTIALS, **_NAMED}), - ({'connection_name': 'c'}, None, {'connection_name': 'c', **_NAMED}), + (_CREDENTIALS, None, 'a', 'nvidia/kumo-sdm'), + ({**_CREDENTIALS, 'application': 'theirs'}, None, 'a', 'theirs'), + (_CREDENTIALS, 'SF_PARTNER', 'a', 'their_partner'), + (_CREDENTIALS, 'streamlit', 'a', 'streamlit'), + ({'connection_name': 'named'}, None, 'from_named', 'from_profile'), + ({'connection_name': 'plain'}, None, 'from_plain', 'nvidia/kumo-sdm'), ( - {'connections_file_path': 'f.toml'}, + {'connections_file_path': _THE_FILE}, None, - {'connections_file_path': 'f.toml', **_DEFAULT, **_NAMED}, + 'from_default', + 'nvidia/kumo-sdm', ), - ( - {**_CREDENTIALS, 'application': 'theirs'}, - None, - {**_CREDENTIALS, 'application': 'theirs'}, - ), - ({}, 'their_partner', _DEFAULT), - (_CREDENTIALS, 'their_partner', _CREDENTIALS), + ({}, None, 'from_default', 'nvidia/kumo-sdm'), ], ) -def test_snowflake_names_the_sessions_it_opens( - monkeypatch, arguments, partner, sent +def test_snowflake_names_a_session_only_where_nothing_else_does( + monkeypatch, tmp_path, arguments, detected, account, application ): - r"""The driver reads its default connection only when given nothing but a - connections file, so that connection is named alongside the application.""" + r"""The name is added to what the driver resolves, so an application the + caller, a connections.toml profile or the driver's own detection gives + is kept, and the default connection is still found.""" pytest.importorskip('snowflake.connector') from sdm_connectors.backends import snowflake as backend - seen = {} - monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: None) + connections = tmp_path / 'connections.toml' + connections.write_text(_PROFILES) + connections.chmod(0o600) + manager = backend.snowflake_connector.config_manager.CONFIG_MANAGER + monkeypatch.setattr(manager, 'file_path', tmp_path / 'config.toml') monkeypatch.setattr( - backend.snowflake_connector.config_manager, - 'CONFIG_MANAGER', - {'default_connection_name': 'default'}, - ) - monkeypatch.setattr( - backend.snowflake_connector, - 'connect', - lambda **kwargs: seen.update(kwargs) or _Conn(), + manager, + '_slices', + [ + s._replace(path=connections) if s.section == 'connections' else s + for s in manager._slices + ], ) - if partner is None: - monkeypatch.delenv('SF_PARTNER', raising=False) - else: - monkeypatch.setenv('SF_PARTNER', partner) + monkeypatch.setattr(manager, 'conf_file_cache', None) + for name in ('SF_PARTNER', 'SNOWFLAKE_DEFAULT_CONNECTION_NAME'): + monkeypatch.delenv(name, raising=False) + for module in ('streamlit', 'ipykernel', 'snowbooks'): + monkeypatch.delitem(sys.modules, module, raising=False) + if detected == 'SF_PARTNER': + monkeypatch.setenv('SF_PARTNER', 'their_partner') + elif detected == 'streamlit': + monkeypatch.setitem(sys.modules, 'streamlit', types.ModuleType('s')) + seen = {} + + def _resolved(self, **kwargs): + seen.update(kwargs) + raise _Resolved + + monkeypatch.setattr(backend, '_active_snowpark_connection', lambda: None) + monkeypatch.setattr(backend.Connection, 'connect', _resolved) + arguments = { + name: connections if value == _THE_FILE else value + for name, value in arguments.items() + } - backend.connect(**arguments) + with pytest.raises((_Resolved, ConnectorError)): + backend.connect(**arguments) - assert seen == sent + assert (seen['account'], seen['application']) == (account, application)