From be447422a817d73041de0a6a30292fb6652a581f Mon Sep 17 00:00:00 2001 From: "Visal .In" Date: Mon, 28 Oct 2024 19:29:23 +0700 Subject: [PATCH] fixing bigquery cannot update with null and detect primary key --- src/connections/bigquery.ts | 188 ++++++++++++---------- src/query-builder/dialects/bigquery.ts | 2 + src/query-builder/index.ts | 36 ++++- tests/connections/connection.test.ts | 29 +++- tests/units/query-builder/postgre.test.ts | 11 +- 5 files changed, 173 insertions(+), 93 deletions(-) diff --git a/src/connections/bigquery.ts b/src/connections/bigquery.ts index 5a92a6c..d11c14a 100644 --- a/src/connections/bigquery.ts +++ b/src/connections/bigquery.ts @@ -1,7 +1,7 @@ import { QueryType } from '../query-params'; import { Query } from '../query'; import { QueryResult } from './index'; -import { Database, Table, TableColumn } from '../models/database'; +import { Database, Schema, Table, TableColumn } from '../models/database'; import { BigQueryDialect } from '../query-builder/dialects/bigquery'; import { BigQuery } from '@google-cloud/bigquery'; import { @@ -12,47 +12,35 @@ import { SqlConnection } from './sql-base'; export class BigQueryConnection extends SqlConnection { bigQuery: BigQuery; - - // Default query type to positional for BigQuery - queryType = QueryType.positional; - - // Default dialect for BigQuery dialect = new BigQueryDialect(); - /** - * Creates a new BigQuery object. - * - * @param keyFileName - Path to a .json, .pem, or .p12 key file. - * @param region - Region for your dataset - */ constructor(bigQuery: any) { super(); this.bigQuery = bigQuery; } - /** - * Performs a connect action on the current Connection object. - * In this particular use case, BigQuery has no connect - * So this is a no-op - * - * @param details - Unused in the BigQuery scenario. - * @returns Promise - */ async connect(): Promise { return Promise.resolve(); } - /** - * Performs a disconnect action on the current Connection object. - * In this particular use case, BigQuery has no disconnect - * So this is a no-op - * - * @returns Promise - */ async disconnect(): Promise { return Promise.resolve(); } + createTable( + schemaName: string | undefined, + tableName: string, + columns: TableColumn[] + ): Promise { + // BigQuery does not support PRIMARY KEY. We can remove if here + const tempColumns = structuredClone(columns); + for (const column of tempColumns) { + delete column.definition.references; + } + + return super.createTable(schemaName, tableName, tempColumns); + } + /** * Triggers a query action on the current Connection object. * @@ -86,70 +74,104 @@ export class BigQueryConnection extends SqlConnection { } } - createTable( - schemaName: string | undefined, - tableName: string, - columns: TableColumn[] - ): Promise { - // BigQuery does not support PRIMARY KEY. We can remove if here - const tempColumns = structuredClone(columns); - for (const column of tempColumns) { - delete column.definition.primaryKey; - delete column.definition.references; - } - - return super.createTable(schemaName, tableName, tempColumns); - } - public async fetchDatabaseSchema(): Promise { - const database: Database = {}; - - // Fetch all datasets - const [datasets] = await this.bigQuery.getDatasets(); - if (datasets.length === 0) { - throw new Error('No datasets found in the project.'); - } - - // Iterate over each dataset - for (const dataset of datasets) { - const datasetId = dataset.id; - if (!datasetId) continue; + const [datasetList] = await this.bigQuery.getDatasets(); + + // Construct the query to get all the table in one go + const sql = datasetList + .map((dataset) => { + const schemaPath = `${this.bigQuery.projectId}.${dataset.id}`; + + return `( + SELECT + a.table_schema, + a.table_name, + a.column_name, + a.data_type, + b.constraint_schema, + b.constraint_name, + c.constraint_type + FROM \`${schemaPath}.INFORMATION_SCHEMA.COLUMNS\` AS a LEFT JOIN \`${schemaPath}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE\` AS b ON ( + a.table_schema = b.table_schema AND + a.table_name = b.table_name AND + a.column_name = b.column_name + ) LEFT JOIN \`${schemaPath}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS\` AS c ON ( + b.constraint_schema = c.constraint_schema AND + b.constraint_name = c.constraint_name + ) +)`; + }) + .join(' UNION ALL '); + + const { data } = await this.query<{ + table_schema: string; + table_name: string; + column_name: string; + data_type: string; + constraint_schema: string; + constraint_name: string; + constraint_type: null | 'PRIMARY KEY' | 'FOREIGN KEY'; + }>({ query: sql }); + + // Group the database schema by table + const database: Database = datasetList.reduce( + (acc, dataset) => { + acc[dataset.id ?? ''] = {}; + return acc; + }, + {} as Record + ); + + // Group the table by database + data.forEach((row) => { + const schema = database[row.table_schema]; + if (!schema) { + return; + } - const [tables] = await dataset.getTables(); + const table = schema[row.table_name] ?? { + name: row.table_name, + columns: [], + indexes: [], + constraints: [], + }; - if (!database[datasetId]) { - database[datasetId] = {}; // Initialize schema in the database + if (!schema[row.table_name]) { + schema[row.table_name] = table; } - for (const table of tables) { - const [metadata] = await table.getMetadata(); - - const columns = metadata.schema.fields.map( - (field: any, index: number): TableColumn => { - return { - name: field.name, - position: index, - definition: { - type: field.type, - nullable: field.mode === 'NULLABLE', - default: null, // BigQuery does not support default values in the schema metadata - primaryKey: false, // BigQuery does not have a concept of primary keys - unique: false, // BigQuery does not have a concept of unique constraints - }, - }; - } + // Add the column to the table + table.columns.push({ + name: row.column_name, + definition: { + type: row.data_type, + primaryKey: row.constraint_type === 'PRIMARY KEY', + }, + }); + + // Add the constraint to the table + if (row.constraint_name && row.constraint_type === 'PRIMARY KEY') { + let constraint = table.constraints.find( + (c) => c.name === row.constraint_name ); - const currentTable: Table = { - name: table.id ?? '', - columns: columns, - indexes: [], // BigQuery does not support indexes - constraints: [], // BigQuery does not support primary keys, foreign keys, or unique constraints - }; - - database[datasetId][table.id ?? ''] = currentTable; + if (!constraint) { + constraint = { + name: row.constraint_name, + schema: row.constraint_schema, + tableName: row.table_name, + type: row.constraint_type, + columns: [], + }; + + table.constraints.push(constraint); + } + + constraint.columns.push({ + columnName: row.column_name, + }); } - } + }); return database; } diff --git a/src/query-builder/dialects/bigquery.ts b/src/query-builder/dialects/bigquery.ts index b9ac7a0..88e0f5f 100644 --- a/src/query-builder/dialects/bigquery.ts +++ b/src/query-builder/dialects/bigquery.ts @@ -1,5 +1,7 @@ import { MySQLDialect } from './mysql'; export class BigQueryDialect extends MySQLDialect { + protected ALWAY_NO_ENFORCED_CONSTRAINT = true; + escapeId(identifier: string): string { return `\`${identifier}\``; } diff --git a/src/query-builder/index.ts b/src/query-builder/index.ts index ef45648..f0847dc 100644 --- a/src/query-builder/index.ts +++ b/src/query-builder/index.ts @@ -43,6 +43,10 @@ export abstract class AbstractDialect implements Dialect { protected AUTO_INCREMENT_KEYWORD = 'AUTO_INCREMENT'; protected SUPPORT_COLUMN_COMMENT = true; + // BigQuery does not support enforced constraint + // This flag is primary for BigQuery only. + protected ALWAY_NO_ENFORCED_CONSTRAINT = false; + escapeId(identifier: string): string { return identifier .split('.') @@ -134,6 +138,15 @@ export abstract class AbstractDialect implements Dialect { } return merged; } else { + // BigQuery does not provide easy way to bind NULL value, + // so we will skip binding NULL values and use raw NULL in query + if (where.value === null) { + return [ + `${this.escapeId(where.column)} ${where.operator} NULL`, + [], + ]; + } + return [ `${this.escapeId(where.column)} ${where.operator} ?`, [where.value], @@ -177,6 +190,10 @@ export abstract class AbstractDialect implements Dialect { const bindings: unknown[] = []; const setClauses = columns.map((column) => { + // BigQuery does not provide easy way to bind NULL value, + // so we will skip binding NULL values and use raw NULL in query + if (data[column] === null) return `${this.escapeId(column)} = NULL`; + bindings.push(data[column]); return `${this.escapeId(column)} = ?`; }); @@ -199,11 +216,20 @@ export abstract class AbstractDialect implements Dialect { const bindings: unknown[] = []; const columnNames = columns.map((column) => { - bindings.push(data[column]); + // BigQuery does not provide easy way to bind NULL value, + // so we will skip binding NULL values and use raw NULL in query + if (data[column] !== null) bindings.push(data[column]); return this.escapeId(column); }); - const placeholders = columns.map(() => '?').join(', '); + const placeholders = columns + .map((column) => { + // BigQuery does not provide easy way to bind NULL value, + // so we will skip binding NULL values and use raw NULL in query + if (data[column] === null) return 'NULL'; + return '?'; + }) + .join(', '); return [ `(${columnNames.join(', ')}) VALUES(${placeholders})`, @@ -218,6 +244,9 @@ export abstract class AbstractDialect implements Dialect { def.nullable === false ? 'NOT NULL' : '', def.invisible ? 'INVISIBLE' : '', // This is for MySQL case def.primaryKey ? 'PRIMARY KEY' : '', + def.primaryKey && this.ALWAY_NO_ENFORCED_CONSTRAINT + ? 'NOT ENFORCED' + : '', def.unique ? 'UNIQUE' : '', def.default ? `DEFAULT ${this.escapeValue(def.default)}` : '', def.defaultExpression ? `DEFAULT (${def.defaultExpression})` : '', @@ -278,7 +307,7 @@ export abstract class AbstractDialect implements Dialect { const tableName = builder.table; if (!tableName) { - throw new Error('Table name is required to build a UPDATE query.'); + throw new Error('Table name is required to build a INSERT query.'); } // Remove all empty value from object and check if there is any data to update @@ -369,6 +398,7 @@ export abstract class AbstractDialect implements Dialect { ref.match ? `MATCH ${ref.match}` : '', ref.onDelete ? `ON DELETE ${ref.onDelete}` : '', ref.onUpdate ? `ON UPDATE ${ref.onUpdate}` : '', + this.ALWAY_NO_ENFORCED_CONSTRAINT ? 'NOT ENFORCED' : '', ] .filter(Boolean) .join(' '); diff --git a/tests/connections/connection.test.ts b/tests/connections/connection.test.ts index 51894ab..8c45dc6 100644 --- a/tests/connections/connection.test.ts +++ b/tests/connections/connection.test.ts @@ -194,6 +194,24 @@ describe('Database Connection', () => { expect(fkConstraint!.referenceTableName).toBe('teams'); expect(fkConstraint!.columns[0].referenceColumnName).toBe('id'); } + + // Check the primary key + if (process.env.CONNECTION_TYPE !== 'mongodb') { + const pkList = Object.values(schemas[DEFAULT_SCHEMA]) + .map((c) => c.constraints) + .flat() + .filter((c) => c.type === 'PRIMARY KEY') + .map((constraint) => + constraint.columns.map( + (column) => + `${constraint.tableName}.${column.columnName}` + ) + ) + .flat() + .sort(); + + expect(pkList).toEqual(['persons.id', 'teams.id']); + } }); test('Select data', async () => { @@ -361,6 +379,10 @@ describe('Database Connection', () => { }); test('Rename table name', async () => { + // Skip BigQuery because you cannot rename table with + // primary key column + if (process.env.CONNECTION_TYPE === 'bigquery') return; + const { error } = await db.renameTable( DEFAULT_SCHEMA, 'persons', @@ -374,12 +396,15 @@ describe('Database Connection', () => { }); expect(cleanup(data).length).toEqual(2); + + // Revert the operation back + await db.renameTable(DEFAULT_SCHEMA, 'people', 'persons'); }); test('Delete a row', async () => { - await db.delete(DEFAULT_SCHEMA, 'people', { id: 1 }); + await db.delete(DEFAULT_SCHEMA, 'persons', { id: 1 }); - const { data } = await db.select(DEFAULT_SCHEMA, 'people', { + const { data } = await db.select(DEFAULT_SCHEMA, 'persons', { orderBy: ['id'], }); diff --git a/tests/units/query-builder/postgre.test.ts b/tests/units/query-builder/postgre.test.ts index 786af46..9f1b06b 100644 --- a/tests/units/query-builder/postgre.test.ts +++ b/tests/units/query-builder/postgre.test.ts @@ -173,12 +173,12 @@ describe('Query Builder - Postgre Dialect', () => { test('Update query without where condition', () => { const { query, parameters } = qb() - .update({ last_name: 'Visal', first_name: 'In' }) + .update({ last_name: 'Visal', banned: null, first_name: 'In' }) .into('persons') .toQuery(); expect(query).toBe( - 'UPDATE "persons" SET "last_name" = ?, "first_name" = ?' + 'UPDATE "persons" SET "last_name" = ?, "banned" = NULL, "first_name" = ?' ); expect(parameters).toEqual(['Visal', 'In']); }); @@ -191,10 +191,11 @@ describe('Query Builder - Postgre Dialect', () => { id: 123, active: 1, }) + .where('banned', 'IS', null) .toQuery(); expect(query).toBe( - 'UPDATE "persons" SET "last_name" = ?, "first_name" = ? WHERE "id" = ? AND "active" = ?' + 'UPDATE "persons" SET "last_name" = ?, "first_name" = ? WHERE "id" = ? AND "active" = ? AND "banned" IS NULL' ); expect(parameters).toEqual(['Visal', 'In', 123, 1]); }); @@ -211,12 +212,12 @@ describe('Query Builder - Postgre Dialect', () => { test('Insert data', () => { const { query, parameters } = qb() - .insert({ last_name: 'Visal', first_name: 'In' }) + .insert({ last_name: 'Visal', banned: null, first_name: 'In' }) .into('persons') .toQuery(); expect(query).toBe( - 'INSERT INTO "persons"("last_name", "first_name") VALUES(?, ?)' + 'INSERT INTO "persons"("last_name", "banned", "first_name") VALUES(?, NULL, ?)' ); expect(parameters).toEqual(['Visal', 'In']); });