Skip to content

Commit

Permalink
feat: primary and foreign keys driver queries (#8115)
Browse files Browse the repository at this point in the history
  • Loading branch information
@vasilev-alex
vasilev-alex committed Apr 23, 2024
1 parent e10df65 commit 35bb1d4
Show file tree
Hide file tree
Showing 10 changed files with 523 additions and 40 deletions.
120 changes: 113 additions & 7 deletions packages/cubejs-base-driver/src/BaseDriver.ts
View file
Expand Up @@ -40,6 +40,8 @@ import {
QueryTablesResult,
QueryColumnsResult,
TableMemoryData,
PrimaryKeysQueryResult,
ForeignKeysQueryResult,
} from './driver.interface';

const sortByKeys = (unordered: any) => {
Expand Down Expand Up @@ -169,7 +171,7 @@ export abstract class BaseDriver implements DriverInterface {
const query = `
SELECT table_schema as ${this.quoteIdentifier('schema_name')},
table_name as ${this.quoteIdentifier('table_name')}
FROM information_schema.tables
FROM information_schema.tables as columns
WHERE table_schema IN (${schemasPlaceholders})
`;
return query;
Expand All @@ -181,18 +183,65 @@ export abstract class BaseDriver implements DriverInterface {
columns.table_name as ${this.quoteIdentifier('table_name')},
columns.table_schema as ${this.quoteIdentifier('schema_name')},
columns.data_type as ${this.quoteIdentifier('data_type')}
FROM information_schema.columns
FROM information_schema.columns as columns
WHERE ${conditionString}
`;

return query;
}

protected primaryKeysQuery(_?: string): string | null {
return null;
}

protected foreignKeysQuery(_?: string): string | null {
return null;
}

protected async primaryKeys(conditionString?: string, params?: string[]): Promise<PrimaryKeysQueryResult[]> {
const query = this.primaryKeysQuery(conditionString);

if (!query) {
return [];
}

try {
return (await this.query<PrimaryKeysQueryResult>(query, params));
} catch (error: any) {
if (this.logger) {
this.logger('Primary Keys Query failed. Primary Keys will be defined by heuristics', {
error: (error.stack || error).toString()
});
}
return [];
}
}

protected async foreignKeys(conditionString?: string, params?: string[]): Promise<ForeignKeysQueryResult[]> {
const query = this.foreignKeysQuery(conditionString);

if (!query) {
return [];
}

try {
return (await this.query<ForeignKeysQueryResult>(query, params));
} catch (error: any) {
if (this.logger) {
this.logger('Foreign Keys Query failed. Joins will be defined by heuristics', {
error: (error.stack || error).toString()
});
}
return [];
}
}

protected getColumnNameForSchemaName() {
return 'table_schema';
return 'columns.table_schema';
}

protected getColumnNameForTableName() {
return 'table_name';
return 'columns.table_name';
}

protected getSslOptions(dataSource: string): TLSConnectionOptions | undefined {
Expand Down Expand Up @@ -316,7 +365,11 @@ export abstract class BaseDriver implements DriverInterface {
let schema = (result[i.table_schema] || {});
const tables = (schema[i.table_name] || []);

tables.push({ name: i.column_name, type: i.data_type, attributes: i.key_type ? ['primaryKey'] : [] });
tables.push({
name: i.column_name,
type: i.data_type,
attributes: i.key_type ? ['primaryKey'] : []
});

tables.sort();
schema[i.table_name] = tables;
Expand All @@ -332,6 +385,41 @@ export abstract class BaseDriver implements DriverInterface {
return this.query(query).then(data => reduce(this.informationColumnsSchemaReducer, {}, data));
}

// Extended version of tablesSchema containing primary and foreign keys
public async tablesSchemaV2() {
const query = this.informationSchemaQuery();

const tablesSchema = await this.query(query).then(data => reduce(this.informationColumnsSchemaReducer, {}, data));
const [primaryKeys, foreignKeys] = await Promise.all([this.primaryKeys(), this.foreignKeys()]);

for (const pk of primaryKeys) {
if (Array.isArray(tablesSchema?.[pk.table_schema]?.[pk.table_name])) {
tablesSchema[pk.table_schema][pk.table_name] = tablesSchema[pk.table_schema][pk.table_name].map((it: any) => {
if (it.name === pk.column_name) {
it.attributes = ['primaryKey'];
}
return it;
});
}
}

for (const foreignKey of foreignKeys) {
if (Array.isArray(tablesSchema?.[foreignKey.table_schema]?.[foreignKey.table_name])) {
tablesSchema[foreignKey.table_schema][foreignKey.table_name] = tablesSchema[foreignKey.table_schema][foreignKey.table_name].map((it: any) => {
if (it.name === foreignKey.column_name) {
it.foreignKeys = [...(it.foreignKeys || []), {
targetTable: foreignKey.target_table,
targetColumn: foreignKey.target_column
}];
}
return it;
});
}
}

return tablesSchema;
}

public async createSchemaIfNotExists(schemaName: string): Promise<void> {
const schemas = await this.query(
`SELECT schema_name FROM information_schema.schemata WHERE schema_name = ${this.param(0)}`,
Expand All @@ -356,7 +444,7 @@ export abstract class BaseDriver implements DriverInterface {
return this.query<QueryTablesResult>(query, schemaNames);
}

public getColumnsForSpecificTables(tables: QueryTablesResult[]) {
public async getColumnsForSpecificTables(tables: QueryTablesResult[]) {
const groupedBySchema: Record<string, string[]> = {};
tables.forEach((t) => {
if (!groupedBySchema[t.schema_name]) {
Expand All @@ -381,8 +469,26 @@ export abstract class BaseDriver implements DriverInterface {
const conditionString = conditions.join(' OR ');

const query = this.getColumnsForSpecificTablesQuery(conditionString);

const [primaryKeys, foreignKeys] = await Promise.all([
this.primaryKeys(conditionString, parameters),
this.foreignKeys(conditionString, parameters)
]);

const columns = await this.query<QueryColumnsResult>(query, parameters);

for (const column of columns) {
if (primaryKeys.some(pk => pk.table_schema === column.schema_name && pk.table_name === column.table_name && pk.column_name === column.column_name)) {
column.attributes = ['primaryKey'];
}

column.foreign_keys = foreignKeys.filter(fk => fk.table_schema === column.schema_name && fk.table_name === column.table_name && fk.column_name === column.column_name).map(fk => ({
target_table: fk.target_table,
target_column: fk.target_column
}));
}

return this.query<QueryColumnsResult>(query, parameters);
return columns;
}

public getTablesQuery(schemaName: string) {
Expand Down
36 changes: 36 additions & 0 deletions packages/cubejs-base-driver/src/driver.interface.ts
View file
Expand Up @@ -9,12 +9,21 @@ export interface TableColumn {
attributes?: string[]
}

export type ForeignKey = {
// eslint-disable-next-line camelcase
target_table: string;
// eslint-disable-next-line camelcase
target_column: string;
};

export interface TableColumnQueryResult {
// eslint-disable-next-line camelcase
column_name: string;
// eslint-disable-next-line camelcase
data_type: GenericDataBaseType;
attributes?: string[]
// eslint-disable-next-line camelcase
foreign_keys?: ForeignKey[]
}

export type TableStructure = TableColumn[];
Expand Down Expand Up @@ -181,6 +190,33 @@ export type QueryTablesResult = { schema_name: string, table_name: string };
// eslint-disable-next-line camelcase
export type QueryColumnsResult = { schema_name: string, table_name: string } & TableColumnQueryResult;

export type PrimaryKeysQueryResult = {
// eslint-disable-next-line camelcase
table_schema: string
// eslint-disable-next-line camelcase
table_name: string
// eslint-disable-next-line camelcase
column_name: string
};

export type ForeignKeysQueryResult = {
// eslint-disable-next-line camelcase
table_schema: string
// eslint-disable-next-line camelcase
table_name: string
// eslint-disable-next-line camelcase
column_name: string
// eslint-disable-next-line camelcase
target_table: string
// eslint-disable-next-line camelcase
target_column: string
};

export type TableKeysFilter = {
tableSchema: string,
tableName: string[]
};

export interface DriverInterface {
createSchemaIfNotExists(schemaName: string): Promise<void>;
uploadTableWithIndexes(
Expand Down
35 changes: 35 additions & 0 deletions packages/cubejs-mysql-driver/src/MySqlDriver.ts
View file
Expand Up @@ -178,6 +178,41 @@ export class MySqlDriver extends BaseDriver implements DriverInterface {
});
}

protected primaryKeysQuery(conditionString?: string): string | null {
return `SELECT
TABLE_SCHEMA as ${this.quoteIdentifier('table_schema')},
TABLE_NAME as ${this.quoteIdentifier('table_name')},
COLUMN_NAME as ${this.quoteIdentifier('column_name')}
FROM
information_schema.KEY_COLUMN_USAGE
WHERE
CONSTRAINT_NAME = 'PRIMARY'
AND TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
${conditionString ? ` AND (${conditionString})` : ''}
ORDER BY
TABLE_SCHEMA,
TABLE_NAME,
ORDINAL_POSITION;`;
}

protected foreignKeysQuery(conditionString?: string): string | null {
return `SELECT
tc.table_schema as ${this.quoteIdentifier('table_schema')},
tc.table_name as ${this.quoteIdentifier('table_name')},
kcu.column_name as ${this.quoteIdentifier('column_name')},
columns.table_name as ${this.quoteIdentifier('target_table')},
columns.column_name as ${this.quoteIdentifier('target_column')}
FROM
information_schema.table_constraints AS tc
JOIN information_schema.key_column_usage AS kcu
ON tc.constraint_name = kcu.constraint_name
JOIN information_schema.key_column_usage AS columns
ON columns.constraint_name = tc.constraint_name
WHERE
columns.table_name NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
AND tc.constraint_type = 'FOREIGN KEY'${conditionString ? ` AND (${conditionString})` : ''};`;
}

public readOnly() {
return !!this.config.readOnly;
}
Expand Down
32 changes: 32 additions & 0 deletions packages/cubejs-postgres-driver/src/PostgresDriver.ts
View file
Expand Up @@ -142,6 +142,38 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
this.enabled = true;
}

protected primaryKeysQuery(conditionString?: string): string | null {
return `SELECT
columns.table_schema as ${this.quoteIdentifier('table_schema')},
columns.table_name as ${this.quoteIdentifier('table_name')},
columns.column_name as ${this.quoteIdentifier('column_name')}
FROM information_schema.table_constraints tc
JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_name)
JOIN information_schema.columns AS columns ON columns.table_schema = tc.constraint_schema
AND tc.table_name = columns.table_name AND ccu.column_name = columns.column_name
WHERE constraint_type = 'PRIMARY KEY' AND columns.table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA')${conditionString ? ` AND (${conditionString})` : ''}`;
}

protected foreignKeysQuery(conditionString?: string): string | null {
return `SELECT
tc.table_schema as ${this.quoteIdentifier('table_schema')},
tc.table_name as ${this.quoteIdentifier('table_name')},
kcu.column_name as ${this.quoteIdentifier('column_name')},
columns.table_name as ${this.quoteIdentifier('target_table')},
columns.column_name as ${this.quoteIdentifier('target_column')}
FROM
information_schema.table_constraints AS tc
JOIN information_schema.key_column_usage AS kcu
ON tc.constraint_name = kcu.constraint_name
JOIN information_schema.constraint_column_usage AS columns
ON columns.constraint_name = tc.constraint_name
WHERE
constraint_type = 'FOREIGN KEY'
AND ${this.getColumnNameForSchemaName()} NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys', 'INFORMATION_SCHEMA')
${conditionString ? ` AND (${conditionString})` : ''}
`;
}

/**
* The easiest way how to add additional configuration from env variables, because
* you cannot call method in RedshiftDriver.constructor before super.
Expand Down
8 changes: 8 additions & 0 deletions packages/cubejs-redshift-driver/src/RedshiftDriver.ts
View file
Expand Up @@ -72,6 +72,14 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
super(options);
}

protected primaryKeysQuery() {
return null;
}

protected foreignKeysQuery() {
return null;
}

/**
* @override
*/
Expand Down

0 comments on commit 35bb1d4

Please sign in to comment.