feat(ai-proxy): add Snowflake integration with PAT authentication (#1573)

* feat(ai-proxy): add Snowflake integration with PAT authentication

Exposes 3 MCP tools backed by Snowflake REST API v2 (Cortex Search,
Cortex Analyst, read-only SQL execution), authenticated via Programmatic
Access Tokens. The execute-query tool enforces a defense-in-depth
read-only SQL guard on top of Snowflake role privileges.


---------

Co-authored-by: Brun Christophe <christophe.brun@forestadmin.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: christophebrun-forest

packages/ai-proxy/src/forest-integration-client.ts

@@ -5,11 +5,13 @@ import type { Logger } from '@forestadmin/datasource-toolkit';
 import { AIBadRequestError } from './errors';
 import getKolarTools, { type KolarConfig } from './integrations/kolar/tools';
 import { validateKolarConfig } from './integrations/kolar/utils';
+import getSnowflakeTools, { type SnowflakeConfig } from './integrations/snowflake/tools';
+import { validateSnowflakeConfig } from './integrations/snowflake/utils';
 import getZendeskTools, { type ZendeskConfig } from './integrations/zendesk/tools';
 import { validateZendeskConfig } from './integrations/zendesk/utils';
 
-export type CustomConfig = ZendeskConfig | KolarConfig;
-export type ForestIntegrationName = 'Zendesk' | 'Kolar';
+export type CustomConfig = ZendeskConfig | KolarConfig | SnowflakeConfig;
+export type ForestIntegrationName = 'Zendesk' | 'Kolar' | 'Snowflake';
 
 export interface ForestIntegrationConfig {
   integrationName: ForestIntegrationName;
@@ -45,6 +47,9 @@ export default class ForestIntegrationClient implements ToolProvider {
         case 'Kolar':
           tools.push(...getKolarTools(config as KolarConfig));
           break;
+        case 'Snowflake':
+          tools.push(...getSnowflakeTools(config as SnowflakeConfig));
+          break;
         default:
           this.logger?.('Warn', `Unsupported integration: ${integrationName}`);
       }
@@ -61,6 +66,8 @@ export default class ForestIntegrationClient implements ToolProvider {
             return validateZendeskConfig(config as ZendeskConfig);
           case 'Kolar':
             return validateKolarConfig(config as KolarConfig);
+          case 'Snowflake':
+            return validateSnowflakeConfig(config as SnowflakeConfig);
           default:
             throw new AIBadRequestError(`Unsupported integration: ${integrationName}`);
         }

packages/ai-proxy/src/integrations/snowflake/tools.ts

@@ -0,0 +1,33 @@
+import type RemoteTool from '../../remote-tool';
+
+import ServerRemoteTool from '../../server-remote-tool';
+import createCortexAnalystTool from './tools/cortex-analyst';
+import createCortexSearchTool from './tools/cortex-search';
+import createExecuteQueryTool from './tools/execute-query';
+import { getSnowflakeAuthHeaders, getSnowflakeBaseUrl } from './utils';
+
+export interface SnowflakeConfig {
+  accountIdentifier: string;
+  programmaticAccessToken: string;
+  defaultWarehouse?: string;
+  defaultDatabase?: string;
+  defaultSchema?: string;
+  defaultRole?: string;
+}
+
+export default function getSnowflakeTools(config: SnowflakeConfig): RemoteTool[] {
+  const headers = getSnowflakeAuthHeaders(config);
+  const baseUrl = getSnowflakeBaseUrl(config.accountIdentifier);
+
+  return [
+    createCortexSearchTool(headers, baseUrl),
+    createCortexAnalystTool(headers, baseUrl),
+    createExecuteQueryTool(headers, baseUrl, config),
+  ].map(
+    tool =>
+      new ServerRemoteTool({
+        sourceId: 'snowflake',
+        tool,
+      }),
+  );
+}

packages/ai-proxy/src/integrations/snowflake/tools/cortex-analyst.ts

@@ -0,0 +1,66 @@
+import { DynamicStructuredTool } from '@langchain/core/tools';
+import { z } from 'zod';
+
+import { AIBadRequestError } from '../../../errors';
+import { assertResponseOk } from '../utils';
+
+export default function createCortexAnalystTool(
+  headers: Record<string, string>,
+  baseUrl: string,
+): DynamicStructuredTool {
+  return new DynamicStructuredTool({
+    name: 'snowflake_cortex_analyst',
+    description:
+      'Ask a natural-language analytical question against a Snowflake semantic model or ' +
+      'semantic view using Cortex Analyst. Returns a generated SQL query and the answer. ' +
+      'Exactly one of `semantic_model_file` or `semantic_view` must be provided.',
+    schema: z.object({
+      question: z.string().min(1).describe('The natural-language analytical question'),
+      semantic_model_file: z
+        .string()
+        .optional()
+        .describe(
+          'Stage path to the YAML semantic model file ' +
+            '(e.g. "@my_db.my_schema.my_stage/model.yaml")',
+        ),
+      semantic_view: z
+        .string()
+        .optional()
+        .describe('Fully qualified semantic view name (e.g. "my_db.my_schema.my_view")'),
+    }),
+    func: async ({ question, semantic_model_file, semantic_view }) => {
+      if (!semantic_model_file && !semantic_view) {
+        throw new AIBadRequestError(
+          'Either `semantic_model_file` or `semantic_view` must be provided.',
+        );
+      }
+
+      if (semantic_model_file && semantic_view) {
+        throw new AIBadRequestError(
+          'Provide only one of `semantic_model_file` or `semantic_view`, not both.',
+        );
+      }
+
+      const body: Record<string, unknown> = {
+        messages: [
+          {
+            role: 'user',
+            content: [{ type: 'text', text: question }],
+          },
+        ],
+        ...(semantic_model_file && { semantic_model_file }),
+        ...(semantic_view && { semantic_view }),
+      };
+
+      const response = await fetch(`${baseUrl}/api/v2/cortex/analyst/message`, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body),
+      });
+
+      await assertResponseOk(response, 'cortex analyst');
+
+      return JSON.stringify(await response.json());
+    },
+  });
+}

