Skip to content

Commit

Permalink
refactor: move database code
Browse files Browse the repository at this point in the history
  • Loading branch information
actions-user authored and benfdking committed Sep 5, 2024
1 parent 3dbf1f6 commit 5f34859
Show file tree
Hide file tree
Showing 5 changed files with 132 additions and 85 deletions.
59 changes: 20 additions & 39 deletions js/packages/quary-extension-bus/src/databaseBigQuery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -92,7 +91,7 @@ export async function makeBigQueryRequest<T>(
}
}

export const runBigQueryStatement = async (
const runBigQueryStatement = async (
accessToken: string,
query: string,
projectId: string,
Expand Down Expand Up @@ -165,7 +164,7 @@ export const runBigQueryStatement = async (
return Ok(out)
}

interface BigQueryOptions {
export interface BigQueryOptions {
projectId: string
datasetId: string
}
Expand All @@ -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<{
Expand Down Expand Up @@ -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,
Expand All @@ -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()
Expand Down
64 changes: 24 additions & 40 deletions js/packages/quary-extension-bus/src/databaseSnowflake.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -55,7 +54,7 @@ export async function makeSnowflakeRequest<T>(
}
}

export async function snowflakeRunStatement(
async function snowflakeRunStatement(
accessToken: string,
accountUrl: string,
database: string,
Expand Down Expand Up @@ -95,7 +94,7 @@ export async function snowflakeRunStatement(
return Ok(out)
}

interface SnowflakeBaseOptions {
export interface SnowflakeBaseOptions {
accountUrl: string
clientId: string
clientSecret: string
Expand All @@ -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>> => {
Expand All @@ -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',
Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand All @@ -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()
Expand Down Expand Up @@ -374,3 +352,9 @@ export class Snowflake extends SnowflakeBase implements ServicesDatabase {
}
}
}

export interface SnowflakeOptions extends SnowflakeBaseOptions {
database: string
schema: string
warehouse: string
}
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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.
Expand Down
39 changes: 37 additions & 2 deletions js/packages/quary-extension/src/web/servicesDatabaseBigQuery.ts
Original file line number Diff line number Diff line change
@@ -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()
}
51 changes: 49 additions & 2 deletions js/packages/quary-extension/src/web/servicesDatabaseSnowflake.ts
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit 5f34859

Please sign in to comment.