From dee7cdabe74058434e4d630846f066dc82fdf512 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 23 May 2024 08:14:02 +0200 Subject: [PATCH] feat: allow DML batches in transactions to execute analyzeUpdate (#3114) Executing analyzeUpdate in a DML Batch inside a transaction should be possible, as the parent transaction can be used for that. --- .../connection/AbstractBaseUnitOfWork.java | 2 -- .../AbstractMultiUseTransaction.java | 2 +- .../cloud/spanner/connection/DdlBatch.java | 2 +- .../cloud/spanner/connection/DmlBatch.java | 11 +++--- .../connection/SingleUseTransaction.java | 2 +- .../cloud/spanner/connection/UnitOfWork.java | 3 ++ .../connection/AnalyzeStatementsTest.java | 34 +++++++++++++++++-- 7 files changed, 45 insertions(+), 11 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java index ec9e7a636ed..ef624ea8416 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java @@ -176,8 +176,6 @@ private Void endUnitOfWorkSpan() { return null; } - abstract boolean isSingleUse(); - /** * Returns a descriptive name for the type of transaction / unit of work. This is used in error * messages. diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java index f84f379e79a..ca78c3e5aea 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java @@ -96,7 +96,7 @@ public String toString() { } @Override - boolean isSingleUse() { + public boolean isSingleUse() { return false; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java index ca5da153f6a..f21c7e2cdeb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java @@ -97,7 +97,7 @@ private DdlBatch(Builder builder) { } @Override - boolean isSingleUse() { + public boolean isSingleUse() { return false; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java index 1551a186e7f..fde3c609cb3 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java @@ -78,12 +78,12 @@ static Builder newBuilder() { private DmlBatch(Builder builder) { super(builder); - this.transaction = builder.transaction; + this.transaction = Preconditions.checkNotNull(builder.transaction); this.statementTag = builder.statementTag; } @Override - boolean isSingleUse() { + public boolean isSingleUse() { return false; } @@ -174,8 +174,11 @@ public ApiFuture executeUpdateAsync( @Override public ApiFuture analyzeUpdateAsync( CallType callType, ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches."); + if (transaction.isSingleUse()) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches."); + } + return transaction.analyzeUpdateAsync(callType, update, analyzeMode, options); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java index 05c05ae9e16..3eeffae47b0 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java @@ -180,7 +180,7 @@ private SingleUseTransaction(Builder builder) { } @Override - boolean isSingleUse() { + public boolean isSingleUse() { return true; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java index 1c8bc6a29c0..577cd39d201 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java @@ -77,6 +77,9 @@ public boolean isActive() { /** @return true if this unit of work is still active. */ boolean isActive(); + /** Returns true if this transaction can only be used for a single statement. */ + boolean isSingleUse(); + /** * Commits the changes in this unit of work to the database. For read-only transactions, this only * closes the {@link ReadContext}. This method will throw a {@link SpannerException} if called for diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java index 3adb50b170f..6f76963fef7 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java @@ -355,8 +355,9 @@ public void testAnalyzeUpdateStatementDdlBatch() { } @Test - public void testAnalyzeUpdateDmlBatch() { + public void testAnalyzeUpdateDmlBatch_AutoCommit() { try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.startBatchDml(); SpannerException exception = @@ -371,8 +372,23 @@ public void testAnalyzeUpdateDmlBatch() { } @Test - public void testAnalyzeUpdateStatementDmlBatch() { + public void testAnalyzeUpdateDmlBatch_Transactional() { try (Connection connection = createConnection()) { + connection.setAutocommit(false); + connection.startBatchDml(); + + assertNotNull(connection.analyzeUpdate(PLAN_UPDATE, QueryAnalyzeMode.PLAN)); + assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT)); + connection.runBatch(); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + } + + @Test + public void testAnalyzeUpdateStatementDmlBatch_AutoCommit() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.startBatchDml(); SpannerException exception = @@ -385,4 +401,18 @@ public void testAnalyzeUpdateStatementDmlBatch() { assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class)); } + + @Test + public void testAnalyzeUpdateStatementDmlBatch_Transactional() { + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + connection.startBatchDml(); + + connection.analyzeUpdateStatement(PLAN_UPDATE, QueryAnalyzeMode.PLAN); + assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT)); + connection.runBatch(); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + } }