dialect adapter changes

Author: VardhanThigle

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/mysql/MysqlDialectAdapter.java

@@ -727,6 +727,24 @@ public boolean checkForTimeout(SQLException exception) {
     return false;
   }
 
+  @Override
+  public boolean supportsApproximateCounts() {
+    return true;
+  }
+
+  @Override
+  public String getApproximateCountQuery(String tableName, ImmutableList<String> partitionColumns) {
+    return addWhereClause("EXPLAIN SELECT * FROM " + tableName, partitionColumns);
+  }
+
+  @Override
+  public long parseApproximateCount(ResultSet rs) throws SQLException {
+    if (rs.next()) {
+      return rs.getLong("rows");
+    }
+    return -1L;
+  }
+
   /**
    * Get Query that returns order of collation. The query must return all the characters in the
    * character set with the columns listed in {@link CollationsOrderQueryColumns}.

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/postgresql/PostgreSQLDialectAdapter.java

@@ -533,6 +533,30 @@ public boolean checkForTimeout(SQLException exception) {
         && TIMEOUT_SQL_STATES.contains(exception.getSQLState().toUpperCase());
   }
 
+  @Override
+  public boolean supportsApproximateCounts() {
+    return true;
+  }
+
+  @Override
+  public String getApproximateCountQuery(String tableName, ImmutableList<String> partitionColumns) {
+    return addWhereClause("EXPLAIN SELECT * FROM " + tableName, partitionColumns);
+  }
+
+  @Override
+  public long parseApproximateCount(ResultSet rs) throws SQLException {
+    if (rs.next()) {
+      String explainPlan = rs.getString(1);
+      // PostgreSQL default EXPLAIN format contains "rows=N"
+      java.util.regex.Pattern pattern = java.util.regex.Pattern.compile("rows=(\\d+)");
+      java.util.regex.Matcher matcher = pattern.matcher(explainPlan);
+      if (matcher.find()) {
+        return Long.parseLong(matcher.group(1));
+      }
+    }
+    return -1L;
+  }
+
   /**
    * Ref <a href="https://www.db-fiddle.com/f/sJyGyFpqfnoxYFpEXPxR1/0"></a> Get Query that returns
    * order of collation. The query must return all the characters in the character set with the

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/mysql/MysqlDialectAdapterTest.java

@@ -889,6 +889,40 @@ public void testExtractBoundaryDuration() throws Exception {
         .isEqualTo(java.time.Duration.ofHours(-838).minusMinutes(59).minusSeconds(59));
   }
 
+  @Test
+  public void testSupportsApproximateCounts() {
+    assertThat(new MysqlDialectAdapter(MySqlVersion.DEFAULT).supportsApproximateCounts()).isTrue();
+  }
+
+  @Test
+  public void testGetApproximateCountQuery() {
+    String testTable = "testTable";
+    ImmutableList<String> cols = ImmutableList.of("col_1", "col_2");
+    String query =
+        new MysqlDialectAdapter(MySqlVersion.DEFAULT).getApproximateCountQuery(testTable, cols);
+    assertThat(query).contains("EXPLAIN SELECT * FROM testTable");
+    assertThat(query).contains("col_1");
+    assertThat(query).contains("col_2");
+  }
+
+  @Test
+  public void testParseApproximateCount() throws SQLException {
+    ResultSet mockResultSet = mock(ResultSet.class);
+    when(mockResultSet.next()).thenReturn(true);
+    when(mockResultSet.getLong("rows")).thenReturn(1234L);
+    assertThat(new MysqlDialectAdapter(MySqlVersion.DEFAULT).parseApproximateCount(mockResultSet))
+        .isEqualTo(1234L);
+  }
+
+  @Test
+  public void testParseApproximateCount_emptyResultSet() throws SQLException {
+    ResultSet mockResultSet = mock(ResultSet.class);
+    when(mockResultSet.next()).thenReturn(false);
+    assertThat(new MysqlDialectAdapter(MySqlVersion.DEFAULT).parseApproximateCount(mockResultSet))
+        .isEqualTo(-1L);
+  }
+
+  @Test
   public void testGetCollationsOrderQuery() {
     MysqlDialectAdapter adapter = new MysqlDialectAdapter(MySqlVersion.DEFAULT);
     String dbCharset = "utf8mb4";

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/postgresql/PostgreSQLDialectAdapterTest.java

@@ -18,6 +18,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import com.google.cloud.teleport.v2.source.reader.io.exception.RetriableSchemaDiscoveryException;
@@ -597,4 +598,42 @@ public void testDiscoverTableIndexes_skipsExtraTables()
     assertThat(result).containsKey(table1);
     assertThat(result).doesNotContainKey(extraTable);
   }
+
+  @Test
+  public void testSupportsApproximateCounts() {
+    assertThat(adapter.supportsApproximateCounts()).isTrue();
+  }
+
+  @Test
+  public void testGetApproximateCountQuery() {
+    String tableName = "testTable";
+    ImmutableList<String> partitionColumns = ImmutableList.of("col1", "col2");
+    String query = adapter.getApproximateCountQuery(tableName, partitionColumns);
+    assertThat(query).contains("EXPLAIN SELECT * FROM testTable");
+    assertThat(query).contains("col1");
+    assertThat(query).contains("col2");
+  }
+
+  @Test
+  public void testParseApproximateCount() throws SQLException {
+    ResultSet mockRs = mock(ResultSet.class);
+    when(mockRs.next()).thenReturn(true);
+    when(mockRs.getString(1)).thenReturn("Seq Scan on testTable  (cost=0.00..1.01 rows=42 width=4)");
+    assertThat(adapter.parseApproximateCount(mockRs)).isEqualTo(42L);
+  }
+
+  @Test
+  public void testParseApproximateCount_emptyResultSet() throws SQLException {
+    ResultSet mockRs = mock(ResultSet.class);
+    when(mockRs.next()).thenReturn(false);
+    assertThat(adapter.parseApproximateCount(mockRs)).isEqualTo(-1L);
+  }
+
+  @Test
+  public void testParseApproximateCount_parsingFailure() throws SQLException {
+    ResultSet mockRs = mock(ResultSet.class);
+    when(mockRs.next()).thenReturn(true);
+    when(mockRs.getString(1)).thenReturn("Some unrecognized output");
+    assertThat(adapter.parseApproximateCount(mockRs)).isEqualTo(-1L);
+  }
 }