JCES-1896 added compatibleWithPreviousVersion for safety rollout

atlassian-labs · Mar 10, 2021 · 09757be · 09757be
1 parent e69f2a9
commit 09757be
Show file tree

Hide file tree

Showing 7 changed files with 237 additions and 334 deletions.
diff --git a/src/main/java/com/atlassian/db/replica/api/DualConnection.java b/src/main/java/com/atlassian/db/replica/api/DualConnection.java
@@ -74,7 +74,7 @@ private DualConnection(
     @Override
     public Statement createStatement() throws SQLException {
         checkClosed();
-        return ReplicaStatement.builder(connectionProvider, consistency, databaseCall, readOnlyFunctions).build();
+        return ReplicaStatement.builder(connectionProvider, consistency, databaseCall, readOnlyFunctions, compatibleWithPreviousVersion).build();
     }
 
     @Override
@@ -85,7 +85,8 @@ public PreparedStatement prepareStatement(String sql) throws SQLException {
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).build();
     }
 
@@ -97,7 +98,8 @@ public CallableStatement prepareCall(String sql) throws SQLException {
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).build();
     }
 
@@ -201,7 +203,7 @@ public void clearWarnings() throws SQLException {
     public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException {
         checkClosed();
         return ReplicaStatement
-            .builder(connectionProvider, consistency, databaseCall, readOnlyFunctions)
+            .builder(connectionProvider, consistency, databaseCall, readOnlyFunctions, compatibleWithPreviousVersion)
             .resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .build();
@@ -219,7 +221,8 @@ public PreparedStatement prepareStatement(
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .build();
@@ -229,7 +232,7 @@ public PreparedStatement prepareStatement(
     public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
         checkClosed();
         return new ReplicaCallableStatement
-            .Builder(connectionProvider, consistency, databaseCall, sql, readOnlyFunctions)
+            .Builder(connectionProvider, consistency, databaseCall, sql, readOnlyFunctions, compatibleWithPreviousVersion)
             .resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .build();
@@ -290,7 +293,7 @@ public Statement createStatement(
         int resultSetHoldability
     ) throws SQLException {
         checkClosed();
-        return ReplicaStatement.builder(connectionProvider, consistency, databaseCall, readOnlyFunctions)
+        return ReplicaStatement.builder(connectionProvider, consistency, databaseCall, readOnlyFunctions, compatibleWithPreviousVersion)
             .resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .resultSetHoldability(resultSetHoldability)
@@ -310,7 +313,8 @@ public PreparedStatement prepareStatement(
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .resultSetHoldability(resultSetHoldability)
@@ -326,7 +330,7 @@ public CallableStatement prepareCall(
     ) throws SQLException {
         checkClosed();
         return new ReplicaCallableStatement
-            .Builder(connectionProvider, consistency, databaseCall, sql, readOnlyFunctions)
+            .Builder(connectionProvider, consistency, databaseCall, sql, readOnlyFunctions, compatibleWithPreviousVersion)
             .resultSetType(resultSetType)
             .resultSetConcurrency(resultSetConcurrency)
             .resultSetHoldability(resultSetHoldability)
@@ -341,7 +345,8 @@ public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) thr
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).autoGeneratedKeys(autoGeneratedKeys)
             .build();
     }
@@ -354,7 +359,8 @@ public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throw
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).columnIndexes(columnIndexes)
             .build();
     }
@@ -367,7 +373,8 @@ public PreparedStatement prepareStatement(String sql, String[] columnNames) thro
             consistency,
             databaseCall,
             sql,
-            readOnlyFunctions
+            readOnlyFunctions,
+            compatibleWithPreviousVersion
         ).columnNames(columnNames)
             .build();
     }

diff --git a/src/main/java/com/atlassian/db/replica/api/reason/RouteDecision.java b/src/main/java/com/atlassian/db/replica/api/reason/RouteDecision.java
@@ -39,7 +39,7 @@ public Optional<RouteDecision> getCause() {
     }
 
     /**
-     * @return true if the accompanying {@link SqlCall#call()} would fail when run on replica.
+     * @return true if the accompanying {@link com.atlassian.db.replica.api.SqlCall#call()} would fail when run on replica.
      */
     public boolean mustRunOnMain() {
         return reason.isWrite();

diff --git a/src/main/java/com/atlassian/db/replica/internal/ReplicaCallableStatement.java b/src/main/java/com/atlassian/db/replica/internal/ReplicaCallableStatement.java
@@ -38,7 +38,8 @@ public ReplicaCallableStatement(
         Integer resultSetType,
         Integer resultSetConcurrency,
         Integer resultSetHoldability,
-        Set<String> readOnlyFunctions
+        Set<String> readOnlyFunctions,
+        boolean compatibleWithPreviousVersion
     ) {
         super(
             connectionProvider,
@@ -48,8 +49,8 @@ public ReplicaCallableStatement(
             resultSetType,
             resultSetConcurrency,
             resultSetHoldability,
-            readOnlyFunctions
-        );
+            readOnlyFunctions,
+            compatibleWithPreviousVersion);
         this.sql = sql;
         this.resultSetType = resultSetType;
         this.resultSetConcurrency = resultSetConcurrency;
@@ -637,6 +638,7 @@ public static class Builder {
         private final DatabaseCall databaseCall;
         private final String sql;
         private final Set<String> readOnlyFunctions;
+        private final boolean compatibleWithPreviousVersion;
         private Integer resultSetType;
         private Integer resultSetConcurrency;
         private Integer resultSetHoldability;
@@ -646,13 +648,15 @@ public Builder(
             ReplicaConsistency consistency,
             DatabaseCall databaseCall,
             String sql,
-            Set<String> readOnlyFunctions
+            Set<String> readOnlyFunctions,
+            boolean compatibleWithPreviousVersion
         ) {
             this.connectionProvider = connectionProvider;
             this.consistency = consistency;
             this.databaseCall = databaseCall;
             this.sql = sql;
             this.readOnlyFunctions = readOnlyFunctions;
+            this.compatibleWithPreviousVersion = compatibleWithPreviousVersion;
         }
 
         public ReplicaCallableStatement.Builder resultSetType(int resultSetType) {
@@ -679,8 +683,8 @@ public ReplicaCallableStatement build() {
                 resultSetType,
                 resultSetConcurrency,
                 resultSetHoldability,
-                readOnlyFunctions
-            );
+                readOnlyFunctions,
+                compatibleWithPreviousVersion);
         }
     }
 }
diff --git a/src/main/java/com/atlassian/db/replica/internal/ReplicaPreparedStatement.java b/src/main/java/com/atlassian/db/replica/internal/ReplicaPreparedStatement.java
@@ -34,7 +34,8 @@ protected ReplicaPreparedStatement(
         Integer autoGeneratedKeys,
         String[] columnNames,
         int[] columnIndexes,
-        Set<String> readOnlyFunctions
+        Set<String> readOnlyFunctions,
+        boolean compatibleWithPreviousVersion
     ) {
         super(
             consistency,
@@ -43,8 +44,8 @@ protected ReplicaPreparedStatement(
             resultSetType,
             resultSetConcurrency,
             resultSetHoldability,
-            readOnlyFunctions
-        );
+            readOnlyFunctions,
+            compatibleWithPreviousVersion);
         this.sql = sql;
         this.resultSetType = resultSetType;
         this.resultSetConcurrency = resultSetConcurrency;
@@ -62,7 +63,8 @@ protected ReplicaPreparedStatement(
         Integer resultSetType,
         Integer resultSetConcurrency,
         Integer resultSetHoldability,
-        Set<String> readOnlyFunctions
+        Set<String> readOnlyFunctions,
+        boolean compatibleWithPreviousVersion
     ) {
         super(
             consistency,
@@ -71,8 +73,8 @@ protected ReplicaPreparedStatement(
             resultSetType,
             resultSetConcurrency,
             resultSetHoldability,
-            readOnlyFunctions
-        );
+            readOnlyFunctions,
+            compatibleWithPreviousVersion);
         this.sql = sql;
         this.resultSetType = resultSetType;
         this.resultSetConcurrency = resultSetConcurrency;
@@ -586,6 +588,7 @@ public static class Builder {
         private final DatabaseCall databaseCall;
         private final String sql;
         private final Set<String> readOnlyFunctions;
+        private final boolean compatibleWithPreviousVersion;
         private Integer resultSetType;
         private Integer resultSetConcurrency;
         private Integer resultSetHoldability;
@@ -598,13 +601,15 @@ public Builder(
             ReplicaConsistency consistency,
             DatabaseCall databaseCall,
             String sql,
-            Set<String> readOnlyFunctions
+            Set<String> readOnlyFunctions,
+            boolean compatibleWithPreviousVersion
         ) {
             this.connectionProvider = connectionProvider;
             this.consistency = consistency;
             this.databaseCall = databaseCall;
             this.sql = sql;
             this.readOnlyFunctions = readOnlyFunctions;
+            this.compatibleWithPreviousVersion = compatibleWithPreviousVersion;
         }
 
         public ReplicaPreparedStatement.Builder resultSetType(int resultSetType) {
@@ -649,8 +654,8 @@ public ReplicaPreparedStatement build() {
                 autoGeneratedKeys,
                 columnNames,
                 columnIndexes,
-                readOnlyFunctions
-            );
+                readOnlyFunctions,
+                compatibleWithPreviousVersion);
         }
     }
 }
diff --git a/src/main/java/com/atlassian/db/replica/internal/ReplicaStatement.java b/src/main/java/com/atlassian/db/replica/internal/ReplicaStatement.java
@@ -39,14 +39,17 @@ public Statement create() throws Exception {
         }
     };
 
