new-tx-i: new Tx interface pattern

repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryCustomQuery.java

@@ -0,0 +1,9 @@
+package tech.ydb.yoj.repository.test.inmemory;
+
+import tech.ydb.yoj.repository.db.Tx;
+
+import java.util.List;
+
+public interface InMemoryCustomQuery<T> extends Tx.CustomQuery<T> {
+    Stream<T> query(PublicInMemoryStorage storage, InMemoryTxLockWatcher watcher);
+}

repository-test/src/main/java/tech/ydb/yoj/repository/test/RepositoryTest.java

@@ -2,6 +2,9 @@
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterators;
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.RequiredArgsConstructor;
 import org.assertj.core.api.Assertions;
 import org.junit.Assert;
 import org.junit.Test;
@@ -29,6 +32,8 @@
 import tech.ydb.yoj.repository.db.list.ListResult;
 import tech.ydb.yoj.repository.db.readtable.ReadTableParams;
 import tech.ydb.yoj.repository.test.entity.TestEntities;
+import tech.ydb.yoj.repository.test.inmemory.InMemoryCustomQuery;
+import tech.ydb.yoj.repository.test.inmemory.InMemoryTxLockWatcher;
 import tech.ydb.yoj.repository.test.sample.TestDb;
 import tech.ydb.yoj.repository.test.sample.TestDbImpl;
 import tech.ydb.yoj.repository.test.sample.TestEntityOperations;
@@ -51,6 +56,7 @@
 import tech.ydb.yoj.repository.test.sample.model.TypeFreak.B;
 import tech.ydb.yoj.repository.test.sample.model.TypeFreak.Embedded;
 import tech.ydb.yoj.repository.test.sample.model.WithUnflattenableField;
+import tech.ydb.yoj.repository.ydb.YdbCustomQuery;
 
 import java.time.Instant;
 import java.util.ArrayList;
@@ -67,6 +73,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import static java.time.temporal.ChronoUnit.MILLIS;
 import static java.util.Arrays.asList;