packages/ai-proxy/src/integrations/snowflake/tools/cortex-search.ts

@@ -0,0 +1,59 @@
+import { DynamicStructuredTool } from '@langchain/core/tools';
+import { z } from 'zod';
+
+import { assertResponseOk } from '../utils';
+
+export default function createCortexSearchTool(
+  headers: Record<string, string>,
+  baseUrl: string,
+): DynamicStructuredTool {
+  return new DynamicStructuredTool({
+    name: 'snowflake_cortex_search',
+    description:
+      'Search data using a Snowflake Cortex Search service. ' +
+      'Returns semantic search results ranked by relevance. ' +
+      'The caller must provide the fully qualified service identifier (database, schema, service).',
+    schema: z.object({
+      database: z.string().min(1).describe('Name of the database containing the search service'),
+      schema: z.string().min(1).describe('Name of the schema containing the search service'),
+      service: z.string().min(1).describe('Name of the Cortex Search service'),
+      query: z.string().min(1).describe('The natural-language search query'),
+      columns: z
+        .array(z.string().min(1))
+        .optional()
+        .describe('Columns to include in the search results'),
+      filter: z
+        .record(z.string(), z.unknown())
+        .optional()
+        .describe('Optional filter expression as a JSON object'),
+      limit: z
+        .number()
+        .int()
+        .positive()
+        .max(1000)
+        .optional()
+        .describe('Maximum number of results to return (default: service-defined)'),
+    }),
+    func: async ({ database, schema, service, query, columns, filter, limit }) => {
+      const url =
+        `${baseUrl}/api/v2/databases/${encodeURIComponent(database)}` +
+        `/schemas/${encodeURIComponent(schema)}` +
+        `/cortex-search-services/${encodeURIComponent(service)}:query`;
+
+      const body: Record<string, unknown> = { query };
+      if (columns) body.columns = columns;
+      if (filter) body.filter = filter;
+      if (limit !== undefined) body.limit = limit;
+
+      const response = await fetch(url, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body),
+      });
+
+      await assertResponseOk(response, 'cortex search');
+
+      return JSON.stringify(await response.json());
+    },
+  });
+}

packages/ai-proxy/src/integrations/snowflake/tools/execute-query.ts

@@ -0,0 +1,58 @@
+import type { SnowflakeConfig } from '../tools';
+
+import { DynamicStructuredTool } from '@langchain/core/tools';
+import { z } from 'zod';
+
+import { assertReadOnlySql, assertResponseOk, buildSessionContext } from '../utils';
+
+export default function createExecuteQueryTool(
+  headers: Record<string, string>,
+  baseUrl: string,
+  config: SnowflakeConfig,
+): DynamicStructuredTool {
+  return new DynamicStructuredTool({
+    name: 'snowflake_execute_query',
+    description:
+      'Execute a read-only SQL query on Snowflake and return the results. ' +
+      'Only SELECT, SHOW, DESCRIBE, and EXPLAIN statements are allowed. ' +
+      'Multiple statements are not permitted.',
+    schema: z.object({
+      statement: z.string().min(1).describe('The read-only SQL statement to execute'),
+      warehouse: z
+        .string()
+        .optional()
+        .describe('Warehouse to use for this query (overrides the default)'),
+      database: z
+        .string()
+        .optional()
+        .describe('Database to use for this query (overrides the default)'),
+      schema: z
+        .string()
+        .optional()
+        .describe('Schema to use for this query (overrides the default)'),
+      role: z.string().optional().describe('Role to use for this query (overrides the default)'),
+    }),
+    func: async ({ statement, warehouse, database, schema, role }) => {
+      assertReadOnlySql(statement);
+
+      const body = {
+        statement,
+        ...buildSessionContext(config),
+        ...(warehouse && { warehouse }),
+        ...(database && { database }),
+        ...(schema && { schema }),
+        ...(role && { role }),
+      };
+
+      const response = await fetch(`${baseUrl}/api/v2/statements`, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body),
+      });
+
+      await assertResponseOk(response, 'execute query');
+
+      return JSON.stringify(await response.json());
+    },
+  });
+}