-
Notifications
You must be signed in to change notification settings - Fork 1.1k
fix(Spanner): Added read-only validation for Mysql shards. This does not validate the user authority. #3658
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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."); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is it possible to add some helpful debugging data about the shard which was in read-only mode? The current error does not tell the user where the error is occurring? |
||
| } | ||
| } catch (SQLException e) { | ||
| throw new RuntimeException(e); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public Connection getConnection(String connectionRequestKey) throws ConnectionException { | ||
| try { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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<HikariDataSource> 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<HikariConfig> mockedConfigConstruction = | ||
| mockConstruction(HikariConfig.class)) { | ||
| connectionHelper.init(mockRequest); | ||
|
|
@@ -136,4 +150,87 @@ public void testInitConnectionPool() { | |
| } | ||
| } | ||
| } | ||
|
|
||
| @Test | ||
| public void testInit_readOnlyShard_throwsException() throws SQLException { | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit: Lets use the same naming convention as the existing tests. (camelCase?) |
||
| 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<Shard> 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<HikariDataSource> 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<Shard> 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<HikariDataSource> mockedDsConstruction = | ||
| mockConstruction( | ||
| HikariDataSource.class, | ||
| (mock, context) -> { | ||
| when(mock.getConnection()).thenReturn(mockConnection); | ||
| })) { | ||
| try (MockedConstruction<HikariConfig> mockedConfigConstruction = | ||
| mockConstruction(HikariConfig.class)) { | ||
| // No exception should be thrown | ||
| connectionHelper.init(mockRequest); | ||
| assertTrue(connectionHelper.isConnectionPoolInitialized()); | ||
| } | ||
| } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this class specific to MySQL? How does this work for PG and other JDBC sources?