Feat/additional schema level objects (Snowflake-Labs#59)

src/extension.ts

@@ -1,5 +1,5 @@
 import * as vscode from 'vscode';
-import { IExtension, IExtensionPlugin, IDriverExtensionApi } from '@sqltools/types';
+import { IExtension, IExtensionPlugin, IDriverExtensionApi, NSDatabase, MConnectionExplorer, IBaseQueries, QueryBuilder, ContextValue } from '@sqltools/types';
 import { ExtensionContext } from 'vscode';
 import { DRIVER_ALIASES } from './constants';
 const { publisher, name } = require('../package.json');
@@ -79,3 +79,45 @@ export async function activate(extContext: ExtensionContext): Promise<IDriverExt
 }
 
 export function deactivate() {}
+
+export interface ISnowflakeQueries extends IBaseQueries {
+  fetchMaterializedViews: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IMaterializedView>;
+  fetchStages: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IStage>;
+  fetchSFFunctions: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IFunction>;
+  fetchPipes: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IPipe>;
+  fetchStreams: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IStream>;
+  fetchTasks: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.ITask>;
+  fetchProcedures: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IProcedure>;
+  fetchFileFormats: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.IFileFormat>;
+  fetchSequences: QueryBuilder<NSDatabase.ISchema, SnowflakeDatabase.ISequence>;
+}
+
+export namespace SnowflakeDatabase {
+  export interface ISnowflakeConstruct {
+    name: string;
+    schema_name: string;
+    database_name: string;
+    arguments: string;
+  }
+
+  export interface IMaterializedView extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { childType: ContextValue.COLUMN }
+  export interface IFunction extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { }
+  export interface IStage extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { 
+    url: string;
+  }
+
+  export interface IPipe extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { }
+  export interface IStream extends MConnectionExplorer.IChildItem, ISnowflakeConstruct {
+    table_name: string;
+  }
+
+  export interface ITask extends MConnectionExplorer.IChildItem, ISnowflakeConstruct {
+    schedule: string;
+  }
+
+  export interface IProcedure extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { }
+  export interface IFileFormat extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { }
+  export interface ISequence extends MConnectionExplorer.IChildItem, ISnowflakeConstruct { 
+    next_value: number; 
+  }
+}

src/ls/driver.ts

@@ -3,6 +3,7 @@ import queries from './queries';
 import { IConnectionDriver, MConnectionExplorer, NSDatabase, ContextValue, Arg0 } from '@sqltools/types';
 import { v4 as generateId } from 'uuid';
 import { Snowflake } from 'snowflake-promise';
+import { SnowflakeDatabase } from '../extension';
 
 type DriverLib = any;
 type DriverOptions = any;
@@ -194,6 +195,138 @@ export default class SnowflakeDriver extends AbstractDriver<DriverLib, DriverOpt
     }));
   }
 
+  private async getMaterializedViews(parent): Promise<NSDatabase.ITable[]> {
+    const results = await this.queryResults(this.queries.fetchMaterializedViews(parent));
+    return results.map(mv => ({
+      label: mv.name,
+      database: mv.database_name,
+      schema: mv.schema_name,
+      iconId: 'table',
+      type: ContextValue.TABLE,
+      isView: true,
+    }));
+  }
+
+  private async getFunctions(parent): Promise<SnowflakeDatabase.IFunction[]> {
+    const results = await this.queryResults(this.queries.fetchSFFunctions(parent));
+    var thing = results.map(func => ({
+      ...func,
+      name: func.name,
+      label: func.name,
+      database: func.database_name,
+      schema: func.schema_name,
+      signature: func.arguments,
+      args: func.arguments
+              .substring(func.arguments.indexOf("(") + 1, func.arguments.lastIndexOf(")"))
+              .split(',')
+              .map(arg => { return arg.trim(); }),
+      resultType: func.arguments.split('RETURN ')[1],
+      detail: func.arguments.substring(func.arguments.indexOf("("), func.arguments.lastIndexOf(")") + 1),
+      type: ContextValue.FUNCTION,
+      childType: ContextValue.NO_CHILD
+    }));
+
+    return thing;
+  }
+
+  private async getStages(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.IStage[]> {
+    const results = await this.queryResults(this.queries.fetchStages(parent));
+    return results.map(stage => ({
+      ...stage,
+      label: stage.name.replace(/"/g, ''),
+      database: stage.database_name,
+      schema: stage.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'layers',
+      detail: stage.url
+    }));
+  }
+
+  private async getPipes(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.IPipe[]> {
+    const results = await this.queryResults(this.queries.fetchPipes(parent));
+    return results.map(pipe => ({
+      ...pipe,
+      label: pipe.name,
+      database: pipe.database_name,
+      schema: pipe.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'export'
+    }));
+  }
+
+  // TODO: Validate that this is correct
+  private async getStreams(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.IStream[]> {
+    const results = await this.queryResults(this.queries.fetchStreams(parent));
+    return results.map(stream => ({
+      ...stream,
+      label: stream.name,
+      database: stream.database_name,
+      schema: stream.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'debug-line-by-line',
+      detail: stream.table_name
+    }));
+  }
+
+  private async getTasks(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.ITask[]> {
+    const results = await this.queryResults(this.queries.fetchTasks(parent));
+    return results.map(task => ({
+      ...task,
+      label: task.name,
+      database: task.database_name,
+      schema: task.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'check',
+      detail: task.schedule
+    }));
+  }
+
+  private async getProcedures(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.IProcedure[]> {
+    const results = await this.queryResults(this.queries.fetchProcedures(parent));
+    return results.map(proc => ({
+      ...proc,
+      label: proc.name,
+      database: parent.database,
+      schema: proc.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'circuit-board',
+      detail: proc.arguments.substring(proc.arguments.indexOf("("), proc.arguments.lastIndexOf(")") + 1)
+    }));
+  }
+
+  private async getFileFormats(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.IFileFormat[]> {
+    const results = await this.queryResults(this.queries.fetchFileFormats(parent));
+    return results.map(fileFormat => ({
+      ...fileFormat,
+      label: fileFormat.name.replace(/"/g, ''),
+      database: fileFormat.database_name,
+      schema: fileFormat.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'file-code',
+      detail: fileFormat.type
+    }));
+  }
+
+  private async getSequences(parent: NSDatabase.ISchema): Promise<SnowflakeDatabase.ISequence[]> {
+    const results = await this.queryResults(this.queries.fetchSequences(parent));
+    return results.map(seq => ({
+      ...seq,
+      label: seq.name,
+      database: seq.database_name,
+      schema: seq.schema_name,
+      type: ContextValue.RESOURCE_GROUP,
+      childType: ContextValue.NO_CHILD,
+      iconId: 'list-ordered',
+      detail: seq.next_value.toString()
+    }));
+  }
+
   public async testConnection() {
     await this.open();
     const testSelect = await this.query('SELECT 1', {});
@@ -215,7 +348,16 @@ export default class SnowflakeDriver extends AbstractDriver<DriverLib, DriverOpt
       case ContextValue.SCHEMA:
         return <MConnectionExplorer.IChildItem[]>[
           { label: 'Tables', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.TABLE },
+          { label: 'Materialized Views', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.MATERIALIZED_VIEW },
           { label: 'Views', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.VIEW },
+          { label: 'Stages', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'Pipes', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'Streams', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'Tasks', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'Functions', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.FUNCTION },
+          { label: 'Procedures', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'File Formats', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
+          { label: 'Sequences', type: ContextValue.RESOURCE_GROUP, iconId: 'folder', childType: ContextValue.RESOURCE_GROUP },
         ];
       case ContextValue.TABLE:
       case ContextValue.VIEW:
@@ -234,7 +376,28 @@ export default class SnowflakeDriver extends AbstractDriver<DriverLib, DriverOpt
         return this.queryResults(this.queries.fetchTables(parent as NSDatabase.ISchema));
       case ContextValue.VIEW:
         return this.queryResults(this.queries.fetchViews(parent as NSDatabase.ISchema));
-    }
+      case ContextValue.MATERIALIZED_VIEW:
+        return this.getMaterializedViews(parent as NSDatabase.ISchema);
+      case ContextValue.FUNCTION:
+        return this.getFunctions(parent as NSDatabase.ISchema);
+      case ContextValue.RESOURCE_GROUP:
+        switch (item.label) {
+          case 'Stages':
+            return this.getStages(parent as NSDatabase.ISchema);
+          case 'Pipes':
+            return this.getPipes(parent as NSDatabase.ISchema);
+          case 'Streams':
+            return this.getStreams(parent as NSDatabase.ISchema);
+          case 'Tasks':
+            return this.getTasks(parent as NSDatabase.ISchema);
+          case 'Procedures':
+            return this.getProcedures(parent as NSDatabase.ISchema);
+          case 'File Formats':
+            return this.getFileFormats(parent as NSDatabase.ISchema);
+          case 'Sequences':
+            return this.getSequences(parent as NSDatabase.ISchema);
+        }
+      }
     return [];
   }
 

src/ls/queries.ts

@@ -1,5 +1,6 @@
 import { IBaseQueries, ContextValue } from '@sqltools/types';
 import queryFactory from '@sqltools/base-driver/dist/lib/factory';
+import { ISnowflakeQueries } from '../extension';
 
 /** write your queries here go fetch desired data. This queries are just examples copied from SQLite driver */
 
@@ -69,9 +70,6 @@ WHERE table_schema = '${p => p.schema}'
 ORDER BY table_name
 `;
 
-const fetchTables: IBaseQueries['fetchTables'] = fetchTablesAndViews(ContextValue.TABLE);
-const fetchViews: IBaseQueries['fetchTables'] = fetchTablesAndViews(ContextValue.VIEW , 'VIEW');
-
 const fetchSchemas: IBaseQueries['fetchSchemas'] = queryFactory`
 SELECT
   schema_name as "label",
@@ -85,8 +83,48 @@ WHERE
 ORDER BY 2
 `;
 
+const fetchTables: IBaseQueries['fetchTables'] = fetchTablesAndViews(ContextValue.TABLE);
+
+const fetchMaterializedViews: ISnowflakeQueries['fetchMaterializedViews'] = queryFactory`
+SHOW MATERIALIZED VIEWS IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchViews: IBaseQueries['fetchTables'] = fetchTablesAndViews(ContextValue.VIEW , 'VIEW');
+
+const fetchStages: ISnowflakeQueries['fetchStages'] = queryFactory`
+SHOW STAGES IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchPipes: ISnowflakeQueries['fetchPipes'] = queryFactory`
+SHOW PIPES IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchStreams: ISnowflakeQueries['fetchStreams'] = queryFactory`
+SHOW STREAMS IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchTasks: ISnowflakeQueries['fetchTasks'] = queryFactory`
+SHOW TASKS IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchSFFunctions: ISnowflakeQueries['fetchSFFunctions'] = queryFactory`
+SHOW USER FUNCTIONS IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchProcedures: ISnowflakeQueries['fetchProcedures'] = queryFactory`
+SHOW USER PROCEDURES IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchFileFormats: ISnowflakeQueries['fetchFileFormats'] = queryFactory`
+SHOW FILE FORMATS IN ${p => p.database}.${p => p.schema}
+`;
+
+const fetchSequences: ISnowflakeQueries['fetchSequences'] = queryFactory`
+SHOW SEQUENCES IN ${p => p.database}.${p => p.schema}
+`;
+
 /**
- * Intellisense code completion is not implemnted
+ * Intellisense code completion is not implemented
  * 
  * Auto completion should use SHOW TABLES <object> LIKE <pattern> IN <object_type>
  * because selecting from INFORMATION SCHEMA directly is too slow.
@@ -106,7 +144,7 @@ WHERE 1 = 0
 `;
 
 /**
- * Intellisense code completion is not implemnted
+ * Intellisense code completion is not implemented
  * 
  * Auto completion should use SHOW TABLES <object> LIKE <pattern> IN <object_type>
  * because selecting from INFORMATION SCHEMA directly is too slow.
@@ -124,15 +162,25 @@ SELECT 'no-column' as "label",
 WHERE 1 = 0
 `;
 
+
 export default {
-  describeTable,
   countRecords,
+  describeTable,
   fetchColumns,
-  fetchRecords,
   fetchDatabases,
+  fetchFileFormats,
+  fetchSFFunctions,
+  fetchMaterializedViews,
+  fetchPipes,
+  fetchProcedures,
+  fetchRecords,
+  fetchSchemas,
+  fetchSequences,
+  fetchStages,
+  fetchStreams,
   fetchTables,
+  fetchTasks,
   fetchViews,
-  fetchSchemas,
-  searchTables,
-  searchColumns
+  searchColumns,
+  searchTables
 }