Add session-scoped connection support. by sfc-gh-jkinkead · Pull Request #13506 · streamlit/streamlit

sfc-gh-jkinkead · 2026-01-05T22:57:03Z

Describe your changes

Add class-level scope method to BaseConnection. Use it when caching instances created by st.connection.

Add instance-level close method to BaseConnection. Register it as the on_release hook when caching instances created by st.connection.

This is part of the session-scoped connections plan, and will be used to create a Snowflake RCR connection.

Testing Plan

See unit tests.

Contribution License Agreement

By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

Add class-level `scope` method to `BaseConnection`. Use it when caching instances created by `st.connection`. Add instance-level `close` method to `BaseConnection`. Register it as the `on_release` hook when caching instances created by `st.connection`.

snyk-io · 2026-01-05T22:57:16Z

✅ Snyk checks have passed. No issues have been found so far.

Status	Scanner	Critical	High	Medium	Low	Total (0)
✅	Open Source Security	0	0	0	0	0 issues
✅	Licenses	0	0	0	0	0 issues

💻 Catch issues earlier using the plugins for VS Code, JetBrains IDEs, Visual Studio, and Eclipse.

github-actions · 2026-01-05T22:57:19Z

✅ PR preview is ready!

Name	Link
📦 Wheel file	https://core-previews.s3-us-west-2.amazonaws.com/pr-13506/streamlit-1.52.2-py3-none-any.whl
📦 `@streamlit/component-v2-lib`	Download from artifacts
🕹️ Preview app	pr-13506.streamlit.app (☁️ Deploy here if not accessible)

lib/tests/streamlit/runtime/connection_factory_test.py

Copilot

Pull request overview

This PR adds support for session-scoped connections by introducing two new methods to BaseConnection: a scope() class method to define connection scope ("global" or "session") and a close() instance method for cleanup when connections are evicted from the cache.

Key Changes

BaseConnection gains scope() class method (returns "global" by default) and close() instance method (no-op by default)
connection_factory validates the scope, passes it to cache_resource, and registers close() as the on_release hook
Three new unit tests verify scope validation, scope propagation to cache, and close callback registration

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
lib/streamlit/connections/base_connection.py	Adds `scope()` class method and `close()` instance method to BaseConnection
lib/streamlit/runtime/connection_factory.py	Validates connection scope, passes scope and on_release hook to cache_resource, updates docstring
lib/tests/streamlit/runtime/connection_factory_test.py	Adds tests for scope validation, scope parameter passing, and close callback

lib/streamlit/connections/base_connection.py

lib/tests/streamlit/runtime/connection_factory_test.py

lib/streamlit/connections/base_connection.py

lib/streamlit/runtime/connection_factory.py

Add init call.

lib/tests/streamlit/runtime/connection_factory_test.py

cursor

✅ Bugbot reviewed your changes and found no bugs!

github-actions · 2026-01-06T01:21:02Z

Summary

This PR adds session-scoped connection support to Streamlit's BaseConnection class as part of the approved spec for session-scoped connections and Snowflake RCR (Restricted Caller's Rights). The changes introduce:

A new scope() classmethod on BaseConnection that returns "global" (default) or "session" to control how connections are cached
A new close() method on BaseConnection that serves as a cleanup hook when connections are released from the cache
Updates to connection_factory to pass these values to the underlying cache_resource decorator
Comprehensive unit tests covering the new functionality

Code Quality

The code follows Streamlit's established patterns and conventions well:

lib/streamlit/connections/base_connection.py:

✅ Proper use of from __future__ import annotations
✅ Type annotations using Literal["global", "session"]
✅ NumPy-style docstrings explaining the methods' purposes
✅ Default implementations maintain backwards compatibility (scope() returns "global", close() is a no-op)

lib/streamlit/runtime/connection_factory.py:

✅ Good validation of scope values with a user-friendly error message (lines 100-104)
✅ Clean wrapper function for on_release (lines 106-107)
✅ Docstring updated to clarify that only "global" scoped connections are shared between sessions (line 236-237)

