安全地把 Gemini Enterprise 连接到 Snowflake 数据

USE ROLE ACCOUNTADMIN;
CREATE OR REPLACE SECURITY INTEGRATION <SECURITY_INTEGRATION_NAME>
TYPE = OAUTH
ENABLED = TRUE
OAUTH_CLIENT = CUSTOM
OAUTH_CLIENT_TYPE = 'CONFIDENTIAL'
-- Google Cloud's official handshake receiver address:
OAUTH_REDIRECT_URI = 'https://vertexaisearch.cloud.google.com/oauth-redirect'
OAUTH_ISSUE_REFRESH_TOKENS = TRUE
-- Let the connection stay alive for 90 days (max Snowflake allowed limit)
OAUTH_REFRESH_TOKEN_VALIDITY = 7776000;

-- Retrieve OAuth Client ID and Client Secret
SELECT SYSTEM$SHOW_OAUTH_CLIENT_SECRETS('<SECURITY_INTEGRATION_NAME>');

USE ROLE SECURITYADMIN;
-- 1. Create a dedicated role
CREATE ROLE IF NOT EXISTS <GENAI_ROLE_NAME>;
GRANT ROLE <GENAI_ROLE_NAME> TO USER <BUSINESS_USER_NAME>;
USE ROLE SYSADMIN;
-- 2. Grant structural access
GRANT USAGE ON DATABASE <DATABASE_NAME> TO ROLE <GENAI_ROLE_NAME>;
GRANT USAGE ON SCHEMA <DATABASE_NAME>.<SCHEMA_NAME> TO ROLE <GENAI_ROLE_NAME>;
-- 3. Grant compute access
GRANT USAGE ON WAREHOUSE <WAREHOUSE_NAME> TO ROLE <GENAI_ROLE_NAME>;
-- 4 .Grant access to the semantic layer
GRANT SELECT, REFERENCES ON SEMANTIC VIEW <DATABASE_NAME>.<SCHEMA_NAME>.<SEMANTIC_VIEW_NAME> TO ROLE <GENAI_ROLE_NAME>;
-- 5. Grant select privileges on underlying raw data tables
GRANT SELECT ON ALL TABLES IN SCHEMA <DATABASE_NAME>.<SCHEMA_NAME> TO ROLE <GENAI_ROLE_NAME>;

ALTER USER <BUSINESS_USER_NAME> SET DEFAULT_ROLE = '<GENAI_ROLE_NAME>';
ALTER USER <BUSINESS_USER_NAME> SET DEFAULT_WAREHOUSE = '<WAREHOUSE_NAME>';

USE DATABASE <DATABASE_NAME>;
USE SCHEMA <SCHEMA_NAME>;
USE ROLE SYSADMIN;
​
-- 1. Create the Cortex Agent over your semantic view
CREATE OR REPLACE AGENT <AGENT_NAME>
  COMMENT = 'Governed data assistant for business queries'
  PROFILE = '{"cortex_analyst": {"semantic_models": ["<DATABASE_NAME>.<SCHEMA_NAME>.<SEMANTIC_VIEW_NAME>"]}}';
​
-- 2. Provision the MCP Server exposing the Cortex Agent
CREATE OR REPLACE MCP SERVER <SERVER_NAME>
FROM SPECIFICATION $$
tools:
  - name: "run_cortex_agent"
    type: "CORTEX_AGENT_RUN"
    identifier: "<DATABASE_NAME>.<SCHEMA_NAME>.<AGENT_NAME>"
    title: "Cortex Agent Assistant"
    description: "Converts natural language questions into SQL and executes them directly within Snowflake using the caller's OAuth session permissions."
$$;

-- 1. Grant access to the Cortex Agent and the MCP server
USE ROLE SECURITYADMIN;
GRANT USAGE ON AGENT <DATABASE_NAME>.<SCHEMA_NAME>.<AGENT_NAME> TO ROLE <GENAI_ROLE_NAME>;
GRANT USAGE ON MCP SERVER <DATABASE_NAME>.<SCHEMA_NAME>.<SERVER_NAME> TO ROLE <GENAI_ROLE_NAME>;
-- 2. View the MCP URL
DESCRIBE MCP SERVER <DATABASE_NAME>.<SCHEMA_NAME>.<SERVER_NAME>;
-- Or:
DESC MCP SERVER <DATABASE_NAME>.<SCHEMA_NAME>.<SERVER_NAME>;

