feat: add --lazy-auth flag to defer Snowflake connection until first tool use#172
Open
Finndersen wants to merge 1 commit intoSnowflake-Labs:mainfrom
Open
feat: add --lazy-auth flag to defer Snowflake connection until first tool use#172Finndersen wants to merge 1 commit intoSnowflake-Labs:mainfrom
Finndersen wants to merge 1 commit intoSnowflake-Labs:mainfrom
Conversation
…tool use - Add `lazy_auth` parameter to `SnowflakeService.__init__()` (default False) - Add `_ensure_connection()` for idempotent lazy connection initialisation - Replace duplicate inline connection block in `get_connection()` with `_ensure_connection()` - Guard `get_api_headers()` and `get_api_host()` with `_ensure_connection()` - Add `--lazy-auth` CLI flag with `SNOWFLAKE_MCP_LAZY_AUTH` env var support - Fix `object_manager/tools.py` to resolve `root` at call time, not registration time
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
By default the MCP server eagerly authenticates with Snowflake during startup (
SnowflakeService.__init__()). This can be undesirable in environments where the server should start quickly and only authenticate when a tool is actually invoked (e.g. browser-based OAuth flows, or to reduce startup latency).Key Changes
SnowflakeService.__init__()— accepts a newlazy_auth: bool = Falseparameter. WhenTrue, skips connection creation at startup._ensure_connection()— new idempotent method that creates the Snowflake connection andRootobject on first call, and is a no-op thereafter.get_connection()— simplified to use_ensure_connection(), removing a duplicate (and buggy — usedclient_session_keep_alive=False) inline connection block.get_api_headers()/get_api_host()— call_ensure_connection()before accessingself.connectionin the external-env branch, ensuring lazy auth works for REST API calls too.--lazy-authCLI flag — added toparse_arguments()withSNOWFLAKE_MCP_LAZY_AUTHenvironment variable fallback.object_manager/tools.py— fixed a bug whererootwas captured eagerly at tool registration time (would beNonewith lazy auth). Each tool closure now readssnowflake_service.rootat call time.Usage
Tests
test_parse_arguments_default_lazy_auth— default isFalsetest_parse_arguments_lazy_auth_flag—--lazy-authsetsTruetest_parse_arguments_lazy_auth_from_env— env var setsTruetest_snowflake_service_eager_auth_connects_on_init— default behaviour unchangedtest_snowflake_service_lazy_auth_defers_connection—connectnot called on init whenlazy_auth=Truetest_snowflake_service_lazy_auth_connects_on_ensure—_ensure_connection()triggers authtest_snowflake_service_ensure_connection_idempotent— called twice, only one connection created