added defaultValue snapshotting

liquibase · Oct 4, 2024 · e7992c5 · e7992c5
1 parent 9563f17
commit e7992c5
Showing 1 changed file with 37 additions and 7 deletions.
diff --git a/src/main/java/liquibase/ext/databricks/snapshot/jvm/ColumnSnapshotGeneratorDatabricks.java b/src/main/java/liquibase/ext/databricks/snapshot/jvm/ColumnSnapshotGeneratorDatabricks.java
@@ -1,29 +1,35 @@
 package liquibase.ext.databricks.snapshot.jvm;
 
+import liquibase.Scope;
 import liquibase.database.Database;
 import liquibase.exception.DatabaseException;
+import liquibase.executor.ExecutorService;
 import liquibase.ext.databricks.database.DatabricksDatabase;
 import liquibase.snapshot.CachedRow;
+import liquibase.snapshot.DatabaseSnapshot;
 import liquibase.snapshot.SnapshotGenerator;
 import liquibase.snapshot.jvm.ColumnSnapshotGenerator;
+import liquibase.statement.core.RawParameterizedSqlStatement;
 import liquibase.structure.DatabaseObject;
 import liquibase.structure.core.Column;
 import liquibase.structure.core.DataType;
 
+import java.util.List;
+import java.util.Map;
+
 public class ColumnSnapshotGeneratorDatabricks extends ColumnSnapshotGenerator {
 
     @Override
     public int getPriority(Class<? extends DatabaseObject> objectType, Database database) {
-        if (database instanceof DatabricksDatabase) {
-            return super.getPriority(objectType, database) + PRIORITY_DATABASE;
-        } else {
-            return PRIORITY_NONE;
-        }
+        if (database instanceof DatabricksDatabase)
+            return PRIORITY_DATABASE;
+        return PRIORITY_NONE;
+
     }
 
     @Override
     public Class<? extends SnapshotGenerator>[] replaces() {
-        return new Class[] { ColumnSnapshotGenerator.class };
+        return new Class[]{ColumnSnapshotGenerator.class};
     }
 
     /**
@@ -43,4 +49,28 @@ protected DataType readDataType(CachedRow columnMetadataResultSet, Column column
         }
         return super.readDataType(columnMetadataResultSet, column, database);
     }
-}
+
+    @Override
+    protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot snapshot) throws DatabaseException {
+        //This should work after fix on Databricks side
+        if (example instanceof Column) {
+            Column column = (Column) super.snapshotObject(example, snapshot);
+            Database database = snapshot.getDatabase();
+
+            String query = String.format("SELECT column_default from %s.COLUMNS where table_name = '%s' AND table_schema='%s' AND column_name ='%s';",
+                    database.getSystemSchema(),
+                    column.getRelation().getName(),
+                    column.getRelation().getSchema().getName(),
+                    column.getName());
+            List<Map<String, ?>> tablePropertiesResponse = Scope.getCurrentScope().getSingleton(ExecutorService.class)
+                    .getExecutor("jdbc", database).queryForList(new RawParameterizedSqlStatement(query));
+            for (Map<String, ?> tableProperty : tablePropertiesResponse) {
+                column.setDefaultValue(tableProperty.get("COLUMN_DEFAULT"));
+            }
+            return column;
+        } else {
+            return example;
+        }
+    }
+
+}