refactor: move database code

js/packages/quary-extension-bus/src/databaseBigQuery.ts

@@ -10,13 +10,12 @@ import {
 } from '@quary/proto/quary/service/v1/connection_response'
 import { columnsValuesToQueryResult } from './shared'
 import { QueryResult } from '@quary/proto/quary/service/v1/query_result'
-import { ModifiedConnectionConfig, ServicesDatabase } from './database'
-import { DatabaseDependentSettings, SqlLanguage } from './config'
-import * as vscode from 'vscode'
 import { TableAddress } from '@quary/proto/quary/service/v1/table_address'
 import { ProjectFileSource } from '@quary/proto/quary/service/v1/project_file'
+import { SqlLanguage } from './config'
+import { ModifiedConnectionConfig, ServicesDatabase } from './database'
 
-export async function makeBigQueryRequest<T>(
+async function makeBigQueryRequest<T>(
   accessToken: string,
   url: string,
   method: 'GET' | 'POST' | 'PUT' | 'DELETE' = 'GET',
@@ -92,7 +91,7 @@ export async function makeBigQueryRequest<T>(
   }
 }
 
-export const runBigQueryStatement = async (
+const runBigQueryStatement = async (
   accessToken: string,
   query: string,
   projectId: string,
@@ -165,7 +164,7 @@ export const runBigQueryStatement = async (
   return Ok(out)
 }
 
-interface BigQueryOptions {
+export interface BigQueryOptions {
   projectId: string
   datasetId: string
 }
@@ -183,22 +182,7 @@ export abstract class BigQueryBase {
     return makeBigQueryRequest(accessToken.value, url, method, body)
   }
 
-  protected getAccessToken = async (): Promise<Result<string>> => {
-    const session = await vscode.authentication.getSession(
-      'quaryBigQuery',
-      [],
-      {
-        createIfNone: true,
-      },
-    )
-    if (!session) {
-      return Err({
-        code: ErrorCodes.INTERNAL,
-        message: 'Failed to get BigQuery session',
-      })
-    }
-    return Ok(session.accessToken)
-  }
+  abstract getAccessToken(): Promise<Result<string>>
 
   listProjects = async (): Promise<Result<BigQueryProject[]>> => {
     const response = await this.makeBigQueryRequest<{
@@ -400,18 +384,20 @@ export abstract class BigQueryBase {
   }
 }
 
-export class BigQueryOAuth extends BigQueryBase implements ServicesDatabase {
-  private readonly projectId: string
-  private readonly datasetId: string
+export abstract class BigQueryBaseWithLocation
+  extends BigQueryBase
+  implements ServicesDatabase
+{
+  readonly projectId: string
+  readonly datasetId: string
 
-  constructor(options: BigQueryOptions) {
+  protected constructor(props: BigQueryOptions) {
     super()
-    this.listColumnsRoot = this.listColumnsRoot.bind(this)
-    this.projectId = options.projectId
-    this.datasetId = options.datasetId
+    this.projectId = props.projectId
+    this.datasetId = props.datasetId
   }
 
-  returnDatabaseConfiguration: () => DatabaseDependentSettings = () => ({
+  returnDatabaseConfiguration = () => ({
     runQueriesByDefault: false,
     lookForCacheViews: true,
     importSourcesAfterOnboarding: true,
@@ -425,17 +411,12 @@ export class BigQueryOAuth extends BigQueryBase implements ServicesDatabase {
     return `${this.projectId}.${this.datasetId}`
   }
 
-  listTables = async () => {
-    return this.listTablesRoot(this.projectId, this.datasetId)
-  }
+  listTables = async () => this.listTablesRoot(this.projectId, this.datasetId)
 
-  listViews = async () => {
-    return this.listViewsRoot(this.projectId, this.datasetId)
-  }
+  listViews = async () => this.listViewsRoot(this.projectId, this.datasetId)
 
-  listColumns = async (tableName: string) => {
-    return this.listColumnsRoot(tableName, this.projectId, this.datasetId)
-  }
+  listColumns = async (tableName: string) =>
+    this.listColumnsRoot(tableName, this.projectId, this.datasetId)
 
   runStatement = async (query: string) => {
     const accessToken = await this.getAccessToken()

js/packages/quary-extension-bus/src/databaseSnowflake.ts

@@ -5,9 +5,8 @@ import { TableAddress } from '@quary/proto/quary/service/v1/table_address'
 import { ProjectFileSource } from '@quary/proto/quary/service/v1/project_file'
 import { ModifiedConnectionConfig, ServicesDatabase } from './database'
 import { DatabaseDependentSettings, SqlLanguage } from './config'
-import * as vscode from 'vscode'
 
-export async function makeSnowflakeRequest<T>(
+async function makeSnowflakeRequest<T>(
   accessToken: string,
   accountUrl: string,
   body?: object,
@@ -55,7 +54,7 @@ export async function makeSnowflakeRequest<T>(
   }
 }
 
-export async function snowflakeRunStatement(
+async function snowflakeRunStatement(
   accessToken: string,
   accountUrl: string,
   database: string,
@@ -95,7 +94,7 @@ export async function snowflakeRunStatement(
   return Ok(out)
 }
 
-interface SnowflakeBaseOptions {
+export interface SnowflakeBaseOptions {
   accountUrl: string
   clientId: string
   clientSecret: string
@@ -115,6 +114,8 @@ export abstract class SnowflakeBase {
     this.role = options.role
   }
 
+  abstract getAccessToken: () => Promise<Result<string>>
+
   protected makeSnowflakeRequest = async <T>(
     body?: object,
   ): Promise<Result<T>> => {
@@ -125,23 +126,6 @@ export abstract class SnowflakeBase {
     return makeSnowflakeRequest(accessToken.value, this.accountUrl, body)
   }
 
-  protected getAccessToken = async (): Promise<Result<string>> => {
-    const session = await vscode.authentication.getSession(
-      'quarySnowflake',
-      [this.accountUrl, this.clientId, this.clientSecret, this.role],
-      {
-        createIfNone: true,
-      },
-    )
-    if (!session) {
-      return Err({
-        code: ErrorCodes.INTERNAL,
-        message: 'Unable to authenticate with Snowflake.',
-      })
-    }
-    return Ok(session.accessToken)
-  }
-
   listDatabases = async (): Promise<Result<string[]>> => {
     const listDatabasesResponse = await this.makeSnowflakeRequest({
       statement: 'SHOW DATABASES',
@@ -300,16 +284,13 @@ export abstract class SnowflakeBase {
   }
 }
 
-interface SnowflakeOptions extends SnowflakeBaseOptions {
-  database: string
-  schema: string
-  warehouse: string
-}
-
-export class Snowflake extends SnowflakeBase implements ServicesDatabase {
-  private readonly database: string
-  private readonly schema: string
-  private readonly warehouse: string
+export abstract class SnowflakeBaseWithLocation
+  extends SnowflakeBase
+  implements ServicesDatabase
+{
+  protected database: string
+  protected schema: string
+  protected warehouse: string
 
   constructor(options: SnowflakeOptions) {
     super(options)
@@ -318,6 +299,8 @@ export class Snowflake extends SnowflakeBase implements ServicesDatabase {
     this.warehouse = options.warehouse
   }
 
+  abstract getAccessToken: () => Promise<Result<string>>
+
   returnDatabaseConfiguration: () => DatabaseDependentSettings = () => ({
     runQueriesByDefault: false,
     lookForCacheViews: true,
@@ -332,17 +315,12 @@ export class Snowflake extends SnowflakeBase implements ServicesDatabase {
     return `${this.database}.${this.schema}`
   }
 
-  listTables = async () => {
-    return this.listTablesRoot(this.database, this.schema)
-  }
+  listTables = async () => this.listTablesRoot(this.database, this.schema)
 
-  listViews = async () => {
-    return this.listViewsRoot(this.database, this.schema)
-  }
+  listViews = async () => this.listViewsRoot(this.database, this.schema)
 
-  listColumns = async (tableName: string) => {
-    return this.listColumnsRoot(tableName, this.database, this.schema)
-  }
+  listColumns = async (tableName: string) =>
+    this.listColumnsRoot(tableName, this.database, this.schema)
 
   runStatement = async (statement: string) => {
     const accessToken = await this.getAccessToken()
@@ -374,3 +352,9 @@ export class Snowflake extends SnowflakeBase implements ServicesDatabase {
     }
   }
 }
+
+export interface SnowflakeOptions extends SnowflakeBaseOptions {
+  database: string
+  schema: string
+  warehouse: string
+}

js/packages/quary-extension/src/web/createDatabaseFromConfig.ts

@@ -2,8 +2,6 @@ import * as vscode from 'vscode'
 import { Err, ErrorCodes, isErr, Ok, Result } from '@shared/result'
 import { ConnectionConfig } from '@quary/proto/quary/service/v1/connection_config'
 import { ServicesDatabase } from '@shared/database'
-import { BigQueryOAuth } from '@shared/databaseBigQuery'
-import { Snowflake } from '@shared/databaseSnowflake'
 import { InMemorySqlite, PathBasedSqlite } from './servicesDatabaseSqlite'
 import { servicesDatabaseBigQueryNode } from './servicesDatabaseBigQueryNode'
 import { ServicesDatabaseDuckDBInMemory } from './servicesDatabaseDuckDB'
@@ -21,6 +19,8 @@ import { ServicesDatabaseRedshiftNode } from './servicesDatabaseRedshiftNode'
 import { ServicesDatabasePostgresNode } from './servicesDatabasePostgresNode'
 import { ServicesDatabaseClickhouseNode } from './servicesDatabaseClickhouseNode'
 import { ServicesDatabaseDremioNode } from './servicesDatabaseDremioNode'
+import { BigQueryOAuth } from './servicesDatabaseBigQuery'
+import { Snowflake } from './servicesDatabaseSnowflake'
 
 /**
  * Creates a database instance from a given configuration.

js/packages/quary-extension/src/web/servicesDatabaseBigQuery.ts

@@ -1,6 +1,41 @@
-import { BigQueryBase } from '@shared/databaseBigQuery'
+import {
+  BigQueryBase,
+  BigQueryOptions,
+  BigQueryBaseWithLocation,
+} from '@shared/databaseBigQuery'
 import { SourcesLister } from '@shared/database'
+import { Err, ErrorCodes, Ok, Result } from '@shared/result'
+import * as vscode from 'vscode'
+
+const getAccessToken = async () => {
+  const session = await vscode.authentication.getSession('quaryBigQuery', [], {
+    createIfNone: true,
+  })
+  if (!session) {
+    return Err({
+      code: ErrorCodes.INTERNAL,
+      message: 'Failed to get BigQuery session',
+    })
+  }
+  return Ok(session.accessToken)
+}
+
+export class BigQueryOAuth extends BigQueryBaseWithLocation {
+  readonly projectId: string
+  readonly datasetId: string
+
+  constructor(options: BigQueryOptions) {
+    super(options)
+    this.projectId = options.projectId
+    this.datasetId = options.datasetId
+  }
+
+  getAccessToken = async (): Promise<Result<string>> => getAccessToken()
+}
 
 export class BigQueryOauthHeadless
   extends BigQueryBase
-  implements SourcesLister {}
+  implements SourcesLister
+{
+  getAccessToken = () => getAccessToken()
+}

js/packages/quary-extension/src/web/servicesDatabaseSnowflake.ts

@@ -1,4 +1,51 @@
-import { SnowflakeBase } from '@shared/databaseSnowflake'
+import {
+  SnowflakeOptions,
+  SnowflakeBaseWithLocation,
+  SnowflakeBase,
+} from '@shared/databaseSnowflake'
 import { SourcesLister } from '@shared/database'
+import { Err, ErrorCodes, Ok, Result } from '@shared/result'
+import * as vscode from 'vscode'
 
-export class SnowflakeHeadless extends SnowflakeBase implements SourcesLister {}
+const getAccessToken = async (
+  accountUrl: string,
+  clientId: string,
+  clientSecret: string,
+  role: string,
+): Promise<Result<string>> => {
+  const session = await vscode.authentication.getSession(
+    'quarySnowflake',
+    [accountUrl, clientId, clientSecret, role],
+    {
+      createIfNone: true,
+    },
+  )
+  if (!session) {
+    return Err({
+      code: ErrorCodes.INTERNAL,
+      message: 'Unable to authenticate with Snowflake.',
+    })
+  }
+  return Ok(session.accessToken)
+}
+
+export class Snowflake extends SnowflakeBaseWithLocation {
+  readonly database: string
+  readonly schema: string
+  readonly warehouse: string
+
+  constructor(options: SnowflakeOptions) {
+    super(options)
+    this.database = options.database
+    this.schema = options.schema
+    this.warehouse = options.warehouse
+  }
+
+  getAccessToken = async (): Promise<Result<string>> =>
+    getAccessToken(this.accountUrl, this.clientId, this.clientSecret, this.role)
+}
+
+export class SnowflakeHeadless extends SnowflakeBase implements SourcesLister {
+  getAccessToken = () =>
+    getAccessToken(this.accountUrl, this.clientId, this.clientSecret, this.role)
+}