Skip to content

Commit

Permalink
Add PrimaryKey support
Browse files Browse the repository at this point in the history
  • Loading branch information
vitaliimak committed May 29, 2024
1 parent a0a49b1 commit 388a993
Show file tree
Hide file tree
Showing 7 changed files with 183 additions and 37 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
package liquibase.ext.bigquery.snapshot.jvm;

import liquibase.CatalogAndSchema;
import liquibase.Scope;
import liquibase.database.Database;
import liquibase.exception.DatabaseException;
import liquibase.executor.Executor;
import liquibase.executor.ExecutorService;
import liquibase.ext.bigquery.database.BigqueryDatabase;
import liquibase.snapshot.DatabaseSnapshot;
import liquibase.snapshot.InvalidExampleException;
import liquibase.snapshot.SnapshotGenerator;
import liquibase.snapshot.jvm.PrimaryKeySnapshotGenerator;
import liquibase.statement.core.RawParameterizedSqlStatement;
import liquibase.structure.DatabaseObject;
import liquibase.structure.core.Column;
import liquibase.structure.core.Index;
import liquibase.structure.core.PrimaryKey;
import liquibase.structure.core.Schema;
import liquibase.structure.core.Table;

import java.util.List;
import java.util.Map;
import java.util.Objects;

public class BigQueryPrimaryKeySnapshotGenerator extends PrimaryKeySnapshotGenerator {

@Override
public int getPriority(Class<? extends DatabaseObject> objectType, Database database) {
if (!(database instanceof BigqueryDatabase)) {
return PRIORITY_NONE;
}
int priority = super.getPriority(objectType, database);
if (priority > PRIORITY_NONE) {
priority += PRIORITY_DATABASE;
}
return priority;
}

@Override
public Class<? extends SnapshotGenerator>[] replaces() {
return new Class[]{PrimaryKeySnapshotGenerator.class};
}

@Override
protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot snapshot) throws DatabaseException, InvalidExampleException {
Database database = snapshot.getDatabase();
Schema schema = example.getSchema();
String searchTableName = null;
if (((PrimaryKey) example).getTable() != null) {
searchTableName = ((PrimaryKey) example).getTable().getName();
searchTableName = database.correctObjectName(searchTableName, Table.class);
}
PrimaryKey returnKey = null;

String keyColumnUsageStatement = String.format("SELECT * FROM %s.INFORMATION_SCHEMA.KEY_COLUMN_USAGE WHERE table_name = ?", schema.getSchema());
Executor executor = Scope.getCurrentScope().getSingleton(ExecutorService.class).getExecutor("jdbc", database);
List<Map<String, ?>> maps = executor.queryForList(new RawParameterizedSqlStatement(keyColumnUsageStatement, searchTableName));
String columnName;
for (Map<String, ?> map : maps) {
columnName = Objects.toString(map.get("COLUMN_NAME"), null);
int position = ((Long) map.get("ORDINAL_POSITION")).intValue();

if (returnKey == null) {
returnKey = new PrimaryKey();
String catalogName = (String) map.get("TABLE_CATALOG");
String schemaName = (String) map.get("TABLE_SCHEMA");
CatalogAndSchema tableSchema = new CatalogAndSchema(catalogName, schemaName);
returnKey.setTable((Table) new Table().setName(Objects.toString(map.get("TABLE_NAME"), null)).setSchema(new Schema(tableSchema.getCatalogName(), tableSchema.getSchemaName())));
returnKey.setName(Objects.toString(map.get("CONSTRAINT_NAME"), null));
}

returnKey.addColumn(position - 1, new Column(columnName)
.setRelation(((PrimaryKey) example).getTable()));
}
if (returnKey != null) {
Index exampleIndex = new Index().setRelation(returnKey.getTable());
exampleIndex.setColumns(returnKey.getColumns());
returnKey.setBackingIndex(exampleIndex);
}
return returnKey;
}

