Skip to content

feat: add --lazy-auth flag to defer Snowflake connection until first tool use#172

Open
Finndersen wants to merge 1 commit intoSnowflake-Labs:mainfrom
Finndersen:feat/lazy-auth
Open

feat: add --lazy-auth flag to defer Snowflake connection until first tool use#172
Finndersen wants to merge 1 commit intoSnowflake-Labs:mainfrom
Finndersen:feat/lazy-auth

Conversation

@Finndersen
Copy link

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 new lazy_auth: bool = False parameter. When True, skips connection creation at startup.
  • _ensure_connection() — new idempotent method that creates the Snowflake connection and Root object on first call, and is a no-op thereafter.
  • get_connection() — simplified to use _ensure_connection(), removing a duplicate (and buggy — used client_session_keep_alive=False) inline connection block.
  • get_api_headers() / get_api_host() — call _ensure_connection() before accessing self.connection in the external-env branch, ensuring lazy auth works for REST API calls too.
  • --lazy-auth CLI flag — added to parse_arguments() with SNOWFLAKE_MCP_LAZY_AUTH environment variable fallback.
  • object_manager/tools.py — fixed a bug where root was captured eagerly at tool registration time (would be None with lazy auth). Each tool closure now reads snowflake_service.root at call time.

Usage

# Defer authentication until first tool invocation
snowflake-labs-mcp --lazy-auth ...

# Or via environment variable
SNOWFLAKE_MCP_LAZY_AUTH=1 snowflake-labs-mcp ...

Tests

  • test_parse_arguments_default_lazy_auth — default is False
  • test_parse_arguments_lazy_auth_flag--lazy-auth sets True
  • test_parse_arguments_lazy_auth_from_env — env var sets True
  • test_snowflake_service_eager_auth_connects_on_init — default behaviour unchanged
  • test_snowflake_service_lazy_auth_defers_connectionconnect not called on init when lazy_auth=True
  • test_snowflake_service_lazy_auth_connects_on_ensure_ensure_connection() triggers auth
  • test_snowflake_service_ensure_connection_idempotent — called twice, only one connection created

…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant