@malloydata/db-snowflake 0.0.126-dev240305182920

package/src/snowflake_connection.ts

@@ -0,0 +1,359 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+import * as crypto from 'crypto';
+import {
+  RunSQLOptions,
+  MalloyQueryData,
+  QueryRunStats,
+  Connection,
+  PersistSQLResults,
+  StreamingConnection,
+  PooledConnection,
+  SQLBlock,
+  StructDef,
+  QueryDataRow,
+  SnowflakeDialect,
+  NamedStructDefs,
+} from '@malloydata/malloy';
+import {SnowflakeExecutor} from './snowflake_executor';
+import {
+  FetchSchemaOptions,
+  TestableConnection,
+} from '@malloydata/malloy/dist/runtime_types';
+import {ConnectionOptions} from 'snowflake-sdk';
+import {Options as PoolOptions} from 'generic-pool';
+
+type namespace = {database: string; schema: string};
+
+export interface SnowflakeConnectionOptions {
+  // snowflake sdk connection options
+  connOptions?: ConnectionOptions;
+  // generic pool options to help maintain a pool of connections to snowflake
+  poolOptions?: PoolOptions;
+
+  // the database and schema where we can perform temporary table operations.
+  // for example, if we want to create a temp table for fetching schema of an sql block
+  // we could use this database & schema instead of the main database & schema
+  scratchSpace?: namespace;
+
+  queryOptions?: RunSQLOptions;
+}
+
+export class SnowflakeConnection
+  implements
+    Connection,
+    PersistSQLResults,
+    StreamingConnection,
+    TestableConnection
+{
+  private readonly dialect = new SnowflakeDialect();
+  private executor: SnowflakeExecutor;
+  private schemaCache = new Map<
+    string,
+    | {schema: StructDef; error?: undefined; timestamp: number}
+    | {error: string; schema?: undefined; timestamp: number}
+  >();
+  private sqlSchemaCache = new Map<
+    string,
+    | {
+        structDef: StructDef;
+        error?: undefined;
+        timestamp: number;
+      }
+    | {error: string; structDef?: undefined; timestamp: number}
+  >();
+
+  // the database & schema where we do temporary operations like creating a temp table
+  private scratchSpace?: namespace;
+  private queryOptions: RunSQLOptions;
+
+  constructor(
+    public readonly name: string,
+    options?: SnowflakeConnectionOptions
+  ) {
+    let connOptions = options?.connOptions;
+    if (connOptions === undefined) {
+      // try to get connection options from ~/.snowflake/connections.toml
+      connOptions = SnowflakeExecutor.getConnectionOptionsFromToml();
+    }
+    this.executor = new SnowflakeExecutor(connOptions, options?.poolOptions);
+    this.scratchSpace = options?.scratchSpace;
+    this.queryOptions = options?.queryOptions ?? {};
+  }
+
+  get dialectName(): string {
+    return 'snowflake';
+  }
+
+  // TODO: make it support nesting soon
+  public get supportsNesting(): boolean {
+    return false;
+  }
+
+  public isPool(): this is PooledConnection {
+    return true;
+  }
+
+  public canPersist(): this is PersistSQLResults {
+    return true;
+  }
+
+  public canStream(): this is StreamingConnection {
+    return true;
+  }
+
+  public async estimateQueryCost(_sqlCommand: string): Promise<QueryRunStats> {
+    return {};
+  }
+
+  async close(): Promise<void> {
+    await this.executor.done();
+  }
+
+  private getTempTableName(sqlCommand: string): string {
+    const hash = crypto.createHash('md5').update(sqlCommand).digest('hex');
+    let tableName = `tt${hash}`;
+    if (this.scratchSpace) {
+      tableName = `${this.scratchSpace.database}.${this.scratchSpace.schema}.${tableName}`;
+    }
+    return tableName;
+  }
+
+  public async runSQL(
+    sql: string,
+    options?: RunSQLOptions
+  ): Promise<MalloyQueryData> {
+    const rowLimit = options?.rowLimit ?? this.queryOptions?.rowLimit;
+    let rows = await this.executor.batch(sql);
+    if (rowLimit !== undefined && rows.length > rowLimit) {
+      rows = rows.slice(0, rowLimit);
+    }
+    return {rows, totalRows: rows.length};
+  }
+
+  public async *runSQLStream(
+    sqlCommand: string,
+    options: RunSQLOptions = {}
+  ): AsyncIterableIterator<QueryDataRow> {
+    const streamQueryOptions = {
+      ...this.queryOptions,
+      ...options,
+    };
+
+    for await (const row of await this.executor.stream(
+      sqlCommand,
+      streamQueryOptions
+    )) {
+      yield row;
+    }
+  }
+
+  public async test(): Promise<void> {
+    await this.executor.batch('SELECT 1');
+  }
+
+  private async schemaFromQuery(
+    infoQuery: string,
+    structDef: StructDef
+  ): Promise<void> {
+    const rows = await this.executor.batch(infoQuery);
+    for (const row of rows) {
+      const snowflakeDataType = row['DATA_TYPE'] as string;
+      const s = structDef;
+      const malloyType = this.dialect.sqlTypeToMalloyType(snowflakeDataType);
+      const name = row['COLUMN_NAME'] as string;
+      if (malloyType) {
+        s.fields.push({...malloyType, name});
+      } else {
+        s.fields.push({
+          type: 'unsupported',
+          rawType: snowflakeDataType,
+          name,
+        });
+      }
+    }
+  }
+
+  private async getTableSchema(
+    tableKey: string,
+    tablePath: string
+  ): Promise<StructDef> {
+    // looks like snowflake:schemaName.tableName
+    tableKey = tableKey.toLowerCase();
+
+    let [schema, tableName] = ['', tablePath];
+    const schema_and_table = tablePath.split('.');
+    if (schema_and_table.length === 2) {
+      [schema, tableName] = schema_and_table;
+    }
+
+    const structDef: StructDef = {
+      type: 'struct',
+      dialect: 'snowflake',
+      name: tableKey,
+      structSource: {type: 'table', tablePath},
+      structRelationship: {
+        type: 'basetable',
+        connectionName: this.name,
+      },
+      fields: [],
+    };
+    // This is how we get variant information
+
+    // WITH tbl as (
+    //   SELECT * FROM malloytest.ga_sample
+    //  )
+    //  SELECT regexp_replace(PATH, '\\[.*\\]', '[]') as PATH, lower(TYPEOF(value)) as type
+    //  FROM (select object_construct(*) o from  tbl limit 100)
+    //      ,table(flatten(input => o, recursive => true)) as meta
+    //  WHERE lower(TYPEOF(value)) <> 'array'
+    //  GROUP BY 1,2
+    //  ORDER BY PATH
+
+    const infoQuery = `
+  SELECT
+    column_name, -- LOWER(COLUMN_NAME) AS column_name,
+    LOWER(DATA_TYPE) as data_type
+  FROM
+    INFORMATION_SCHEMA.COLUMNS
+  WHERE
+    table_schema = UPPER('${schema}')
+    AND table_name = UPPER('${tableName}');
+    `;
+
+    await this.schemaFromQuery(infoQuery, structDef);
+    return structDef;
+  }
+
+  public async fetchSchemaForTables(
+    missing: Record<string, string>,
+    {refreshTimestamp}: FetchSchemaOptions
+  ): Promise<{
+    schemas: Record<string, StructDef>;
+    errors: Record<string, string>;
+  }> {
+    const schemas: NamedStructDefs = {};
+    const errors: {[name: string]: string} = {};
+
+    for (const tableKey in missing) {
+      let inCache = this.schemaCache.get(tableKey);
+      if (
+        !inCache ||
+        (refreshTimestamp && refreshTimestamp > inCache.timestamp)
+      ) {
+        const tablePath = missing[tableKey];
+        const timestamp = refreshTimestamp || Date.now();
+        try {
+          inCache = {
+            schema: await this.getTableSchema(tableKey, tablePath),
+            timestamp,
+          };
+          this.schemaCache.set(tableKey, inCache);
+        } catch (error) {
+          inCache = {error: error.message, timestamp};
+        }
+      }
+      if (inCache.schema !== undefined) {
+        schemas[tableKey] = inCache.schema;
+      } else {
+        errors[tableKey] = inCache.error || 'Unknown schema fetch error';
+      }
+    }
+    return {schemas, errors};
+  }
+
+  private async getSQLBlockSchema(sqlRef: SQLBlock): Promise<StructDef> {
+    const structDef: StructDef = {
+      type: 'struct',
+      dialect: 'snowflake',
+      name: sqlRef.name,
+      structSource: {
+        type: 'sql',
+        method: 'subquery',
+        sqlBlock: sqlRef,
+      },
+      structRelationship: {
+        type: 'basetable',
+        connectionName: this.name,
+      },
+      fields: [],
+    };
+
+    // create temp table with same schema as the query
+    const tempTableName = this.getTempTableName(sqlRef.selectStr);
+    this.runSQL(
+      `
+      CREATE OR REPLACE TEMP TABLE ${tempTableName} as SELECT * FROM (
+        ${sqlRef.selectStr}
+      ) as x WHERE false;
+      `
+    );
+
+    const infoQuery = `
+  SELECT
+    column_name, -- LOWER(column_name) as column_name,
+    LOWER(data_type) as data_type
+  FROM
+    INFORMATION_SCHEMA.COLUMNS
+  WHERE
+    table_name = UPPER('${tempTableName}');
+  `;
+    await this.schemaFromQuery(infoQuery, structDef);
+    return structDef;
+  }
+
+  public async fetchSchemaForSQLBlock(
+    sqlRef: SQLBlock,
+    {refreshTimestamp}: FetchSchemaOptions
+  ): Promise<
+    | {structDef: StructDef; error?: undefined}
+    | {error: string; structDef?: undefined}
+  > {
+    const key = sqlRef.name;
+    let inCache = this.sqlSchemaCache.get(key);
+    if (
+      !inCache ||
+      (refreshTimestamp && refreshTimestamp > inCache.timestamp)
+    ) {
+      const timestamp = refreshTimestamp ?? Date.now();
+      try {
+        inCache = {
+          structDef: await this.getSQLBlockSchema(sqlRef),
+          timestamp,
+        };
+      } catch (error) {
+        inCache = {error: error.message, timestamp};
+      }
+      this.sqlSchemaCache.set(key, inCache);
+    }
+    return inCache;
+  }
+
+  public async manifestTemporaryTable(sqlCommand: string): Promise<string> {
+    const tableName = this.getTempTableName(sqlCommand);
+    const cmd = `CREATE OR REPLACE TEMP TABLE ${tableName} AS (${sqlCommand});`;
+    await this.runSQL(cmd);
+    return tableName;
+  }
+}

package/src/snowflake_executor.spec.ts

@@ -0,0 +1,108 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+import {SnowflakeExecutor} from './snowflake_executor';
+import {QueryData, RunSQLOptions} from '@malloydata/malloy';
+
+const envDatabases =
+  process.env['MALLOY_DATABASES'] || process.env['MALLOY_DATABASE'];
+
+let describe = globalThis.describe;
+if (envDatabases && !envDatabases.includes('snowflake')) {
+  describe = describe.skip;
+}
+
+class SnowflakeExecutorTestSetup {
+  private executor_: SnowflakeExecutor;
+  constructor(private executor: SnowflakeExecutor) {
+    this.executor_ = executor;
+  }
+
+  async runBatch(sqlText: string): Promise<QueryData> {
+    let ret: QueryData = [];
+    await (async () => {
+      const rows = await this.executor_.batch(sqlText);
+      return rows;
+    })().then((rows: QueryData) => {
+      ret = rows;
+    });
+    return ret;
+  }
+
+  async runStreaming(sqlText: string, queryOptions?: RunSQLOptions) {
+    const rows: QueryData = [];
+    await (async () => {
+      for await (const row of await this.executor_.stream(
+        sqlText,
+        queryOptions
+      )) {
+        rows.push(row);
+      }
+    })();
+    return rows;
+  }
+
+  async done() {
+    await this.executor_.done();
+  }
+}
+
+describe('db:SnowflakeExecutor', () => {
+  let db: SnowflakeExecutorTestSetup;
+  let query: string;
+
+  beforeAll(() => {
+    const connOptions =
+      SnowflakeExecutor.getConnectionOptionsFromEnv() ||
+      SnowflakeExecutor.getConnectionOptionsFromToml();
+    const executor = new SnowflakeExecutor(connOptions);
+    db = new SnowflakeExecutorTestSetup(executor);
+    query = `
+    select
+    *
+  from
+    (
+      values
+        (1, 'one'),
+        (2, 'two'),
+        (3, 'three'),
+        (4, 'four'),
+        (5, 'five')
+    );
+    `;
+  });
+
+  afterAll(async () => {
+    await db.done();
+  });
+
+  it('verifies batch execute', async () => {
+    const rows = await db.runBatch(query);
+    expect(rows.length).toBe(5);
+  });
+
+  it('verifies stream iterable', async () => {
+    const rows = await db.runStreaming(query, {rowLimit: 2});
+    expect(rows.length).toBe(2);
+  });
+});

package/src/snowflake_executor.ts

@@ -0,0 +1,229 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+import snowflake, {
+  SnowflakeError,
+  Statement,
+  Connection,
+  ConnectionOptions,
+} from 'snowflake-sdk';
+import {Pool, Options as PoolOptions} from 'generic-pool';
+import * as toml from 'toml';
+import * as fs from 'fs';
+import * as path from 'path';
+import {Readable} from 'stream';
+import {
+  toAsyncGenerator,
+  QueryData,
+  QueryDataRow,
+  RunSQLOptions,
+} from '@malloydata/malloy';
+
+export interface ConnectionConfigFile {
+  // a toml file with snowflake connection settings
+  // if not provided, we will try to read ~/.snowflake/config
+  config_file_path?: string;
+  // the name of connection in the config file
+  // if not provided, we will try to use the "default" connection
+  connection_name?: string;
+}
+
+// function columnNameToLowerCase(row: QueryDataRow): QueryDataRow {
+//   const ret: QueryDataRow = {};
+//   for (const key in row) {
+//     ret[key.toLowerCase()] = row[key];
+//   }
+//   return ret;
+// }
+
+export class SnowflakeExecutor {
+  private static defaultPoolOptions_: PoolOptions = {
+    min: 1,
+    max: 1,
+    // ensure we validate a connection before giving it to a client
+    testOnBorrow: true,
+    testOnReturn: true,
+  };
+  private static defaultConnectionOptions = {
+    clientSessionKeepAlive: true, // default = false
+    clientSessionKeepAliveHeartbeatFrequency: 900, // default = 3600
+  };
+
+  private pool_: Pool<Connection>;
+  constructor(connOptions: ConnectionOptions, poolOptions?: PoolOptions) {
+    this.pool_ = snowflake.createPool(connOptions, {
+      ...SnowflakeExecutor.defaultPoolOptions_,
+      ...(poolOptions ?? {}),
+    });
+  }
+
+  public static getConnectionOptionsFromEnv(): ConnectionOptions | undefined {
+    const account = process.env['SNOWFLAKE_ACCOUNT'];
+    if (account) {
+      const username = process.env['SNOWFLAKE_USER'];
+      const password = process.env['SNOWFLAKE_PASSWORD'];
+      const warehouse = process.env['SNOWFLAKE_WAREHOUSE'];
+      const database = process.env['SNOWFLAKE_DATABASE'];
+      const schema = process.env['SNOWFLAKE_SCHEMA'];
+      return {
+        account,
+        username,
+        password,
+        warehouse,
+        database,
+        schema,
+      };
+    }
+    return undefined;
+  }
+
+  public static getConnectionOptionsFromToml(
+    options?: ConnectionConfigFile
+  ): ConnectionOptions {
+    let location: string | undefined = options?.config_file_path;
+    if (location === undefined) {
+      const homeDir = process.env['HOME'] || process.env['USERPROFILE'];
+      if (homeDir === undefined) {
+        throw new Error('could not find a path to connections.toml');
+      }
+      location = path.join(homeDir, '.snowflake', 'connections.toml');
+    }
+
+    if (!fs.existsSync(location)) {
+      throw new Error(
+        `provided snowflake connection config file: ${location} does not exist`
+      );
+    }
+
+    const tomlData = fs.readFileSync(location, 'utf-8');
+    const connections = toml.parse(tomlData);
+    const tomlConnectionName = options?.connection_name ?? 'default';
+    const connection = connections[tomlConnectionName];
+    if (connection === undefined) {
+      throw new Error(
+        `provided snowflake connection name: ${tomlConnectionName} does not exist at ${location}`
+      );
+    }
+
+    // sometimes the connection file uses "user" instead of "username"
+    // because the python api expects 'user'
+    connection['username'] = connection['username'] ?? connection['user'];
+    if (!connection || !connection.account || !connection.username) {
+      throw new Error(
+        `provided snowflake connection config file at ${location} is not valid`
+      );
+    }
+
+    return {
+      // some basic options we configure by default but can be overriden
+      ...SnowflakeExecutor.defaultConnectionOptions,
+      account: connection.account,
+      username: connection.username,
+      ...connection,
+    };
+  }
+
+  public async done() {
+    await this.pool_.drain().then(() => {
+      this.pool_.clear();
+    });
+  }
+
+  public async _execute(sqlText: string, conn: Connection): Promise<QueryData> {
+    return new Promise((resolve, reject) => {
+      const _statment = conn.execute({
+        sqlText,
+        complete: (
+          err: SnowflakeError | undefined,
+          _stmt: Statement,
+          rows?: QueryData
+        ) => {
+          if (err) {
+            reject(err);
+          } else if (rows) {
+            resolve(rows);
+          }
+        },
+      });
+    });
+  }
+
+  private async _setSessionParams(conn: Connection) {
+    // set some default session parameters
+    // this is quite imporant for snowflake because malloy tends to add quotes to all database identifiers
+    // and snowflake is case sensitive by with quotes but matches against all caps identifiers without quotes
+    // await this._execute(
+    //   'ALTER SESSION SET QUOTED_IDENTIFIERS_IGNORE_CASE = true;',
+    //   conn
+    // );
+    // set utc as the default timezone which is the malloy convention
+    await this._execute("ALTER SESSION SET TIMEZONE = 'UTC';", conn);
+    // ensure week starts on Sunday which is the malloy convention
+    await this._execute('ALTER SESSION SET WEEK_START = 7;', conn);
+  }
+
+  public async batch(sqlText: string): Promise<QueryData> {
+    return await this.pool_.use(async (conn: Connection) => {
+      await this._setSessionParams(conn);
+      return await this._execute(sqlText, conn);
+    });
+  }
+
+  public async stream(
+    sqlText: string,
+    options?: RunSQLOptions
+  ): Promise<AsyncIterableIterator<QueryDataRow>> {
+    const pool: Pool<Connection> = this.pool_;
+    return await pool.acquire().then(async (conn: Connection) => {
+      await this._setSessionParams(conn);
+      const stmt: Statement = conn.execute({
+        sqlText,
+        streamResult: true,
+      });
+      const stream: Readable = stmt.streamRows();
+      function streamSnowflake(
+        onError: (error: Error) => void,
+        onData: (data: QueryDataRow) => void,
+        onEnd: () => void
+      ) {
+        function handleEnd() {
+          onEnd();
+          pool.release(conn);
+        }
+
+        let index = 0;
+        function handleData(this: Readable, row: QueryDataRow) {
+          onData(row);
+          index += 1;
+          if (options?.rowLimit !== undefined && index >= options.rowLimit) {
+            onEnd();
+          }
+        }
+        stream.on('error', onError);
+        stream.on('data', handleData);
+        stream.on('end', handleEnd);
+      }
+      return Promise.resolve(toAsyncGenerator<QueryDataRow>(streamSnowflake));
+    });
+  }
+}

package/tsconfig.json

@@ -0,0 +1,13 @@
+{
+  "extends": "../../tsconfig.packages.json",
+  "compilerOptions": {
+    "rootDir": "src",
+    "outDir": "dist",
+    "composite": true
+  },
+  "references": [
+    {
+      "path": "../malloy"
+    }
+  ]
+}