Design Note: The spec document shows scope as an instance property, but the implementation uses a classmethod. The classmethod approach is actually better design because scope is a characteristic of the connection class, not individual instances. This allows the factory to query the scope before creating instances.

Test Coverage

The test coverage in lib/tests/streamlit/runtime/connection_factory_test.py is comprehensive:

Test	Coverage
`test_friendly_error_for_bad_scope` (lines 213-227)	Tests that invalid scope values raise a clear `StreamlitAPIException`
`test_scope_is_passed_to_cache` (lines 229-262)	Tests both `"session"` and `"global"` scopes are correctly passed to `cache_resource`
`test_close_is_passed_to_cache` (lines 264-287)	Tests that `close()` is registered as the `on_release` callback

The tests appropriately:

✅ Use pytest fixtures and assertions
✅ Include docstrings explaining the test purpose
✅ Mock cache_resource to verify arguments without testing the caching system itself (which has its own tests)
✅ Test error conditions with descriptive error message matching

E2E Tests: No e2e tests were added, which is appropriate. This PR adds infrastructure for future features (session-scoped connections) and the actual functionality is a pass-through to existing cache_resource behavior that already has test coverage.

Backwards Compatibility

Excellent backward compatibility:

Default scope() returns "global" - All existing BaseConnection subclasses will continue to work exactly as before with global caching
Default close() is a no-op - Existing connections won't break; they simply won't perform cleanup
No changes to st.connection() signature - User code calling st.connection() is unaffected
Existing first-party connections unchanged - SQLConnection, SnowflakeConnection, and SnowparkConnection don't override these methods, maintaining their current behavior

Security & Risk

Low risk assessment:

✅ The changes are additive and don't modify existing behavior paths
✅ Scope validation prevents injection of invalid scope values
✅ The close() method follows the established pattern from cache_resource's on_release parameter
✅ Exception handling in close() is handled by the caching layer (exceptions bubble up appropriately according to existing cache code)

Minor consideration: The close() method could throw exceptions during cleanup. However, reviewing the caching code shows this is already handled - exceptions in on_release functions bubble up as user script errors, and the cache continues clearing other entries. This is consistent with established behavior.

Recommendations

The PR is well-implemented and ready for merge. A few minor optional improvements:

Documentation clarity (optional): The updated docstring in connection_factory mentions that "global" scoped connections are shared between sessions, but doesn't explain what happens with "session" scoped connections. Consider adding: "Connection types with a scope of "session" will be cached per-session and not shared."
Future consideration (not blocking): When implementing actual session-scoped connections (like the planned Snowflake RCR connection), ensure the close() implementations are thread-safe since they may be called during session teardown from different threads.

Verdict

APPROVED: This PR adds well-designed, backwards-compatible infrastructure for session-scoped connections with comprehensive unit test coverage. The implementation follows Streamlit conventions, aligns with the approved product spec, and introduces no regressions to existing functionality.

This is an automated AI review. Please verify the feedback and use your judgment.

lukasmasuch

LGTM 👍

## Describe your changes Fixes an issue with Snowflake connections not getting re-initialized after having been closed. <details> <summary>Claude issue analysis</summary> # Issue Report: Snowflake Connection "Connection is closed" Error ## Summary After recent PRs adding `on_release` to `st.cache_resource` and session-scoped connection support, users may encounter a `snowflake.connector.errors.DatabaseError: 250002 (08003): Connection is closed` error when using `st.connection("snowflake")` with cached data queries. ## Related PRs - **PR #13439**: Add `on_release` to `st.cache_resource` - **PR #13482**: Add session scoping to caches - **PR #13538**: Add `SnowflakeCallersRightsConnection` - **PR #13506**: Add session-scoped connection support ## Error Details ``` Traceback (most recent call last): File ".../streamlit/runtime/scriptrunner/exec_code.py", line 129, in exec_func_with_error_handling result = func() ... File ".../snowflake/snowpark/_internal/server_connection.py", line 205, in _cursor self._thread_store.cursor = self._conn.cursor() File ".../snowflake/connector/connection.py", line 1270, in cursor Error.errorhandler_wrapper(...) snowflake.connector.errors.DatabaseError: 250002 (08003): Connection is closed ``` ## Root Cause Analysis ### Background The PRs mentioned added important functionality: 1. **PR #13439** added the `on_release` callback to `st.cache_resource`, which is called when cache entries are evicted 2. **PR #13506** modified `connection_factory.py` to use this `on_release` callback to call `connection.close()` when a connection is evicted from the cache This is the relevant code in `connection_factory.py`: ```python def on_release_wrapped(connection: ConnectionClass) -> None: connection.close() __create_connection = cache_resource( max_entries=max_entries, show_spinner="Running `st.connection(...)`.", ttl=ttl, scope=scope, on_release=on_release_wrapped, # Calls close() when evicted )(__create_connection) ``` ### The Bug In `BaseSnowflakeConnection.close()`, after calling `self._raw_instance.close()`, the `_raw_instance` attribute was **NOT** reset to `None`: ```python def close(self) -> None: """Closes the underlying Snowflake connection.""" if self._raw_instance is not None: self._raw_instance.close() # BUG: _raw_instance was NOT set to None! ``` This caused the following issue: 1. When `close()` was called (e.g., via `on_release` when a cache entry is evicted), the underlying connection was closed 2. However, `_raw_instance` still referenced the **closed** connection object 3. The `_instance` property checks `if self._raw_instance is None` to decide whether to create a new connection: ```python @Property def _instance(self) -> RawConnectionT: if self._raw_instance is None: self._raw_instance = self._connect(**self._kwargs) return self._raw_instance ``` 4. Since `_raw_instance` wasn't `None`, subsequent access to `_instance` returned the **CLOSED** connection 5. Any operations on the closed connection failed with "Connection is closed" ### When This Bug Manifests The `on_release` callback (which calls `close()`) is triggered when: - Cache entries expire due to TTL - Cache is full and oldest entries are evicted (`max_entries`) - `st.cache_resource.clear()` is called - For session-scoped caches: when a session disconnects For global-scoped connections like `st.connection("snowflake")`, this typically only happens if: - `st.cache_resource.clear()` is called explicitly - TTL is set and expires - `max_entries` is set and exceeded ### Additional Consideration: Snowpark Sessions When users call `conn.session()`, they get a Snowpark Session that internally references `self._instance`. If the underlying connection is closed: ```python def session(self) -> Session: if running_in_sis(): return get_active_session() return Session.builder.configs({"connection": self._instance}).create() ``` Any Snowpark Sessions created from the connection will also fail because they hold a reference to the now-closed underlying connection object. ## Fix The fix is simple: reset `_raw_instance` to `None` after closing the connection: ```python def close(self) -> None: """Closes the underlying Snowflake connection.""" if self._raw_instance is not None: self._raw_instance.close() self._raw_instance = None # Added this line ``` This ensures that after `close()` is called, the next access to `_instance` will create a new connection instead of returning the closed one. ## Files Changed 1. **`lib/streamlit/connections/snowflake_connection.py`** - Fixed `close()` method to reset `_raw_instance = None` after closing 2. **`lib/tests/streamlit/connections/snowflake_connection_test.py`** - Added `TestSnowflakeConnectionClose` test class with: - `test_close_resets_raw_instance`: Verifies that `close()` closes the connection AND resets `_raw_instance` - `test_close_is_noop_when_not_connected`: Verifies that `close()` doesn't fail when `_raw_instance` is already `None` ## Testing ```bash PYTHONPATH=lib pytest lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose -v ``` Output: ``` lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose::test_close_resets_raw_instance PASSED lib/tests/streamlit/connections/snowflake_connection_test.py::TestSnowflakeConnectionClose::test_close_is_noop_when_not_connected PASSED ``` ## Recommendations for Users Until this fix is released, users experiencing this issue can: 1. **Avoid storing Snowpark Sessions long-term**: Instead of caching Snowpark Sessions, create them fresh when needed 2. **Check if using `st.cache_resource.clear()`**: If calling this anywhere in the app, it will close all cached connections 3. **Consider connection TTL settings**: If TTL is set on the connection, it may expire and close ## Impact - **Affected**: Users of `st.connection("snowflake")` and `st.connection("snowflake-callers-rights")` who experience cache eviction scenarios - **Severity**: Medium - The bug causes operations to fail with a confusing error message, but the workaround (restarting the app or avoiding cache clears) is available - **Scope**: Only affects `SnowflakeConnection` and its subclasses; other connection types (`SQLConnection`, `SnowparkConnection`) inherit the no-op `close()` from `BaseConnection` and are not affected </details> ## GitHub Issue Link (if applicable) ## Testing Plan - Added unit test. --- **Contribution License Agreement** By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.

