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 16e6f6926..d76ae114c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,12 @@ these packages. All five distributions share one version and one entry. fitted as the NIM fits it. A request the model cannot answer is refused with a `StructuredDataError` whose code says why. Install it with the `in-process` extra. +- 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]` diff --git a/docs/about/architecture.md b/docs/about/architecture.md index d185a261b..cda67c5dd 100644 --- a/docs/about/architecture.md +++ b/docs/about/architecture.md @@ -283,6 +283,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 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 6f464de60..0f411ecff 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 @@ -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,29 @@ 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. Where + it loads the default connection instead, it drops what it detected, so + the detection is asked again before this name is used. + """ + + _application_defaulted = False + + def connect(self, **kwargs: Any) -> None: + if not self._application_defaulted: + self._application_defaulted = True + if 'application' not in kwargs: + kwargs['application'] = ( + self._detect_application() or APPLICATION + ) + super().connect(**kwargs) def _connect_without_arguments() -> Connection: @@ -134,7 +164,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 ' diff --git a/packages/structured-data-models-connectors/tests/test_backends.py b/packages/structured-data-models-connectors/tests/test_backends.py index eaf9e241a..9dec7e593 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, ) @@ -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,104 @@ 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', + ), + ( + {'connections_file_path': _THE_FILE}, + 'SF_PARTNER', + 'from_default', + 'their_partner', + ), + ({}, None, 'from_default', 'nvidia/kumo-sdm'), + ({}, 'streamlit', 'from_default', 'streamlit'), + ], +) +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)