https://<org_name>-<account_name>.snowflakecomputing.com/api/v2/databases/<database_name>/schemas/<schema_name>/mcp-servers/<server_name>

import os
import re
from typing import Optional
from google.adk.agents.llm_agent import LlmAgent
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.tool_context import ToolContext
​
# Define the Auth ID first in your code.
# You will use this exact name when registering the agent in the Google Cloud Console.
# Gemini Enterprise maps the authorization token to this identifier in the ToolContext.
AUTH_ID = "snowflake_oauth_conn"
​
# 1. OAuth Header Extraction:
# Define a callback to dynamically extract the user's OAuth token injected by Gemini Enterprise
# and generate the Authorization headers required to authenticate with the Snowflake MCP.
def get_auth_headers(tool_context: Optional[ToolContext] = None) -> dict:
    """
    OAuth Handling Callback: Dynamically extracts the user's OAuth token injected by Gemini Enterprise and generates the headers required to authenticate with the Snowflake MCP.
    """
    token = None
    if tool_context and hasattr(tool_context, "state"):
        try:
            # Parse the state dictionary injected by Gemini Enterprise containing authorization states
            state_dict = tool_context.state.to_dict() if hasattr(tool_context.state, 'to_dict') else tool_context.state
            pattern = re.compile(fr".*{re.escape(AUTH_ID)}.*", re.IGNORECASE)
            # Find the active token matching our Auth ID (snowflake_oauth_conn)
            matching_keys = [k for k in state_dict.keys() if pattern.match(k)]
            if matching_keys:
                token = state_dict.get(matching_keys[0])
        except Exception as e:
            print(f"Warning: Failed to read ToolContext OAuth state: {e}")
    # Return authorization headers. If token is None, requests to MCP will fail.
    return {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream",
    }
​
# 2. MCP Initialization:
# Define the Snowflake MCP URL retrieved from 'DESCRIBE MCP SERVER' directly in the code:
MCP_SERVER_URL = "<YOUR_MCP_SERVER_URL>" # Replace with the URL retrieved from 'DESCRIBE MCP SERVER'
​
# Initialize the Snowflake MCP Toolset using this URL.
snowflake_mcp = McpToolset(
    connection_params=StreamableHTTPConnectionParams(
        url=MCP_SERVER_URL
    ),
    header_provider=get_auth_headers
)
​
# 3. Build the Agent:
# Construct the LlmAgent using the Gemini 2.5 Flash model and attach the Snowflake MCP Toolset.
root_agent = LlmAgent(
    model="gemini-2.5-flash",
    name="enterprise_data_analyst",
    instruction=(
        "You are an enterprise data analyst. Use the `enterprise_data_agent` tool "
        "on the Snowflake MCP server to answer business questions. Pass the user's "
        "natural language question to the tool, which translates it into SQL and "
        "executes it directly within Snowflake's governed boundary, then present "
        "the results clearly to the user."
    ),
    tools=[snowflake_mcp]
)

snowflake_agent/
├── agent.py # The agent definition and MCP connection code (from Step 5)
├── __init__.py # Package initializer exposing the agent
├── .env # Environment configurations (Google Cloud Project, etc.)
└── requirements.txt # List of project dependencies (google-adk, etc.)

# Snowflake ADK Agent package
from .agent import root_agent

GOOGLE_CLOUD_PROJECT="<YOUR_GCP_PROJECT_ID>"
GOOGLE_CLOUD_LOCATION="<LOCATION>"
GOOGLE_GENAI_USE_VERTEXAI="TRUE"

google-adk[mcp]==2.3.0
google-genai==2.10.0
pydantic-settings==2.14.2
google-cloud-aiplatform[agent_engines]>=1.50.0

# Deploy the agent to the Google Cloud Agent Runtime
agents-cli deploy --deployment-target agent_runtime

DESCRIBE SECURITY INTEGRATION <SECURITY_INTEGRATION_NAME>;

https://<org_name>-<account_name>.snowflakecomputing.com/oauth/authorize?client_id=YOUR_CLIENT_ID&redirect_uri=https%3A%2F%2Fvertexaisearch.cloud.google.com%2Foauth-redirect&scope=session%3Arole-any&response_type=code