DAT-18790: review changes - optimized DB calls from one per column to…

… one per table for column default values. updated test
liquibase · Oct 25, 2024 · 6178f64 · 6178f64
1 parent 12a4014
commit 6178f64
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 41 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,16 +1,13 @@
 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.DatabaseFunction;
-import liquibase.statement.core.RawParameterizedSqlStatement;
 import liquibase.structure.DatabaseObject;
 import liquibase.structure.core.Column;
 import liquibase.structure.core.DataType;
@@ -66,16 +63,13 @@ protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot
         if (example instanceof Column) {
             Column column = (Column) super.snapshotObject(example, snapshot);
             //These two are used too often, avoiding them? otherwise there would be too much DB calls
-            if (!column.getRelation().getName().equalsIgnoreCase("databasechangelog")
-                    && !column.getRelation().getName().equalsIgnoreCase("databasechangeloglock")) {
-                Database database = snapshot.getDatabase();
-                String query = String.format("SHOW CREATE TABLE %s.%s.%s;",
-                        column.getRelation().getSchema().getCatalog(),
-                        column.getRelation().getSchema().getName(),
-                        column.getRelation().getName());
-                String showCreateTableResponse = Scope.getCurrentScope().getSingleton(ExecutorService.class)
-                        .getExecutor("jdbc", database).queryForObject(new RawParameterizedSqlStatement(query), String.class);
-                String defaultValue = extractDefaultValue(showCreateTableResponse, column.getName());
+            String showCreateRelatedTableQuery = String.format("SHOW CREATE TABLE %s.%s.%s;",
+                    column.getRelation().getSchema().getCatalog(),
+                    column.getRelation().getSchema().getName(),
+                    column.getRelation().getName());
+            if (snapshot.getScratchData(showCreateRelatedTableQuery) != null) {
+                String showCreateTableStatement = (String) snapshot.getScratchData(showCreateRelatedTableQuery);
+                String defaultValue = extractDefaultValue(showCreateTableStatement, column.getName());
                 if (defaultValue != null) {
                     Matcher functionMatcher = FUNCTION_PATTERN.matcher(defaultValue);
                     if (functionMatcher.find()) {

diff --git a/src/main/java/liquibase/ext/databricks/snapshot/jvm/TableSnapshotGeneratorDatabricks.java b/src/main/java/liquibase/ext/databricks/snapshot/jvm/TableSnapshotGeneratorDatabricks.java
@@ -3,6 +3,7 @@
 import liquibase.Scope;
 import liquibase.database.Database;
 import liquibase.exception.DatabaseException;
+import liquibase.executor.Executor;
 import liquibase.executor.ExecutorService;
 import liquibase.ext.databricks.database.DatabricksDatabase;
 import liquibase.snapshot.DatabaseSnapshot;
@@ -45,8 +46,18 @@ protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot
         if (table != null) {
             String query = String.format("DESCRIBE TABLE EXTENDED %s.%s.%s;", database.getDefaultCatalogName(), database.getDefaultSchemaName(),
                     example.getName());
-            List<Map<String, ?>> tablePropertiesResponse = Scope.getCurrentScope().getSingleton(ExecutorService.class)
-                    .getExecutor("jdbc", database).queryForList(new RawParameterizedSqlStatement(query));
+            Executor jdbcExecutor = Scope.getCurrentScope().getSingleton(ExecutorService.class).getExecutor("jdbc", database);
+            List<Map<String, ?>> tablePropertiesResponse = jdbcExecutor.queryForList(new RawParameterizedSqlStatement(query));
+            //Skipping changelog tables default values processing
+            List<String> changelogTableNames = Arrays.asList(database.getDatabaseChangeLogLockTableName(), database.getDatabaseChangeLogTableName());
+            if(!changelogTableNames.contains(table.getName())) {
+                String showCreateTableQuery = String.format("SHOW CREATE TABLE %s.%s.%s;", table.getSchema().getCatalog(),
+                        table.getSchema().getName(), table.getName());
+                if(snapshot.getScratchData(showCreateTableQuery) == null) {
+                    String createTableStatement = jdbcExecutor.queryForObject(new RawParameterizedSqlStatement(showCreateTableQuery), String.class);
+                    snapshot.setScratchData(showCreateTableQuery, createTableStatement);
+                }
+            }
             // DESCRIBE TABLE EXTENDED returns both columns and additional information.
             // We need to make sure "Location" is not column in the table, but table location in s3
             boolean detailedInformationNode = false;

diff --git a/...est/java/liquibase/ext/databricks/snapshot/jvm/ColumnSnapshotGeneratorDatabricksTest.java b/...est/java/liquibase/ext/databricks/snapshot/jvm/ColumnSnapshotGeneratorDatabricksTest.java
@@ -28,30 +28,18 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
 class ColumnSnapshotGeneratorDatabricksTest {
 
-    @Mock
-    private DatabricksDatabase database;
     @Mock
     private JdbcDatabaseSnapshot snapshot;
-    @Mock
-    private JdbcConnection connection;
-    @Mock
-    private PreparedStatement preparedStatement;
-    @Mock
-    private ResultSet resultSet;
-    @Mock
-    private ResultSetMetaData resultSetMetaData;
     @InjectMocks
     private ColumnSnapshotGeneratorDatabricks snapshotGenerator;
     @Captor
-    private ArgumentCaptor<String> statementCaptor;
+    private ArgumentCaptor<String> queryCaptor;
     private DatabaseObject testedColumn;
-    private DatabaseObject testedTable;
 
     private static final Map<String, String> COLUMN_WITH_DEFAULT_NAMES = new HashMap<String , String>() {{
         put("intcolumn", "-101");
@@ -66,7 +54,7 @@ class ColumnSnapshotGeneratorDatabricksTest {
     private static final String TEST_CATALOG_NAME = "main";
     private static final String TEST_SCHEMA_NAME = "liquibase_harness_test_ds";
     private static final String TEST_TABLE_NAME = "tablewithdefaultvalues";
-    private static final String EXPECTED_SHOW_CREATE_STATEMENT = "SHOW CREATE TABLE main.liquibase_harness_test_ds.tablewithdefaultvalues;";
+    private static final String EXPECTED_SHOW_CREATE_QUERY = "SHOW CREATE TABLE main.liquibase_harness_test_ds.tablewithdefaultvalues;";
     private static final String SHOW_CREATE_TABLE_RESPONSE = "CREATE TABLE main.liquibase_harness_test_ds.tablewithdefaultvalues(" +
             "longcolumn BIGINT GENERATED BY DEFAULT AS IDENTITY (START WITH 1 INCREMENT BY 1), " +
             "intcolumn INT DEFAULT -101, " +
@@ -82,39 +70,30 @@ class ColumnSnapshotGeneratorDatabricksTest {
     public void setUp() throws DatabaseException, SQLException {
         snapshotGenerator = new ColumnSnapshotGeneratorDatabricks();
         testedColumn = new Column();
-        testedTable = new Table(TEST_CATALOG_NAME, TEST_SCHEMA_NAME, TEST_TABLE_NAME);
-        testedColumn.setAttribute("relation", testedTable);
-        when(snapshot.getDatabase()).thenReturn(database);
-        when(database.getConnection()).thenReturn(connection);
-        when(connection.prepareStatement(statementCaptor.capture())).thenReturn(preparedStatement);
-        when(preparedStatement.executeQuery()).thenReturn(resultSet);
-        when(resultSet.getMetaData()).thenReturn(resultSetMetaData);
-        when(resultSet.getString(anyInt())).thenReturn(SHOW_CREATE_TABLE_RESPONSE);
-        when(resultSetMetaData.getColumnCount()).thenReturn(1);
+        testedColumn.setAttribute("relation", new Table(TEST_CATALOG_NAME, TEST_SCHEMA_NAME, TEST_TABLE_NAME));
+        when(snapshot.getScratchData(queryCaptor.capture())).thenReturn(SHOW_CREATE_TABLE_RESPONSE);
     }
 
     @Test
     void snapshotObjectTest() throws DatabaseException, SQLException {
         for(Map.Entry<String, String> columnWithDefault : COLUMN_WITH_DEFAULT_NAMES.entrySet()) {
             testedColumn.setName(columnWithDefault.getKey());
             testedColumn.setAttribute("liquibase-complete", true);
-            when(resultSet.next()).thenReturn(true, false);
             DatabaseObject databaseObject = snapshotGenerator.snapshotObject(testedColumn, snapshot);
             assertTrue(databaseObject instanceof Column);
             assertNull(((Column) databaseObject).getComputed());
             assertNotNull(((Column) databaseObject).getDefaultValue());
             assertEquals(columnWithDefault.getValue(), ((Column) databaseObject).getDefaultValue());
         }
-        assertEquals(EXPECTED_SHOW_CREATE_STATEMENT, statementCaptor.getValue());
         for(Map.Entry<String, DatabaseFunction> columnWithDefaultComputed: COLUMN_WITH_DEFAULT_COMPUTED_NAMES.entrySet()) {
             testedColumn.setName(columnWithDefaultComputed.getKey());
             testedColumn.setAttribute("liquibase-complete", true);
-            when(resultSet.next()).thenReturn(true, false);
             DatabaseObject databaseObject = snapshotGenerator.snapshotObject(testedColumn, snapshot);
             assertTrue(databaseObject instanceof Column);
             assertTrue(((Column) databaseObject).getComputed());
             assertNotNull(((Column) databaseObject).getDefaultValue());
             assertEquals(columnWithDefaultComputed.getValue(), ((Column) databaseObject).getDefaultValue());
         }
+        assertEquals(EXPECTED_SHOW_CREATE_QUERY, queryCaptor.getValue());
     }
 }