Name the Snowflake sessions the client opens

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.

CHANGELOG.md

@@ -9,6 +9,12 @@ 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
+  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]`

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/<version>`, 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 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.

packages/structured-data-models-connectors/src/sdm_connectors/backends/snowflake.py

@@ -28,6 +28,10 @@
 
 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'
+
 _AUTH_ARGS = frozenset(
     {
         'account',
@@ -103,6 +107,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``
+    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)
@@ -122,7 +130,24 @@ def connect(
             return borrowed
         if not kwargs:
             return _connect_without_arguments()
-    return snowflake_connector.connect(**kwargs)
+    return _NamedConnection(**kwargs)
+
+
+class _NamedConnection(Connection):
+    r"""A connection that names itself :data:`APPLICATION` if nothing else does.
+
+    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.
+    """
+
+    _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:
@@ -134,7 +159,7 @@ def _connect_without_arguments() -> Connection:
     where it is configured; only the failure is re-stated in the client's terms.
     """
     try:
-        connection = snowflake_connector.connect()
+        connection = _NamedConnection()
     except Exception as error:
         raise ConnectorError(
             f'snowflake connector requires connection arguments (at least '

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,
     )
 
@@ -163,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
 
@@ -182,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()
@@ -687,3 +687,97 @@ 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
     )
+
+
+_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', 'detected', 'account', 'application'),
+    [
+        (_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': _THE_FILE},
+            None,
+            'from_default',
+            'nvidia/kumo-sdm',
+        ),
+        ({}, None, 'from_default', 'nvidia/kumo-sdm'),
+    ],
+)
+def test_snowflake_names_a_session_only_where_nothing_else_does(
+    monkeypatch, tmp_path, arguments, detected, account, 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
+
+    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(
+        manager,
+        '_slices',
+        [
+            s._replace(path=connections) if s.section == 'connections' else s
+            for s in manager._slices
+        ],
+    )
+    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()
+    }
+
+    with pytest.raises((_Resolved, ConnectorError)):
+        backend.connect(**arguments)
+
+    assert (seen['account'], seen['application']) == (account, application)