diff --git a/v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelper.java b/v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelper.java index 1f2f32ec75..46aae3d235 100644 --- a/v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelper.java +++ b/v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelper.java @@ -23,6 +23,9 @@ import java.io.IOException; import java.io.StringReader; import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -84,10 +87,24 @@ public synchronized void init(ConnectionHelperRequest connectionHelperRequest) { config.addDataSourceProperty(key, value); } HikariDataSource ds = new HikariDataSource(config); + validateMysqlHikariConnection(ds); connectionPoolMap.put(sourceConnectionUrl + "/" + shard.getUserName(), ds); } } + private void validateMysqlHikariConnection(HikariDataSource ds) { + try (Statement stmt = ds.getConnection().createStatement(); + ResultSet rs = stmt.executeQuery("SELECT @@global.read_only")) { + + if (rs.next() && rs.getBoolean(1)) { + // Node is read-only. Close connection and throw exception. + throw new SQLException("Connection rejected: MySQL Shard is in read-only mode."); + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + @Override public Connection getConnection(String connectionRequestKey) throws ConnectionException { try { diff --git a/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelperTest.java b/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelperTest.java index 70b7c0e6ad..b2d01dc368 100644 --- a/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelperTest.java +++ b/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/dbutils/connection/JdbcConnectionHelperTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.verify; @@ -31,6 +32,9 @@ import com.zaxxer.hikari.HikariConfig; import com.zaxxer.hikari.HikariDataSource; import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; import java.util.Collections; import java.util.List; import java.util.Map; @@ -112,7 +116,17 @@ public void testInitConnectionPool() { try (MockedConstruction mockedDsConstruction = mockConstruction( HikariDataSource.class, - (mock, context) -> when(mock.getConnection()).thenReturn(mock(Connection.class)))) { + (mock, context) -> { + Connection mockConnection = mock(Connection.class); + Statement mockStatement = mock(Statement.class); + ResultSet mockResultSet = mock(ResultSet.class); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getBoolean(1)).thenReturn(false); // Not read-only + when(mockStatement.executeQuery("SELECT @@global.read_only")) + .thenReturn(mockResultSet); + when(mockConnection.createStatement()).thenReturn(mockStatement); + when(mock.getConnection()).thenReturn(mockConnection); + })) { try (MockedConstruction mockedConfigConstruction = mockConstruction(HikariConfig.class)) { connectionHelper.init(mockRequest); @@ -136,4 +150,87 @@ public void testInitConnectionPool() { } } } + + @Test + public void testInit_readOnlyShard_throwsException() throws SQLException { + ConnectionHelperRequest mockRequest = mock(ConnectionHelperRequest.class); + Shard mockShard = mock(Shard.class); + when(mockShard.getHost()).thenReturn("localhost"); + when(mockShard.getPort()).thenReturn("3306"); + when(mockShard.getDbName()).thenReturn("testdb"); + when(mockShard.getUserName()).thenReturn("testuser"); + when(mockShard.getPassword()).thenReturn("testpassword"); + when(mockRequest.getDriver()).thenReturn("com.mysql.cj.jdbc.Driver"); + when(mockRequest.getMaxConnections()).thenReturn(1); + + List mockShards = Collections.singletonList(mockShard); + when(mockRequest.getShards()).thenReturn(mockShards); + + // Mock the JDBC objects to simulate a read-only database + ResultSet mockResultSet = mock(ResultSet.class); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getBoolean(1)).thenReturn(true); // Read-only is true + + Statement mockStatement = mock(Statement.class); + when(mockStatement.executeQuery("SELECT @@global.read_only")).thenReturn(mockResultSet); + + Connection mockConnection = mock(Connection.class); + when(mockConnection.createStatement()).thenReturn(mockStatement); + + try (MockedConstruction mockedDsConstruction = + mockConstruction( + HikariDataSource.class, + (mock, context) -> { + when(mock.getConnection()).thenReturn(mockConnection); + })) { + + try { + connectionHelper.init(mockRequest); + fail("Expected RuntimeException was not thrown"); + } catch (RuntimeException e) { + assertTrue(e.getCause() instanceof SQLException); + assertEquals( + "Connection rejected: MySQL Shard is in read-only mode.", e.getCause().getMessage()); + } + } + } + + @Test + public void testInit_writableShard_succeeds() throws SQLException { + ConnectionHelperRequest mockRequest = mock(ConnectionHelperRequest.class); + Shard mockShard = mock(Shard.class); + when(mockShard.getHost()).thenReturn("localhost"); + when(mockShard.getPort()).thenReturn("3306"); + when(mockShard.getDbName()).thenReturn("testdb"); + when(mockShard.getUserName()).thenReturn("testuser"); + when(mockShard.getPassword()).thenReturn("testpassword"); + + List mockShards = Collections.singletonList(mockShard); + when(mockRequest.getShards()).thenReturn(mockShards); + + // Mock the JDBC objects to simulate a writable database + ResultSet mockResultSet = mock(ResultSet.class); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getBoolean(1)).thenReturn(false); // Read-only is false + + Statement mockStatement = mock(Statement.class); + when(mockStatement.executeQuery("SELECT @@global.read_only")).thenReturn(mockResultSet); + + Connection mockConnection = mock(Connection.class); + when(mockConnection.createStatement()).thenReturn(mockStatement); + + try (MockedConstruction mockedDsConstruction = + mockConstruction( + HikariDataSource.class, + (mock, context) -> { + when(mock.getConnection()).thenReturn(mockConnection); + })) { + try (MockedConstruction mockedConfigConstruction = + mockConstruction(HikariConfig.class)) { + // No exception should be thrown + connectionHelper.init(mockRequest); + assertTrue(connectionHelper.isConnectionPoolInitialized()); + } + } + } }