@Override
protected void addTo(DatabaseObject foundObject, DatabaseSnapshot snapshot) throws DatabaseException {
if (!snapshot.getSnapshotControl().shouldInclude(PrimaryKey.class)) {
return;
}

if (foundObject instanceof Table) {
Table table = (Table) foundObject;
Database database = snapshot.getDatabase();
Schema schema = table.getSchema();

Executor executor = Scope.getCurrentScope().getSingleton(ExecutorService.class).getExecutor("jdbc", database);
String tableConstraintsStatement = String.format("SELECT * FROM %s.INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE table_name = ?", schema.getSchema());
List<Map<String, ?>> maps = executor.queryForList(new RawParameterizedSqlStatement(tableConstraintsStatement, table.getName()));

for (Map<String, ?> map : maps) {
if (map.containsKey("CONSTRAINT_NAME")) {
String constraintName = Objects.toString(map.get("CONSTRAINT_NAME"), null);
PrimaryKey primaryKey = new PrimaryKey().setName(constraintName);
primaryKey.setTable((Table) foundObject);
if (!database.isSystemObject(primaryKey)) {
table.setPrimaryKey(primaryKey.setTable(table));
}
}
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,66 @@

import static liquibase.ext.bigquery.database.BigqueryDatabase.BIGQUERY_PRIORITY_DATABASE;

import liquibase.Scope;
import liquibase.database.Database;
import liquibase.datatype.DataTypeFactory;
import liquibase.datatype.DatabaseDataType;
import liquibase.ext.bigquery.database.BigqueryDatabase;
import liquibase.sqlgenerator.core.AddColumnGenerator;
import liquibase.statement.AutoIncrementConstraint;
import liquibase.statement.core.AddColumnStatement;

public class BigQueryAddColumnGenerator extends AddColumnGenerator {

@Override
protected String generateSingleColumnSQL(AddColumnStatement statement, Database database) {
return super.generateSingleColumnSQL(statement, database).replace(" ADD ", " ADD COLUMN ");
DatabaseDataType columnType = null;

if (statement.getColumnType() != null) {
columnType = DataTypeFactory.getInstance().fromDescription(statement.getColumnType() + (statement.isAutoIncrement() ? "{autoIncrement:true}" : ""), database).toDatabaseDataType(database);
}

String alterTable = " ADD COLUMN " + database.escapeColumnName(statement.getCatalogName(), statement.getSchemaName(), statement.getTableName(), statement.getColumnName());

if (columnType != null) {
alterTable += " " + columnType;
}

if (statement.isAutoIncrement() && database.supportsAutoIncrement()) {
AutoIncrementConstraint autoIncrementConstraint = statement.getAutoIncrementConstraint();
alterTable += " " + database.getAutoIncrementClause(autoIncrementConstraint.getStartWith(), autoIncrementConstraint.getIncrementBy(), autoIncrementConstraint.getGenerationType(), autoIncrementConstraint.getDefaultOnNull());
}

if (statement.getDefaultValue() != null) {
// TODO add default value
// Add field with default value to an existing table schema is not supported.
// You can achieve the same result by executing the following 3 statements:
// 1. ALTER TABLE tableName ADD COLUMN columnName;
// 2. ALTER TABLE tableName ALTER COLUMN columnName SET DEFAULT columnDefaultValue;
// 3. UPDATE tableName SET columnName = columnDefaultValue WHERE TRUE;
Scope.getCurrentScope().getLog(this.getClass()).fine("Default clauses are not supported during column creation by BigQuery");
}

if (!statement.isNullable()) {
Scope.getCurrentScope().getLog(this.getClass()).fine("Not null constraints are not supported by BigQuery");
}

if (statement.isPrimaryKey()) {
String baseSQL = generateSingleColumBaseSQL(statement, database);
alterTable += "; " + baseSQL + " ADD PRIMARY KEY ("
+ database.escapeColumnName(statement.getCatalogName(), statement.getSchemaName(), statement.getTableName(), statement.getColumnName())
+ ") NOT ENFORCED";
}

if ((statement.getAddBeforeColumn() != null) && !statement.getAddBeforeColumn().isEmpty()) {
Scope.getCurrentScope().getLog(this.getClass()).fine("Before clauses are not supported by BigQuery");
}

if ((statement.getAddAfterColumn() != null) && !statement.getAddAfterColumn().isEmpty()) {
Scope.getCurrentScope().getLog(this.getClass()).fine("After clauses are not supported by BigQuery");
}

return alterTable;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,9 @@ public boolean supports(AddPrimaryKeyStatement statement, Database database) {

@Override
public Sql[] generateSql(AddPrimaryKeyStatement statement, Database database, SqlGeneratorChain sqlGeneratorChain) {
String sql = "SELECT 1";

String sql = String.format("ALTER TABLE %s ADD PRIMARY KEY (%s) NOT ENFORCED",
database.escapeTableName(statement.getCatalogName(), statement.getSchemaName(), statement.getTableName()),
database.escapeColumnNameList(statement.getColumnNames()));
return new Sql[]{
new UnparsedSql(sql, getAffectedPrimaryKey(statement))
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import liquibase.sqlgenerator.SqlGeneratorChain;
import liquibase.sqlgenerator.core.CreateTableGenerator;
import liquibase.statement.DatabaseFunction;
import liquibase.statement.PrimaryKeyConstraint;
import liquibase.statement.core.CreateTableStatement;

import java.util.ArrayList;
Expand Down Expand Up @@ -71,6 +72,22 @@ public Sql[] generateSql(CreateTableStatement statement, Database database, SqlG
}
}

PrimaryKeyConstraint primaryKeyConstraint = statement.getPrimaryKeyConstraint();
if (primaryKeyConstraint != null) {
buffer.append(" PRIMARY KEY (");

for (int i = 0; i < primaryKeyConstraint.getColumns().size(); i++) {
String primaryKeyColumnName = primaryKeyConstraint.getColumns().get(i);
buffer.append(database.escapeColumnName(statement.getCatalogName(), statement.getSchemaName(), statement.getTableName(), primaryKeyColumnName));
if (i < primaryKeyConstraint.getColumns().size() - 1) {
buffer.append(", ");
}
}

buffer.append(")")
.append(" NOT ENFORCED")
.append(",");
}

buffer.append(",");
String sql = buffer.toString().replaceFirst(",\\s*$", "") + ")";
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@ liquibase.ext.bigquery.snapshot.jvm.BigQueryDatasetSnapshotGenerator
liquibase.ext.bigquery.snapshot.jvm.BigQueryUniqueConstraintSnapshotGenerator
liquibase.ext.bigquery.snapshot.jvm.BigQuerySequenceSnapshotGenerator
liquibase.ext.bigquery.snapshot.jvm.BigQueryViewSnapshotGenerator

liquibase.ext.bigquery.snapshot.jvm.BigQueryPrimaryKeySnapshotGenerator
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ liquibase.ext.bigquery.sqlgenerator.BigQueryAddColumnGenerator
liquibase.ext.bigquery.sqlgenerator.BigQueryAddForeignKeyConstraintGenerator
liquibase.ext.bigquery.sqlgenerator.BigQueryAddPrimaryKeyConstraintGenerator
liquibase.ext.bigquery.sqlgenerator.BigQueryDropForeignKeyConstraintGenerator
liquibase.ext.bigquery.sqlgenerator.BigQueryDropPrimaryKeyConstraintGenerator
liquibase.ext.bigquery.sqlgenerator.BigqueryCreateDatabaseChangeLogLockTableGenerator
liquibase.ext.bigquery.sqlgenerator.BigqueryCreateDatabaseChangeLogTableGenerator
liquibase.ext.bigquery.sqlgenerator.BigqueryDeleteGenerator
Expand Down

0 comments on commit 388a993

Please sign in to comment.