#80: Minimal compatibility: Return YDB SDK's YdbTransaction from YOJ'…

…s YdbRepositoryTransaction
ydb-platform · Sep 12, 2024 · 797af83 · 797af83
1 parent 8fb22de
commit 797af83
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 0 deletions.
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepositoryTransaction.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepositoryTransaction.java
@@ -1,5 +1,6 @@
 package tech.ydb.yoj.repository.ydb;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.Iterables;
 import io.grpc.Context;
@@ -8,6 +9,8 @@
 import lombok.NonNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import tech.ydb.common.transaction.TxMode;
+import tech.ydb.common.transaction.YdbTransaction;
 import tech.ydb.core.Result;
 import tech.ydb.core.Status;
 import tech.ydb.core.StatusCode;
@@ -27,6 +30,7 @@
 import tech.ydb.table.values.StructValue;
 import tech.ydb.table.values.TupleValue;
 import tech.ydb.table.values.Value;
+import tech.ydb.yoj.ExperimentalApi;
 import tech.ydb.yoj.repository.BaseDb;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.IsolationLevel;
@@ -59,10 +63,12 @@
 import tech.ydb.yoj.repository.ydb.table.YdbTable;
 import tech.ydb.yoj.util.lang.Interrupts;
 
+import javax.annotation.Nullable;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -575,4 +581,42 @@ public String toString() {
             }
         });
     }
+
+    /**
+     * @return YDB SDK {@link YdbTransaction} wrapping this {@code YdbRepositoryTransaction}
+     */
+    @ExperimentalApi(issue = "https://github.com/ydb-platform/yoj-project/issues/80")
+    public YdbTransaction toSdkTransaction() {
+        return new YdbTransaction() {
+            @Nullable
+            @Override
+            public String getId() {
+                return txId;
+            }
+
+            @Override
+            public TxMode getTxMode() {
+                return switch (options.getIsolationLevel()) {
+                    case SERIALIZABLE_READ_WRITE -> TxMode.SERIALIZABLE_RW;
+                    case ONLINE_CONSISTENT_READ_ONLY -> TxMode.ONLINE_RO;
+                    case ONLINE_INCONSISTENT_READ_ONLY -> TxMode.ONLINE_INCONSISTENT_RO;
+                    case STALE_CONSISTENT_READ_ONLY -> TxMode.STALE_RO;
+                    case SNAPSHOT -> TxMode.SNAPSHOT_RO;
+                    // TxMode.NONE corresponds to DDL statements, and we have no DDL statements in YOJ transactions
+                };
+            }
+
+            @Override
+            public String getSessionId() {
+                Preconditions.checkState(!isBadSession, "No active YDB session (tx closed by YDB side)");
+                Preconditions.checkState(session != null, "No active YDB session");
+                return session.getId();
+            }
+
+            @Override
+            public CompletableFuture<Status> getStatusFuture() {
+                throw new UnsupportedOperationException();
+            }
+        };
+    }
 }
diff --git a/...sitory-ydb-v2/src/test/java/tech/ydb/yoj/repository/ydb/YdbRepositoryIntegrationTest.java b/...sitory-ydb-v2/src/test/java/tech/ydb/yoj/repository/ydb/YdbRepositoryIntegrationTest.java
@@ -19,6 +19,7 @@
 import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
+import tech.ydb.common.transaction.TxMode;
 import tech.ydb.core.grpc.YdbHeaders;
 import tech.ydb.core.utils.Version;
 import tech.ydb.proto.OperationProtos;
@@ -96,6 +97,7 @@
 import static java.util.stream.Collectors.toList;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 import static org.junit.Assert.assertEquals;
 import static tech.ydb.yoj.repository.db.EntityExpressions.newFilterBuilder;
 import static tech.ydb.yoj.repository.db.EntityExpressions.newOrderBuilder;
@@ -963,6 +965,53 @@ public void creatingRepositoryDoesNotConnect() {
         repository.shutdown();
     }
 
+    @Test
+    public void ydbTransactionCompatibility() {
+        db.tx(() -> {
+            // No db tx or session yet!
+            var sdkTx = ((YdbRepositoryTransaction<?>) Tx.Current.get().getRepositoryTransaction()).toSdkTransaction();
+            assertThatIllegalStateException().isThrownBy(sdkTx::getSessionId);
+            assertThat(sdkTx.getId()).isNull();
+            assertThat(sdkTx.getTxMode()).isEqualTo(TxMode.SERIALIZABLE_RW);
+            assertThatExceptionOfType(UnsupportedOperationException.class).isThrownBy(sdkTx::getStatusFuture);
+
+            // Perform any read - session and tx ID appear
+            db.projects().countAll();
+            sdkTx = ((YdbRepositoryTransaction<?>) Tx.Current.get().getRepositoryTransaction()).toSdkTransaction();
+            assertThat(sdkTx.getSessionId()).isNotNull();
+            assertThat(sdkTx.getId()).isNotNull();
+            assertThat(sdkTx.getTxMode()).isEqualTo(TxMode.SERIALIZABLE_RW);
+            assertThatExceptionOfType(UnsupportedOperationException.class).isThrownBy(sdkTx::getStatusFuture);
+        });
+
+        for (var entry : Map.of(
+                IsolationLevel.ONLINE_CONSISTENT_READ_ONLY, TxMode.ONLINE_RO,
+                IsolationLevel.ONLINE_INCONSISTENT_READ_ONLY, TxMode.ONLINE_INCONSISTENT_RO,
+                IsolationLevel.STALE_CONSISTENT_READ_ONLY, TxMode.STALE_RO,
+                IsolationLevel.SNAPSHOT, TxMode.SNAPSHOT_RO
+        ).entrySet()) {
+            var isolationLevel = entry.getKey();
+            var txMode = entry.getValue();
+
+            db.readOnly().withStatementIsolationLevel(isolationLevel).run(() -> {
+                // No db tx or session yet!
+                var sdkTx = ((YdbRepositoryTransaction<?>) Tx.Current.get().getRepositoryTransaction()).toSdkTransaction();
+                assertThatIllegalStateException().isThrownBy(sdkTx::getSessionId);
+                assertThat(sdkTx.getId()).isNull();
+                assertThat(sdkTx.getTxMode()).isEqualTo(txMode);
+                assertThatExceptionOfType(UnsupportedOperationException.class).isThrownBy(sdkTx::getStatusFuture);
+
+                // Perform any read - session and tx ID appear
+                db.projects().countAll();
+                sdkTx = ((YdbRepositoryTransaction<?>) Tx.Current.get().getRepositoryTransaction()).toSdkTransaction();
+                assertThat(sdkTx.getSessionId()).isNotNull();
+                assertThat(sdkTx.getId()).isNull(); // Read transactions have no ID (that's what YDB returns, folks!)
+                assertThat(sdkTx.getTxMode()).isEqualTo(txMode);
+                assertThatExceptionOfType(UnsupportedOperationException.class).isThrownBy(sdkTx::getStatusFuture);
+            });
+        }
+    }
+
     @AllArgsConstructor
     private static class DelegateSchemeServiceImplBase extends SchemeServiceGrpc.SchemeServiceImplBase {
         @Delegate