From 2b38bea590b392cecce4754f455f1ccb1a169ed2 Mon Sep 17 00:00:00 2001
From: Dmitriy Rusov <rusov.d.s@gmail.com>
Date: Sat, 17 May 2025 09:49:19 +0200
Subject: [PATCH] dev

---
 packages/cubejs-backend-shared/src/env.ts     | 20 +++++++++++
 .../src/SnowflakeDriver.ts                    | 33 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts
index c7b0c2e3bd584..191d82009e886 100644
--- a/packages/cubejs-backend-shared/src/env.ts
+++ b/packages/cubejs-backend-shared/src/env.ts
@@ -1625,6 +1625,26 @@ const variables: Record<string, (...args: any) => any> = {
       return false;
     }
   },
+
+  snowflakeExecutionProcedure: ({
+    dataSource
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_SNOWFLAKE_EXECUTION_PROCEDURE', dataSource)
+    ]
+  ),
+
+  snowflakeInformationSchemaProcedure: ({
+    dataSource
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_SNOWFLAKE_INFORMATION_SCHEMA_PROCEDURE', dataSource)
+    ]
+  ),
   /** ****************************************************************
    * Presto Driver                                                   *
    ***************************************************************** */
diff --git a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts
index e9b4c8baef0e1..3d81ef8102ab5 100644
--- a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts
+++ b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts
@@ -163,6 +163,8 @@ interface SnowflakeDriverOptions {
   executionTimeout?: number,
   identIgnoreCase?: boolean,
   application: string,
+  snowflakeExecutionProcedure?: string,
+  snowflakeInformationSchemaProcedure?: string,
   readOnly?: boolean,
 
   /**
@@ -208,6 +210,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
       'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PATH',
       'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PASS',
       'CUBEJS_DB_SNOWFLAKE_QUOTED_IDENTIFIERS_IGNORE_CASE',
+      'CUBEJS_DB_SNOWFLAKE_EXECUTION_PROCEDURE',
+      'CUBEJS_DB_SNOWFLAKE_INFORMATION_SCHEMA_PROCEDURE',
     ];
   }
 
@@ -296,6 +300,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
       identIgnoreCase: getEnv('snowflakeQuotedIdentIgnoreCase', { dataSource }),
       exportBucketCsvEscapeSymbol: getEnv('dbExportBucketCsvEscapeSymbol', { dataSource }),
       application: 'CubeDev_Cube',
+      snowflakeExecutionProcedure: getEnv('snowflakeExecutionProcedure', { dataSource }),
+      snowflakeInformationSchemaProcedure: getEnv('snowflakeInformationSchemaProcedure', { dataSource }),
       ...config
     };
   }
@@ -879,6 +885,24 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
     values?: unknown[],
     rehydrate: boolean = true
   ): Promise<R> {
+    if (
+      this.config.snowflakeExecutionProcedure &&
+      query.toUpperCase().startsWith('SELECT')
+    ) {
+      const escapedQuery = query.replace(/'/g, "\\'");
+      const serializedBinds = `ARRAY_CONSTRUCT(${(values ?? [])
+        .map((v) =>
+          typeof v === 'string'
+            ? `'${String(v).replace(/'/g, "''")}'`
+            : v === null || v === undefined
+              ? "NULL"
+              : v
+        )
+        .join(", ")})`;
+
+      query = `CALL ${this.config.snowflakeExecutionProcedure}('${escapedQuery}', ${serializedBinds})`;
+    }
+
     return new Promise((resolve, reject) => connection.execute({
       sqlText: query,
       binds: <string[] | undefined>values,
@@ -908,6 +932,15 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
   }
 
   public informationSchemaQuery() {
+    if (this.config.snowflakeInformationSchemaProcedure) {
+      return `
+        SELECT COLUMN_NAME as "column_name", 
+               TABLE_NAME as "table_name", 
+               TABLE_SCHEMA as "table_schema", 
+               DATA_TYPE as "data_type" 
+        FROM TABLE(${this.config.snowflakeInformationSchemaProcedure}())`;
+    }
+
     return `
         SELECT COLUMNS.COLUMN_NAME as "column_name",
                COLUMNS.TABLE_NAME as "table_name",