#197: WIP WIP WIP Pass Db and Tx in transaction body lambda parameters WIP WIP WIP

repository-test/src/main/java/tech/ydb/yoj/repository/test/sample/TestDb.java

@@ -1,6 +1,4 @@
 package tech.ydb.yoj.repository.test.sample;
 
-import tech.ydb.yoj.repository.db.TxManager;
-
-public interface TestDb extends TxManager, TestEntityOperations {
+public interface TestDb extends TestEntityOperations {
 }

repository-ydb-v2/src/test/java/tech/ydb/yoj/repository/ydb/list/YdbListingIntegrationTest.java

@@ -32,16 +32,16 @@ public void idPredicatesAreProperlyOrdered() {
         Complex c2 = new Complex(new Complex.Id(999_999, 15L, "UUU", Complex.Status.OK));
         Complex c3 = new Complex(new Complex.Id(999_999, 15L, "KKK", Complex.Status.OK));
         Complex c4 = new Complex(new Complex.Id(999_000, 15L, "AAA", Complex.Status.OK));
-        db.tx(() -> db.complexes().insert(c1, c2, c3, c4));
+        tx.run(db -> db.complexes().insert(c1, c2, c3, c4));
 
         FilterExpression<Complex> filter = newFilterBuilder(Complex.class).where("id.b").eq(15L).and("id.a").eq(999_999).build();
         ListRequest<Complex> listRequest = ListRequest.builder(Complex.class).pageSize(3).filter(filter).build();
 
         YqlPredicate predicate = YqlListingQuery.toYqlPredicate(filter);
         assertThat(predicate.toYql(EntitySchema.of(Complex.class))).isEqualTo("(`id_a` = ?) AND (`id_b` = ?)");
 
-        db.tx(() -> {
-            ListResult<Complex> page = listComplex(listRequest);
+        tx.run(db -> {
+            ListResult<Complex> page = db.complexes().list(listRequest);
             assertThat(page).containsExactly(c3, c2, c1);
             assertThat(page.isLastPage()).isTrue();
         });

repository-ydb-v2/src/test/java/tech/ydb/yoj/repository/ydb/merge/YqlQueryMergerIntegrationTest.java

@@ -3,10 +3,10 @@
 import org.junit.ClassRule;
 import org.junit.Test;
 import tech.ydb.yoj.repository.db.Repository;
+import tech.ydb.yoj.repository.db.ScopedTxManager;
 import tech.ydb.yoj.repository.test.RepositoryTestSupport;
 import tech.ydb.yoj.repository.test.entity.TestEntities;
 import tech.ydb.yoj.repository.test.sample.TestDb;
-import tech.ydb.yoj.repository.test.sample.TestDbImpl;
 import tech.ydb.yoj.repository.test.sample.model.Project;
 import tech.ydb.yoj.repository.ydb.TestYdbRepository;
 import tech.ydb.yoj.repository.ydb.YdbEnvAndTransportRule;
@@ -20,17 +20,17 @@ public class YqlQueryMergerIntegrationTest extends RepositoryTestSupport {
     @ClassRule
     public static final YdbEnvAndTransportRule ydbEnvAndTransport = new YdbEnvAndTransportRule();
 
-    protected TestDb db;
+    protected ScopedTxManager<TestDb> tx;
 
     @Override
     public void setUp() {
         super.setUp();
-        this.db = new TestDbImpl<>(this.repository);
+        this.tx = new ScopedTxManager<>(this.repository, TestDb.class);
     }
 
     @Override
     public void tearDown() {
-        this.db = null;
+        this.tx = null;
         super.tearDown();
     }
 
@@ -41,7 +41,7 @@ protected final Repository createRepository() {
 
     @Test
     public void insertAfterFindByIdWorks() {
-        db.tx(() -> {
+        tx.run(db -> {
             Project.Id id = new Project.Id(RandomUtils.nextString(10));
 
             assertThat(db.projects().find(id)).isNull();

repository/src/main/java/tech/ydb/yoj/repository/db/DelegatingTxManager.java

@@ -18,10 +18,13 @@
  * doRunTx()}</li>
  * </ul>
  *
+ * @param <SELF> This class's type. This allows to have {@code TxManager} configuration calls return specific subtype
+ *               instead of just {@code TxManager}
+ *
  * @see #wrapTxBody(Supplier)
  * @see #doRunTx(Supplier)
  */
-public abstract class DelegatingTxManager implements TxManager {
+public abstract class DelegatingTxManager<SELF extends DelegatingTxManager<SELF>> implements TxManager {
     protected final TxManager delegate;
 
     protected DelegatingTxManager(@NonNull TxManager delegate) {
@@ -61,7 +64,7 @@ protected <T> Supplier<T> wrapTxBody(Supplier<T> supplier) {
      * @param delegate transaction manager to delegate to
      * @return wrapped {@code delegate}
      */
-    protected abstract TxManager createTxManager(TxManager delegate);
+    protected abstract SELF createTxManager(TxManager delegate);
 
     @Override
     public final void tx(Runnable runnable) {
@@ -77,102 +80,102 @@ public final <T> T tx(Supplier<T> supplier) {
     }
 
     @Override
-    public final TxManager withName(String name, String logContext) {
+    public final SELF withName(String name, String logContext) {
         return createTxManager(this.delegate.withName(name, logContext));
     }
 
     @Override
-    public final TxManager withName(String name) {
+    public final SELF withName(String name) {
         return createTxManager(this.delegate.withName(name));
     }
 
     @Override
-    public final TxManager withLogContext(String logContext) {
+    public final SELF withLogContext(String logContext) {
         return createTxManager(this.delegate.withLogContext(logContext));
     }
 
     @Override
-    public final TxManager separate() {
+    public final SELF separate() {
         return createTxManager(this.delegate.separate());
     }
 
     @Override
-    public TxManager delayedWrites() {
+    public final SELF delayedWrites() {
         return createTxManager(this.delegate.delayedWrites());
     }
 
     @Override
-    public final TxManager immediateWrites() {
+    public final SELF immediateWrites() {
         return createTxManager(this.delegate.immediateWrites());
     }
 
     @Override
-    public final TxManager noFirstLevelCache() {
+    public final SELF noFirstLevelCache() {
         return createTxManager(this.delegate.noFirstLevelCache());
     }
 
     @Override
-    public final TxManager failOnUnknownSeparateTx() {
+    public final SELF failOnUnknownSeparateTx() {
         return createTxManager(this.delegate.failOnUnknownSeparateTx());
     }
 
     @Override
-    public final TxManager withMaxRetries(int maxRetries) {
+    public final SELF withMaxRetries(int maxRetries) {
         return createTxManager(this.delegate.withMaxRetries(maxRetries));
     }
 
     @Override
-    public final TxManager withDryRun(boolean dryRun) {
+    public final SELF withDryRun(boolean dryRun) {
         return createTxManager(this.delegate.withDryRun(dryRun));
     }
 
     @Override
-    public final TxManager withLogLevel(TransactionLog.Level level) {
+    public final SELF withLogLevel(TransactionLog.Level level) {
         return createTxManager(this.delegate.withLogLevel(level));
     }
 
     @Override
-    public final TxManager withLogStatementOnSuccess(boolean logStatementOnSuccess) {
+    public final SELF withLogStatementOnSuccess(boolean logStatementOnSuccess) {
         return createTxManager(this.delegate.withLogStatementOnSuccess(logStatementOnSuccess));
     }
 
     @Override
-    public final TxManager withTimeout(@NonNull Duration timeout) {
+    public final SELF withTimeout(@NonNull Duration timeout) {
         return createTxManager(this.delegate.withTimeout(timeout));
     }
 
     @Override
-    public final TxManager withQueryStats(@NonNull QueryStatsMode queryStats) {
+    public final SELF withQueryStats(@NonNull QueryStatsMode queryStats) {
         return createTxManager(this.delegate.withQueryStats(queryStats));
     }
 
     @Override
-    public TxManager withFullQueryTracing() {
+    public final SELF withFullQueryTracing() {
         return createTxManager(this.delegate.withFullQueryTracing());
     }
 
     @Override
-    public TxManager noQueryTracing() {
+    public final SELF noQueryTracing() {
         return createTxManager(this.delegate.noQueryTracing());
     }
 
     @Override
-    public TxManager withTracingFilter(@NonNull QueryTracingFilter tracingFilter) {
+    public final SELF withTracingFilter(@NonNull QueryTracingFilter tracingFilter) {
         return createTxManager(this.delegate.withTracingFilter(tracingFilter));
     }
 
     @Override
-    public final TxManager withVerboseLogging() {
+    public final SELF withVerboseLogging() {
         return createTxManager(this.delegate.withVerboseLogging());
     }
 
     @Override
-    public final TxManager withBriefLogging() {
+    public final SELF withBriefLogging() {
         return createTxManager(this.delegate.withBriefLogging());
     }
 
     @Override
-    public final TxManager noLogging() {
+    public final SELF noLogging() {
         return createTxManager(this.delegate.noLogging());
     }
 

repository/src/main/java/tech/ydb/yoj/repository/db/ScopedTxManager.java

@@ -0,0 +1,49 @@
+package tech.ydb.yoj.repository.db;
+
+import lombok.NonNull;
+import tech.ydb.yoj.ExperimentalApi;
+import tech.ydb.yoj.repository.BaseDb;
+
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+@ExperimentalApi(issue = "https://github.com/ydb-platform/yoj-project/issues/197")
+public final class ScopedTxManager<D extends BaseDb> extends DelegatingTxManager<ScopedTxManager<D>> {
+    private final D db;
+
+    public ScopedTxManager(@NonNull Repository repository, @NonNull Class<D> dbClass) {
+        this(new StdTxManager(repository), dbClass);
+    }
+
+    public ScopedTxManager(@NonNull TxManager delegate, @NonNull Class<D> dbClass) {
+        this(delegate, BaseDb.current(dbClass));
+    }
+
+    private ScopedTxManager(@NonNull TxManager delegate, @NonNull D db) {
+        super(delegate);
+        this.db = db;
+    }
+
+    public <R> R call(@NonNull Function<D, R> function) {
+        return tx(() -> function.apply(db));
+    }
+
+    public <T> T call(@NonNull BiFunction<D, Tx, T> function) {
+        return tx(() -> function.apply(db, Tx.Current.get()));
+    }
+
+    public void run(@NonNull Consumer<D> consumer) {
+        tx(() -> consumer.accept(db));
+    }
+
+    public void run(@NonNull BiConsumer<D, Tx> consumer) {
+        tx(() -> consumer.accept(db, Tx.Current.get()));
+    }
+
+    @Override
+    protected ScopedTxManager<D> createTxManager(TxManager delegate) {
+        return new ScopedTxManager<>(delegate, this.db);
+    }
+}