Copilot AI review requested due to automatic review settings January 5, 2026 22:57

Copilot started reviewing on behalf of sfc-gh-jkinkead January 5, 2026 22:57 View session

github-code-quality bot found potential problems Jan 5, 2026

View reviewed changes

lib/tests/streamlit/runtime/connection_factory_test.py Fixed Show fixed Hide fixed

Copilot AI reviewed Jan 5, 2026

View reviewed changes

Doc addition.

126f1fe

Add init call.

graphite-app bot reviewed Jan 5, 2026

View reviewed changes

lib/tests/streamlit/runtime/connection_factory_test.py Outdated Show resolved Hide resolved

lib/tests/streamlit/runtime/connection_factory_test.py Outdated Show resolved Hide resolved

lib/tests/streamlit/runtime/connection_factory_test.py Outdated Show resolved Hide resolved

Tighten type scopes for Graphite.

9fd2551

cursor bot reviewed Jan 6, 2026

View reviewed changes

lukasmasuch added security-assessment-completed change:feature PR contains new feature or enhancement implementation impact:users PR changes affect end users ai-review If applied to PR or issue will run AI review workflow labels Jan 6, 2026

github-actions bot removed the ai-review If applied to PR or issue will run AI review workflow label Jan 6, 2026

lukasmasuch approved these changes Jan 6, 2026

View reviewed changes

sfc-gh-jkinkead merged commit 4fbaa89 into develop Jan 6, 2026
83 of 89 checks passed

sfc-gh-jkinkead deleted the jkinkead-add-scoped-connections branch January 6, 2026 20:25

lukasmasuch mentioned this pull request Jan 21, 2026

Fix a regression with snowflake connection being closed #13665

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add session-scoped connection support.#13506

Add session-scoped connection support.#13506
sfc-gh-jkinkead merged 3 commits intodevelopfrom
jkinkead-add-scoped-connections

sfc-gh-jkinkead commented Jan 5, 2026 •

edited

Loading

Uh oh!

snyk-io bot commented Jan 5, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Jan 5, 2026 •

edited

Loading

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor bot left a comment

Uh oh!

github-actions bot commented Jan 6, 2026

Uh oh!

lukasmasuch left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

sfc-gh-jkinkead commented Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

Testing Plan

Uh oh!

snyk-io bot commented Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Snyk checks have passed. No issues have been found so far.

Uh oh!

github-actions bot commented Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ PR preview is ready!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Key Changes

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 6, 2026

Summary

Code Quality

Test Coverage

Backwards Compatibility

Security & Risk

Recommendations

Verdict

Uh oh!

lukasmasuch left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sfc-gh-jkinkead commented Jan 5, 2026 •

edited

Loading

snyk-io bot commented Jan 5, 2026 •

edited

Loading

github-actions bot commented Jan 5, 2026 •

edited

Loading