+    private boolean compatibleWithPreviousVersion;
+
     public ReplicaStatement(
         ReplicaConsistency consistency,
         ReplicaConnectionProvider connectionProvider,
         DatabaseCall databaseCall,
         Integer resultSetType,
         Integer resultSetConcurrency,
         Integer resultSetHoldability,
-        Set<String> readOnlyFunctions
+        Set<String> readOnlyFunctions,
+        boolean compatibleWithPreviousVersion
     ) {
         this.consistency = consistency;
         this.connectionProvider = connectionProvider;
@@ -55,6 +58,7 @@ public ReplicaStatement(
         this.resultSetConcurrency = resultSetConcurrency;
         this.resultSetHoldability = resultSetHoldability;
         this.sqlFunction = new SqlFunction(readOnlyFunctions);
+        this.compatibleWithPreviousVersion = compatibleWithPreviousVersion;
     }
 
     @Override
@@ -502,8 +506,14 @@ public long executeLargeUpdate(String sql, String[] columnNames) throws SQLExcep
 
     <T> T execute(final SqlCall<T> call, final RouteDecision routeDecision) throws SQLException {
         final T result = databaseCall.call(call, routeDecision);
-        if (routeDecision.isWrite() && isWriteOperation) {
-            recordWriteAfterQueryExecution();
+        if (compatibleWithPreviousVersion) {
+            if (routeDecision.willRunOnMain() && isWriteOperation) {
+                recordWriteAfterQueryExecution();
+            }
+        } else {
+            if (routeDecision.mustRunOnMain()) {
+                recordWriteAfterQueryExecution();
+            }
         }
         return result;
     }
@@ -541,9 +551,10 @@ public static Builder builder(
         ReplicaConnectionProvider connectionProvider,
         ReplicaConsistency consistency,
         DatabaseCall databaseCall,
-        Set<String> readOnlyFunctions
+        Set<String> readOnlyFunctions,
+        boolean compatibleWithPreviousVersion
     ) {
-        return new Builder(connectionProvider, consistency, databaseCall, readOnlyFunctions);
+        return new Builder(connectionProvider, consistency, databaseCall, readOnlyFunctions, compatibleWithPreviousVersion);
     }
 
     void recordWriteAfterQueryExecution() throws SQLException {
@@ -560,9 +571,17 @@ public Statement getReadStatement(RouteDecisionBuilder decisionBuilder) {
             return prepareWriteStatement(decisionBuilder);
         }
         SqlQuery sqlQuery = new SqlQuery(decisionBuilder.getSql());
-        if (sqlQuery.isWriteOperation(sqlFunction)) {
-            decisionBuilder.reason(WRITE_OPERATION);
-            return prepareWriteStatement(decisionBuilder);
+        if (compatibleWithPreviousVersion) {
+            isWriteOperation = sqlQuery.isWriteOperation(sqlFunction);
+            if (isWriteOperation) {
+                decisionBuilder.reason(WRITE_OPERATION);
+                return prepareWriteStatement(decisionBuilder);
+            }
+        } else {
+            if (sqlQuery.isWriteOperation(sqlFunction)) {
+                decisionBuilder.reason(WRITE_OPERATION);
+                return prepareWriteStatement(decisionBuilder);
+            }
         }
 
         if (sqlQuery.isSelectForUpdate()) {
@@ -575,7 +594,9 @@ public Statement getReadStatement(RouteDecisionBuilder decisionBuilder) {
     }
 
     protected Statement getWriteStatement(RouteDecisionBuilder decisionBuilder) {
-        isWriteOperation = true;
+        if (compatibleWithPreviousVersion) {
+            isWriteOperation = true;
+        }
         return prepareWriteStatement(decisionBuilder);
     }
 
@@ -611,6 +632,7 @@ public static class Builder {
         private final ReplicaConsistency consistency;
         private final DatabaseCall databaseCall;
         private final Set<String> readOnlyFunctions;
+        private final boolean compatibleWithPreviousVersion;
         private Integer resultSetType;
         private Integer resultSetConcurrency;
         private Integer resultSetHoldability;
@@ -619,12 +641,14 @@ private Builder(
             ReplicaConnectionProvider connectionProvider,
             ReplicaConsistency consistency,
             DatabaseCall databaseCall,
-            Set<String> readOnlyFunctions
+            Set<String> readOnlyFunctions,
+            boolean compatibleWithPreviousVersion
         ) {
             this.connectionProvider = connectionProvider;
             this.consistency = consistency;
             this.databaseCall = databaseCall;
             this.readOnlyFunctions = readOnlyFunctions;
+            this.compatibleWithPreviousVersion = compatibleWithPreviousVersion;
         }
 
         public Builder resultSetType(int resultSetType) {
@@ -650,8 +674,8 @@ public ReplicaStatement build() {
                 resultSetType,
                 resultSetConcurrency,
                 resultSetHoldability,
-                readOnlyFunctions
-            );
+                readOnlyFunctions,
+                compatibleWithPreviousVersion);
         }
     }