@@ -288,13 +295,79 @@ public void deferFinallyRollbackRetryable() {
 
     @Test
     public void deferFinallyNotInTxContext() {
-        db.tx(() -> Tx.Current.get().deferFinally(() -> assertFalse(Tx.Current.exists())));
+        db.txC(tx -> tx.deferFinally(() -> assertFalse(Tx.Current.exists())));
+    }
+
+    @Test
+    public void myTest() {
+        db.txC(tx -> {
+            // new. Usefull because you see that table depends on tx;
+            Db1.project(tx).findAll();
+
+            // new. Usefull because you can initialize db one time in tx and use call table without pass tx any time
+            Db2.of(tx).project().findAll();
+
+            // old
+            db.projects().findAll();
+        });
+    }
+
+    // Examples of new DB patterns. User can use one or both;
+
+    @AllArgsConstructor(access = AccessLevel.PRIVATE)
+    final static class Db1 {
+        public static TestEntityOperations.ProjectTable project(Tx tx) {
+            return new TestEntityOperations.ProjectTable(tx.table(Project.class));
+        }
+    }
+
+    @RequiredArgsConstructor(staticName = "of")
+    final static class Db2 {
+        private final Tx tx;
+
+        TestEntityOperations.ProjectTable project() {
+            return new TestEntityOperations.ProjectTable(tx.table(Project.class));
+        }
+    }
+
+    // example of custom query
+
+    @Test
+    public void teatCustomQuery() {
+        db.txC(tx -> {
+            List<Project> projs = tx.customQuery(new MyCustomQuery("my_name")).toList();
+        });
+    }
+
+    @AllArgsConstructor
+    public static class MyCustomQuery implements InMemoryCustomQuery<Project>, YdbCustomQuery<String, Project> {
+        private final String name;
+
+        @Override
+        public Query<String> getQuery() {
+            String query = """
+                --!syntax_v1
+                DECLARE $name AS List<String>;
+                SELECT `id` FROM `@TABLESPACE@/Project`
+                    WHERE (`name` IN $name)
+                    ORDER BY `id` ASC
+            """;
+            return new Query<>(query, name);
+        }
+
+        @Override
+        public Stream<Project> query(PublicInMemoryStorage storage, InMemoryTxLockWatcher watcher) {
+            watcher.markTableRead(Project.class);
+
+            return storage.getDataShard(Project.class).findAll().stream()
+                    .filter(p -> p.getName().equals(name));
+        }
     }
 
     @Test
     public void deferFinallyRollbackNotInTxContext() {
-        assertThatExceptionOfType(RuntimeException.class).isThrownBy(() -> db.tx(() -> {
-            Tx.Current.get().deferFinally(() -> assertFalse(Tx.Current.exists()));
+        assertThatExceptionOfType(RuntimeException.class).isThrownBy(() -> db.txC(tx -> {
+            tx.deferFinally(() -> assertFalse(Tx.Current.exists()));
             throw new RuntimeException();
         }));
     }

repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbCustomQuery.java

@@ -0,0 +1,10 @@
+package tech.ydb.yoj.repository.ydb;
+
+import tech.ydb.yoj.repository.db.Tx;
+
+public interface YdbCustomQuery<PARAMS, T> extends Tx.CustomQuery<T> {
+    record Query<PARAMS>(String taskQuery, PARAMS params) {
+    }
+
+    Query<PARAMS> getQuery();
+}

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

@@ -22,6 +22,7 @@
 import java.time.Duration;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
@@ -181,23 +182,25 @@ public ScanBuilder scan() {
 
     @Override
     public void tx(Runnable runnable) {
-        tx(() -> {
-            runnable.run();
-            return null;
-        });
+        txC(__ -> runnable.run());
     }
 
     @Override
     public <T> T tx(Supplier<T> supplier) {
+        return tx(__ -> supplier.get());
+    }
+
+    @Override
+    public <T> T tx(Function<Tx, T> func) {
         if (name == null) {
-            return withGeneratedNameAndLine().tx(supplier);
+            return withGeneratedNameAndLine().tx(func);
         }
 
         checkSeparatePolicy(separatePolicy, name);
-        return txImpl(supplier);
+        return txImpl(func);
     }
 
-    private <T> T txImpl(Supplier<T> supplier) {
+    private <T> T txImpl(Function<Tx, T> func) {
         RetryableException lastRetryableException = null;
         TxImpl lastTx = null;
         try (Timer ignored = totalDuration.labels(name).startTimer()) {
@@ -207,7 +210,7 @@ private <T> T txImpl(Supplier<T> supplier) {
                     T result;
                     try (var ignored1 = attemptDuration.labels(name).startTimer()) {
                         lastTx = new TxImpl(name, repository.startTransaction(options), options);
-                        result = runAttempt(supplier, lastTx);
+                        result = runAttempt(lastTx, func);
                     }
 
                     if (options.isDryRun()) {
@@ -258,11 +261,11 @@ private String getExceptionNameForMetric(RetryableException e) {
         return Strings.removeSuffix(e.getClass().getSimpleName(), "Exception");
     }
 
-    private <T> T runAttempt(Supplier<T> supplier, TxImpl tx) {
+    private <T> T runAttempt(TxImpl tx, Function<Tx, T> func) {
         try (var ignored2 = MDC.putCloseable("tx", formatTx());
              var ignored3 = MDC.putCloseable("tx-id", formatTxId());
              var ignored4 = MDC.putCloseable("tx-name", formatTxName(false))) {
-            return tx.run(supplier);
+            return tx.run(func);
         }
     }
 

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

@@ -4,14 +4,24 @@
 
 import javax.annotation.Nullable;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 public interface Tx {
+    default <T extends Entity<T>> Table<T> table(Class<T> cls) {
+        return getRepositoryTransaction().table(cls);
+    }
+
     void defer(Runnable runnable);
 
     void deferFinally(Runnable runnable);
 
     void deferBeforeCommit(Runnable runnable);
 
+    <T> Stream<T> customQuery(CustomQuery<T> query);
+
+    interface CustomQuery<T> {
+    }
+
     String getName();
 
     RepositoryTransaction getRepositoryTransaction();

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

@@ -9,7 +9,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Supplier;
+import java.util.function.Function;
 
 final class TxImpl implements Tx {
     private static final Logger log = LoggerFactory.getLogger(TxImpl.class);
@@ -32,10 +32,10 @@ public TxImpl(String name, RepositoryTransaction repositoryTransaction, TxOption
         this.logStatementOnSuccess = options.isLogStatementOnSuccess();
     }
 
-    <R> R run(Supplier<R> supplier) {
+    <R> R run(Function<Tx, R> func) {
         R value;
         try {
-            value = Current.runInTx(this, () -> runImpl(supplier));
+            value = Current.runInTx(this, () -> runImpl(func));
         } catch (Exception e) {
             if (Interrupts.isInterruptException(e)) {
                 Thread.currentThread().interrupt();
@@ -72,11 +72,11 @@ public void deferBeforeCommit(Runnable runnable) {
         deferredBeforeCommit.add(runnable);
     }
 
-    private <R> R runImpl(Supplier<R> supplier) {
+    private <R> R runImpl(Function<Tx, R> func) {
         Stopwatch sw = Stopwatch.createStarted();
         R res;
         try {
-            res = supplier.get();
+            res = func.apply(this);
             deferredBeforeCommit.forEach(Runnable::run);
         } catch (Throwable t) {
             doRollback(isBusinessException(t),

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

@@ -7,6 +7,8 @@
 import tech.ydb.yoj.repository.db.exception.RetryableException;
 
 import java.time.Duration;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 public interface TxManager {
@@ -148,6 +150,15 @@ default TxManager noLogging() {
      */
     void tx(Runnable runnable);
 
+    <T> T tx(Function<Tx, T> supplier);
+
+    default void txC(Consumer<Tx> consumer) {
+        tx(tx -> {
+            consumer.accept(tx);
+            return null;
+        });
+    }
+
     /**
      * Start a transaction-like session of read-only statements. Each statement will be executed <em>separately</em>,
      * with the specified isolation level (online consistent read-only, by default).