From 7ccfc1211bc36aecbbf72d775761fd6407257cac Mon Sep 17 00:00:00 2001
From: Alexander Lavrukov <lavrukov@yandex-team.ru>
Date: Mon, 23 Dec 2024 21:14:18 +0300
Subject: [PATCH] Use TableDescriptor in InMemoryTable

---
 .../test/inmemory/InMemoryDataShard.java      | 27 +++++----
 .../test/inmemory/InMemoryRepository.java     | 20 +++----
 .../InMemoryRepositoryTransaction.java        | 17 ++++--
 .../test/inmemory/InMemoryStorage.java        | 60 +++++++++----------
 .../test/inmemory/InMemoryTable.java          | 53 +++++++++-------
 .../test/inmemory/InMemoryTxLockWatcher.java  | 33 +++++-----
 .../test/RepositoryTestSupport.java           | 11 +---
 .../ydb/yoj/repository/ydb/YdbRepository.java | 28 +++------
 .../ydb/YdbRepositoryTransaction.java         |  6 ++
 .../yoj/repository/ydb/table/YdbTable.java    |  7 +++
 .../java/tech/ydb/yoj/repository/BaseDb.java  |  8 ++-
 .../ydb/yoj/repository/db/Repository.java     | 13 ++--
 .../repository/db/RepositoryTransaction.java  |  6 +-
 13 files changed, 156 insertions(+), 133 deletions(-)

diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java
index 785d8ead..0cfe3abe 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryDataShard.java
@@ -6,6 +6,7 @@
 import tech.ydb.yoj.repository.db.EntitySchema;
 import tech.ydb.yoj.repository.db.Range;
 import tech.ydb.yoj.repository.db.Table;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.ViewSchema;
 import tech.ydb.yoj.repository.db.exception.EntityAlreadyExistsException;
 import tech.ydb.yoj.repository.db.exception.OptimisticLockException;
@@ -20,21 +21,27 @@
 import java.util.TreeMap;
 
 final class InMemoryDataShard<T extends Entity<T>> {
-    private final Class<T> type;
+    private final TableDescriptor<T> tableDescriptor;
     private final EntitySchema<T> schema;
     private final TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines;
     private final Map<Long, Set<Entity.Id<T>>> uncommited = new HashMap<>();
 
     private InMemoryDataShard(
-            Class<T> type, EntitySchema<T> schema, TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines
+            TableDescriptor<T> tableDescriptor,
+            EntitySchema<T> schema,
+            TreeMap<Entity.Id<T>, InMemoryEntityLine> entityLines
     ) {
-        this.type = type;
+        this.tableDescriptor = tableDescriptor;
         this.schema = schema;
         this.entityLines = entityLines;
     }
 
-    public InMemoryDataShard(Class<T> type) {
-        this(type, EntitySchema.of(type), createEmptyLines(type));
+    public InMemoryDataShard(TableDescriptor<T> tableDescriptor) {
+        this(
+                tableDescriptor,
+                EntitySchema.of(tableDescriptor.entityType()),
+                createEmptyLines(tableDescriptor.entityType())
+        );
     }
 
     private static <T extends Entity<T>> TreeMap<Entity.Id<T>, InMemoryEntityLine> createEmptyLines(Class<T> type) {
@@ -42,11 +49,11 @@ private static <T extends Entity<T>> TreeMap<Entity.Id<T>, InMemoryEntityLine> c
     }
 
     public synchronized InMemoryDataShard<T> createSnapshot() {
-        TreeMap<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(type);
+        TreeMap<Entity.Id<T>, InMemoryEntityLine> snapshotLines = createEmptyLines(tableDescriptor.entityType());
         for (Map.Entry<Entity.Id<T>, InMemoryEntityLine> entry : entityLines.entrySet()) {
             snapshotLines.put(entry.getKey(), entry.getValue().createSnapshot());
         }
-        return new InMemoryDataShard<>(type, schema, snapshotLines);
+        return new InMemoryDataShard<>(tableDescriptor, schema, snapshotLines);
     }
 
     public synchronized void commit(long txId, long version) {
@@ -60,14 +67,14 @@ public synchronized void commit(long txId, long version) {
     }
 
     public synchronized void checkLocks(long version, InMemoryTxLockWatcher watcher) {
-        for (Entity.Id<T> lockedId : watcher.getReadRows(type)) {
+        for (Entity.Id<T> lockedId : watcher.getReadRows(tableDescriptor)) {
             InMemoryEntityLine entityLine = entityLines.get(lockedId);
             if (entityLine != null && entityLine.hasYounger(version)) {
                 throw new OptimisticLockException("Row lock failed " + lockedId);
             }
         }
 
-        List<Range<Entity.Id<T>>> lockedRanges = watcher.getReadRanges(type);
+        List<Range<Entity.Id<T>>> lockedRanges = watcher.getReadRanges(tableDescriptor);
         if (lockedRanges.isEmpty()) {
             return;
         }
@@ -79,7 +86,7 @@ public synchronized void checkLocks(long version, InMemoryTxLockWatcher watcher)
 
             for (Range<Entity.Id<T>> lockedRange: lockedRanges) {
                 if (lockedRange.contains(entry.getKey())) {
-                    throw new OptimisticLockException("Table lock failed " + type.getSimpleName());
+                    throw new OptimisticLockException("Table lock failed " + tableDescriptor.toDebugString());
                 }
             }
         }
diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepository.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepository.java
index a133c1b8..21b30de0 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepository.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepository.java
@@ -6,11 +6,11 @@
 import tech.ydb.yoj.repository.db.Repository;
 import tech.ydb.yoj.repository.db.RepositoryTransaction;
 import tech.ydb.yoj.repository.db.SchemaOperations;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.TxOptions;
 import tech.ydb.yoj.repository.db.exception.DropTableException;
 
 import java.util.Map;
-import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -37,36 +37,30 @@ public void loadSnapshot(String id) {
         storage = snapshots.get(id).createSnapshot();
     }
 
-    @Override
-    public Set<Class<? extends Entity<?>>> tables() {
-        return Set.copyOf(storage.tables());
-    }
-
     @Override
     public RepositoryTransaction startTransaction(TxOptions options) {
         return new InMemoryRepositoryTransaction(options, this);
     }
 
-    @Override
-    public <T extends Entity<T>> SchemaOperations<T> schema(Class<T> c) {
+    public <T extends Entity<T>> SchemaOperations<T> schema(TableDescriptor<T> tableDescriptor) {
         return new SchemaOperations<T>() {
             @Override
             public void create() {
-                storage.createTable(c);
+                storage.createTable(tableDescriptor);
             }
 
             @Override
             public void drop() {
-                if (!storage.dropTable(c)) {
-                    throw new DropTableException(
-                            String.format("Can't drop table %s: table doesn't exist", c.getSimpleName())
+                if (!storage.dropTable(tableDescriptor)) {
+                    throw new DropTableException(String.format("Can't drop table %s: table doesn't exist",
+                            tableDescriptor.toDebugString())
                     );
                 }
             }
 
             @Override
             public boolean exists() {
-                return storage.containsTable(c);
+                return storage.containsTable(tableDescriptor);
             }
         };
     }
diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepositoryTransaction.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepositoryTransaction.java
index 89e15003..b154c3cb 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepositoryTransaction.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryRepositoryTransaction.java
@@ -7,6 +7,7 @@
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.RepositoryTransaction;
 import tech.ydb.yoj.repository.db.Table;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.TxOptions;
 import tech.ydb.yoj.repository.db.cache.TransactionLocal;
 import tech.ydb.yoj.repository.db.exception.IllegalTransactionIsolationLevelException;
@@ -59,6 +60,12 @@ public <T extends Entity<T>> Table<T> table(Class<T> c) {
         return new InMemoryTable<>(getMemory(c));
     }
 
+    @Override
+    public <T extends Entity<T>> Table<T> table(TableDescriptor<T> tableDescriptor) {
+        return new InMemoryTable<>(this, tableDescriptor);
+    }
+
+    @Deprecated // use other constructor of InMemoryTable
     public final <T extends Entity<T>> InMemoryTable.DbMemory<T> getMemory(Class<T> c) {
         return new InMemoryTable.DbMemory<>(c, this);
     }
@@ -119,7 +126,7 @@ private boolean isFinalActionNeeded(String action) {
     }
 
     final <T extends Entity<T>> void doInWriteTransaction(
-            String log, Class<T> type, Consumer<WriteTxDataShard<T>> consumer
+            String log, TableDescriptor<T> tableDescriptor, Consumer<WriteTxDataShard<T>> consumer
     ) {
         if (options.isScan()) {
             throw new IllegalTransactionScanException("Mutable operations");
@@ -129,7 +136,7 @@ final <T extends Entity<T>> void doInWriteTransaction(
         }
 
         Runnable query = () -> logTransaction(log, () -> {
-            WriteTxDataShard<T> shard = storage.getWriteTxDataShard(type, txId, getVersion());
+            WriteTxDataShard<T> shard = storage.getWriteTxDataShard(tableDescriptor, txId, getVersion());
             consumer.accept(shard);
 
             hasWrites = true;
@@ -143,11 +150,13 @@ final <T extends Entity<T>> void doInWriteTransaction(
     }
 
     final <T extends Entity<T>, R> R doInTransaction(
-            String action, Class<T> type, Function<ReadOnlyTxDataShard<T>, R> func
+            String action, TableDescriptor<T> tableDescriptor, Function<ReadOnlyTxDataShard<T>, R> func
     ) {
         return logTransaction(action, () -> {
             InMemoryTxLockWatcher findWatcher = hasWrites ? watcher : InMemoryTxLockWatcher.NO_LOCKS;
-            ReadOnlyTxDataShard<T> shard = storage.getReadOnlyTxDataShard(type, txId, getVersion(), findWatcher);
+            ReadOnlyTxDataShard<T> shard = storage.getReadOnlyTxDataShard(
+                    tableDescriptor, txId, getVersion(), findWatcher
+            );
             try {
                 return func.apply(shard);
             } catch (OptimisticLockException e) {
diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryStorage.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryStorage.java
index 353ce90c..f78d86f6 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryStorage.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryStorage.java
@@ -1,6 +1,7 @@
 package tech.ydb.yoj.repository.test.inmemory;
 
 import tech.ydb.yoj.repository.db.Entity;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 import java.util.HashMap;
 import java.util.HashSet;
@@ -8,8 +9,8 @@
 import java.util.Set;
 
 final class InMemoryStorage {
-    private final Map<Class<?>, InMemoryDataShard<?>> shards;
-    private final Map<Long, Set<Class<?>>> uncommited = new HashMap<>();
+    private final Map<TableDescriptor<?>, InMemoryDataShard<?>> shards;
+    private final Map<Long, Set<TableDescriptor<?>>> uncommited = new HashMap<>();
 
     private long currentVersion;
 
@@ -17,7 +18,7 @@ public InMemoryStorage() {
         this(0, new HashMap<>());
     }
 
-    private InMemoryStorage(long version, Map<Class<?>, InMemoryDataShard<?>> shards) {
+    private InMemoryStorage(long version, Map<TableDescriptor<?>, InMemoryDataShard<?>> shards) {
         this.shards = shards;
         this.currentVersion = version;
     }
@@ -27,8 +28,8 @@ public synchronized long getCurrentVersion() {
     }
 
     public synchronized InMemoryStorage createSnapshot() {
-        Map<Class<?>, InMemoryDataShard<?>> snapshotDb = new HashMap<>();
-        for (Map.Entry<Class<?>, InMemoryDataShard<?>> entry : shards.entrySet()) {
+        Map<TableDescriptor<?>, InMemoryDataShard<?>> snapshotDb = new HashMap<>();
+        for (Map.Entry<TableDescriptor<?>, InMemoryDataShard<?>> entry : shards.entrySet()) {
             snapshotDb.put(entry.getKey(), entry.getValue().createSnapshot());
         }
         return new InMemoryStorage(currentVersion, snapshotDb);
@@ -45,42 +46,42 @@ public synchronized void commit(long txId, long version, InMemoryTxLockWatcher w
 
         currentVersion++;
 
-        Set<Class<?>> uncommitedTables = uncommited.remove(txId);
-        for (Class<?> type : uncommitedTables) {
-            shards.get(type).commit(txId, currentVersion);
+        Set<TableDescriptor<?>> uncommitedTables = uncommited.remove(txId);
+        for (TableDescriptor<?> tableDescriptor : uncommitedTables) {
+            shards.get(tableDescriptor).commit(txId, currentVersion);
         }
     }
 
     public synchronized void rollback(long txId) {
-        Set<Class<?>> uncommitedTables = uncommited.remove(txId);
+        Set<TableDescriptor<?>> uncommitedTables = uncommited.remove(txId);
         if (uncommitedTables == null) {
             return;
         }
-        for (Class<?> type : uncommitedTables) {
-            shards.get(type).rollback(txId);
+        for (TableDescriptor<?> tableDescriptor : uncommitedTables) {
+            shards.get(tableDescriptor).rollback(txId);
         }
     }
 
     public synchronized <T extends Entity<T>> WriteTxDataShard<T> getWriteTxDataShard(
-            Class<T> type, long txId, long version
+            TableDescriptor<T> tableDescriptor, long txId, long version
     ) {
-        uncommited.computeIfAbsent(txId, __ -> new HashSet<>()).add(type);
-        return getTxDataShard(type, txId, version, InMemoryTxLockWatcher.NO_LOCKS);
+        uncommited.computeIfAbsent(txId, __ -> new HashSet<>()).add(tableDescriptor);
+        return getTxDataShard(tableDescriptor, txId, version, InMemoryTxLockWatcher.NO_LOCKS);
     }
 
     public synchronized <T extends Entity<T>> ReadOnlyTxDataShard<T> getReadOnlyTxDataShard(
-            Class<T> type, long txId, long version, InMemoryTxLockWatcher watcher
+            TableDescriptor<T> tableDescriptor, long txId, long version, InMemoryTxLockWatcher watcher
     ) {
-        return getTxDataShard(type, txId, version, watcher);
+        return getTxDataShard(tableDescriptor, txId, version, watcher);
     }
 
     private <T extends Entity<T>> TxDataShardImpl<T> getTxDataShard(
-            Class<T> type, long txId, long version, InMemoryTxLockWatcher watcher
+            TableDescriptor<T> tableDescriptor, long txId, long version, InMemoryTxLockWatcher watcher
     ) {
         @SuppressWarnings("unchecked")
-        InMemoryDataShard<T> shard = (InMemoryDataShard<T>) shards.get(type);
+        InMemoryDataShard<T> shard = (InMemoryDataShard<T>) shards.get(tableDescriptor);
         if (shard == null) {
-            throw new InMemoryRepositoryException("Table is not created: " + type.getSimpleName());
+            throw new InMemoryRepositoryException("Table is not created: " + tableDescriptor.toDebugString());
         }
         return new TxDataShardImpl<>(shard, txId, version, watcher);
     }
@@ -89,27 +90,26 @@ public synchronized void dropDb() {
         shards.clear();
     }
 
-    @SuppressWarnings({"unchecked", "rawtypes"})
-    public synchronized Set<Class<? extends Entity<?>>> tables() {
-        return (Set) shards.keySet();
+    public synchronized Set<TableDescriptor<?>> tables() {
+        return shards.keySet();
     }
 
-    public synchronized boolean containsTable(Class<?> type) {
-        return shards.containsKey(type);
+    public synchronized boolean containsTable(TableDescriptor<?> tableDescriptor) {
+        return shards.containsKey(tableDescriptor);
     }
 
-    public synchronized <T extends Entity<T>> void createTable(Class<T> type) {
-        if (containsTable(type)) {
+    public synchronized <T extends Entity<T>> void createTable(TableDescriptor<T> tableDescriptor) {
+        if (containsTable(tableDescriptor)) {
             return;
         }
-        shards.put(type, new InMemoryDataShard<>(type));
+        shards.put(tableDescriptor, new InMemoryDataShard<>(tableDescriptor));
     }
 
-    public synchronized boolean dropTable(Class<?> type) {
-        if (!containsTable(type)) {
+    public synchronized boolean dropTable(TableDescriptor<?> tableDescriptor) {
+        if (!containsTable(tableDescriptor)) {
             return false;
         }
-        shards.remove(type);
+        shards.remove(tableDescriptor);
         return true;
     }
 }
diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTable.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTable.java
index 47b4e308..3779071e 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTable.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTable.java
@@ -33,21 +33,28 @@
 import static java.util.stream.Collectors.toUnmodifiableSet;
 
 public class InMemoryTable<T extends Entity<T>> implements Table<T> {
-    private final Class<T> type;
     private final EntitySchema<T> schema;
     private final TableDescriptor<T> tableDescriptor;
     private final InMemoryRepositoryTransaction transaction;
 
+    @Deprecated // Don't use DbMemory, use other constructor instead
     public InMemoryTable(DbMemory<T> memory) {
-        this.type = memory.type();
-        this.schema = EntitySchema.of(type);
+        this(memory.transaction(), memory.type());
+    }
+
+    public InMemoryTable(InMemoryRepositoryTransaction transaction, Class<T> type) {
+        this(transaction, TableDescriptor.from(EntitySchema.of(type)));
+    }
+
+    public InMemoryTable(InMemoryRepositoryTransaction transaction, TableDescriptor<T> tableDescriptor) {
+        this.schema = EntitySchema.of(tableDescriptor.entityType());
         this.tableDescriptor = TableDescriptor.from(schema);
-        this.transaction = memory.transaction();
+        this.transaction = transaction;
     }
 
     @Override
     public List<T> findAll() {
-        transaction.getWatcher().markTableRead(type);
+        transaction.getWatcher().markTableRead(tableDescriptor);
         return findAll0();
     }
 
@@ -155,7 +162,7 @@ public <V extends ViewId<T>, ID extends Entity.Id<T>> Stream<V> readTable(
 
     @Override
     public Class<T> getType() {
-        return type;
+        return tableDescriptor.entityType();
     }
 
     @Override
@@ -165,7 +172,7 @@ public T find(Entity.Id<T> id) {
         }
         return transaction.getTransactionLocal().firstLevelCache().get(id, __ -> {
             markKeyRead(id);
-            T entity = transaction.doInTransaction("find(" + id + ")", type, shard -> shard.find(id));
+            T entity = transaction.doInTransaction("find(" + id + ")", tableDescriptor, shard -> shard.find(id));
             return postLoad(entity);
         });
     }
@@ -184,13 +191,13 @@ public <V extends View> V find(Class<V> viewType, Entity.Id<T> id) {
         }
 
         markKeyRead(id);
-        return transaction.doInTransaction("find(" + id + ")", type, shard -> shard.find(id, viewType));
+        return transaction.doInTransaction("find(" + id + ")", tableDescriptor, shard -> shard.find(id, viewType));
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public <ID extends Entity.Id<T>> List<T> find(Range<ID> range) {
-        transaction.getWatcher().markRangeRead(type, range);
+        transaction.getWatcher().markRangeRead(tableDescriptor, range);
         return findAll0().stream()
                 .filter(e -> range.contains((ID) e.getId()))
                 .sorted(EntityIdSchema.SORT_ENTITY_BY_ID)
@@ -352,7 +359,7 @@ public <KEY> List<T> find(
         );
 
         for (Map<String, Object> id : keysSet) {
-            transaction.getWatcher().markRangeRead(type, id);
+            transaction.getWatcher().markRangeRead(tableDescriptor, id);
         }
 
         Stream<T> result = getAllEntries().stream()
@@ -388,9 +395,9 @@ private <ID extends Entity.Id<T>> void markKeyRead(ID id) {
         EntityIdSchema<Entity.Id<T>> idSchema = schema.getIdSchema();
         if (idSchema.flattenFieldNames().size() != idSchema.flatten(id).size()) {
             // Partial key, will throw error when not searching by PK prefix
-            transaction.getWatcher().markRangeRead(type, Range.create(id, id));
+            transaction.getWatcher().markRangeRead(tableDescriptor, Range.create(id, id));
         } else {
-            transaction.getWatcher().markRowRead(type, id);
+            transaction.getWatcher().markRowRead(tableDescriptor, id);
         }
     }
 
@@ -406,8 +413,8 @@ public <ID extends Entity.Id<T>> List<ID> findIds(Set<ID> ids) {
     @Override
     public T insert(T tt) {
         T t = tt.preSave();
-        transaction.getWatcher().markRowRead(type, t.getId());
-        transaction.doInWriteTransaction("insert(" + t + ")", type, shard -> shard.insert(t));
+        transaction.getWatcher().markRowRead(tableDescriptor, t.getId());
+        transaction.doInWriteTransaction("insert(" + t + ")", tableDescriptor, shard -> shard.insert(t));
         transaction.getTransactionLocal().firstLevelCache().put(t);
         transaction.getTransactionLocal().projectionCache().save(t);
         return t;
@@ -416,7 +423,7 @@ public T insert(T tt) {
     @Override
     public T save(T tt) {
         T t = tt.preSave();
-        transaction.doInWriteTransaction("save(" + t + ")", type, shard -> shard.save(t));
+        transaction.doInWriteTransaction("save(" + t + ")", tableDescriptor, shard -> shard.save(t));
         transaction.getTransactionLocal().firstLevelCache().put(t);
         transaction.getTransactionLocal().projectionCache().save(t);
         return t;
@@ -424,18 +431,22 @@ public T save(T tt) {
 
     @Override
     public void delete(Entity.Id<T> id) {
-        transaction.doInWriteTransaction("delete(" + id + ")", type, shard -> shard.delete(id));
+        transaction.doInWriteTransaction("delete(" + id + ")", tableDescriptor, shard -> shard.delete(id));
         transaction.getTransactionLocal().firstLevelCache().putEmpty(id);
         transaction.getTransactionLocal().projectionCache().delete(id);
     }
 
     @Override
     public void deleteAll() {
-        transaction.doInWriteTransaction("deleteAll(" + type.getName() + ")", type, WriteTxDataShard::deleteAll);
+        transaction.doInWriteTransaction(
+                "deleteAll(" + tableDescriptor.toDebugString() + ")", tableDescriptor, WriteTxDataShard::deleteAll
+        );
     }
 
     private List<T> getAllEntries() {
-        return transaction.doInTransaction("findAll(" + type.getName() + ")", type, ReadOnlyTxDataShard::findAll);
+        return transaction.doInTransaction(
+                "findAll(" + tableDescriptor.toDebugString() + ")", tableDescriptor, ReadOnlyTxDataShard::findAll
+        );
     }
 
     private List<T> findAll0() {
@@ -494,9 +505,9 @@ public <ID extends Entity.Id<T>> Stream<ID> streamPartialIds(ID partial, int bat
 
     private <ID extends Entity.Id<T>> void markRangeRead(Range<ID> range) {
         if (range == null) {
-            transaction.getWatcher().markTableRead(type);
+            transaction.getWatcher().markTableRead(tableDescriptor);
         } else {
-            transaction.getWatcher().markRangeRead(type, range);
+            transaction.getWatcher().markRangeRead(tableDescriptor, range);
         }
     }
 
@@ -565,6 +576,8 @@ private static <V extends Table.View, T extends Entity<T>> V toView(
         return Columns.fromEntity(schema, entity).toSchema(viewSchema);
     }
 
+
+    @Deprecated // Legacy. Using only for creating InMemoryTable. Use constructor of InMemoryTable instead
     public record DbMemory<T extends Entity<T>>(
             Class<T> type,
             InMemoryRepositoryTransaction transaction
diff --git a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTxLockWatcher.java b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTxLockWatcher.java
index 6d7d7b87..878280b2 100644
--- a/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTxLockWatcher.java
+++ b/repository-inmemory/src/main/java/tech/ydb/yoj/repository/test/inmemory/InMemoryTxLockWatcher.java
@@ -5,6 +5,7 @@
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
 import tech.ydb.yoj.repository.db.Range;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -17,40 +18,40 @@
 public final class InMemoryTxLockWatcher {
     public static final InMemoryTxLockWatcher NO_LOCKS = new InMemoryTxLockWatcher(Map.of(), Map.of());
 
-    private final Map<Class<?>, Set<Entity.Id<?>>> readRows;
-    private final Map<Class<?>, List<Range<?>>> readRanges;
+    private final Map<TableDescriptor<?>, Set<Entity.Id<?>>> readRows;
+    private final Map<TableDescriptor<?>, List<Range<?>>> readRanges;
 
     public InMemoryTxLockWatcher() {
         this(new HashMap<>(), new HashMap<>());
     }
 
-    public <T extends Entity<T>> void markRowRead(Class<T> type, Entity.Id<T> id) {
-        readRows.computeIfAbsent(type, __ -> new HashSet<>()).add(id);
+    public <T extends Entity<T>> void markRowRead(TableDescriptor<T> tableDescriptor, Entity.Id<T> id) {
+        readRows.computeIfAbsent(tableDescriptor, __ -> new HashSet<>()).add(id);
     }
 
-    public <T extends Entity<T>, ID extends Entity.Id<T>> void markRangeRead(Class<T> type, Range<ID> range) {
-        readRanges.computeIfAbsent(type, __ -> new ArrayList<>()).add(range);
+    public <T extends Entity<T>, ID extends Entity.Id<T>> void markRangeRead(TableDescriptor<T> tableDescriptor, Range<ID> range) {
+        readRanges.computeIfAbsent(tableDescriptor, __ -> new ArrayList<>()).add(range);
     }
 
-    public <T extends Entity<T>, ID extends Entity.Id<T>> void markRangeRead(Class<T> type, Map<String, Object> map) {
-        Range<ID> range = Range.create(EntitySchema.of(type).getIdSchema(), map);
-        markRangeRead(type, range);
+    public <T extends Entity<T>, ID extends Entity.Id<T>> void markRangeRead(TableDescriptor<T> tableDescriptor, Map<String, Object> map) {
+        Range<ID> range = Range.create(EntitySchema.of(tableDescriptor.entityType()).getIdSchema(), map);
+        markRangeRead(tableDescriptor, range);
     }
 
-    public <T extends Entity<T>, ID extends Entity.Id<T>> void markTableRead(Class<T> type) {
-        Range<ID> range = Range.create(EntitySchema.of(type).getIdSchema(), Map.of());
-        markRangeRead(type, range);
+    public <T extends Entity<T>, ID extends Entity.Id<T>> void markTableRead(TableDescriptor<T> tableDescriptor) {
+        Range<ID> range = Range.create(EntitySchema.of(tableDescriptor.entityType()).getIdSchema(), Map.of());
+        markRangeRead(tableDescriptor, range);
     }
 
-    public <T extends Entity<T>> Set<Entity.Id<T>> getReadRows(Class<T> type) {
+    public <T extends Entity<T>> Set<Entity.Id<T>> getReadRows(TableDescriptor<T> tableDescriptor) {
         @SuppressWarnings({"unchecked", "rawtypes"})
-        Set<Entity.Id<T>> lockedRows = (Set) readRows.getOrDefault(type, Set.of());
+        Set<Entity.Id<T>> lockedRows = (Set) readRows.getOrDefault(tableDescriptor, Set.of());
         return lockedRows;
     }
 
-    public <T extends Entity<T>> List<Range<Entity.Id<T>>> getReadRanges(Class<T> type) {
+    public <T extends Entity<T>> List<Range<Entity.Id<T>>> getReadRanges(TableDescriptor<T> tableDescriptor) {
         @SuppressWarnings({"unchecked", "rawtypes"})
-        List<Range<Entity.Id<T>>> lockedRanges = (List) readRanges.getOrDefault(type, List.of());
+        List<Range<Entity.Id<T>>> lockedRanges = (List) readRanges.getOrDefault(tableDescriptor, List.of());
         return lockedRanges;
     }
 }
diff --git a/repository-test/src/main/java/tech/ydb/yoj/repository/test/RepositoryTestSupport.java b/repository-test/src/main/java/tech/ydb/yoj/repository/test/RepositoryTestSupport.java
index 6cbe33db..35c47e63 100644
--- a/repository-test/src/main/java/tech/ydb/yoj/repository/test/RepositoryTestSupport.java
+++ b/repository-test/src/main/java/tech/ydb/yoj/repository/test/RepositoryTestSupport.java
@@ -3,8 +3,6 @@
 import org.junit.After;
 import org.junit.Before;
 import tech.ydb.yoj.repository.db.Repository;
-import tech.ydb.yoj.repository.db.StdTxManager;
-import tech.ydb.yoj.repository.db.Tx;
 
 import java.util.IdentityHashMap;
 import java.util.Map;
@@ -30,13 +28,6 @@ public void setUp() {
 
     @After
     public void tearDown() {
-        clearDb(this.repository);
-    }
-
-    @SuppressWarnings({"unchecked", "rawtypes"})
-    public static void clearDb(Repository repo) {
-        var tables = repo.tables();
-        new StdTxManager(repo).tx(() -> tables
-                .forEach(table -> Tx.Current.get().getRepositoryTransaction().table((Class) table).deleteAll()));
+        repository.dropDb();
     }
 }
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
index dcfe87ed..63d38691 100644
--- a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
+++ b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
@@ -18,6 +18,7 @@
 import tech.ydb.yoj.repository.db.Repository;
 import tech.ydb.yoj.repository.db.RepositoryTransaction;
 import tech.ydb.yoj.repository.db.SchemaOperations;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.TxOptions;
 import tech.ydb.yoj.repository.ydb.client.SessionManager;
 import tech.ydb.yoj.repository.ydb.client.YdbPaths;
@@ -34,15 +35,12 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.Objects;
-import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
-import static java.util.stream.Collectors.toUnmodifiableSet;
 import static tech.ydb.yoj.repository.ydb.client.YdbPaths.canonicalDatabase;
 
 public class YdbRepository implements Repository {
@@ -198,14 +196,6 @@ public void createTablespace() {
         getSchemaOperations().createTablespace();
     }
 
-    @Override
-    public Set<Class<? extends Entity<?>>> tables() {
-        return getSchemaOperations().getTableNames().stream()
-                .map(entityClassesByTableName::get)
-                .filter(Objects::nonNull)
-                .collect(toUnmodifiableSet());
-    }
-
     @Override
     public RepositoryTransaction startTransaction(TxOptions options) {
         return new YdbRepositoryTransaction<>(this, options);
@@ -249,12 +239,12 @@ public void dropDb() {
     }
 
     @Override
-    public <T extends Entity<T>> SchemaOperations<T> schema(Class<T> c) {
-        EntitySchema<T> schema = EntitySchema.of(c);
+    public <T extends Entity<T>> SchemaOperations<T> schema(TableDescriptor<T> tableDescriptor) {
+        EntitySchema<T> schema = EntitySchema.of(tableDescriptor.entityType());
         return new SchemaOperations<>() {
             @Override
             public void create() {
-                String tableName = schema.getName();
+                String tableName = tableDescriptor.tableName();
                 getSchemaOperations().createTable(
                         tableName,
                         schema.flattenFields(),
@@ -264,12 +254,12 @@ public void create() {
                         schema.getTtlModifier(),
                         schema.getChangefeeds()
                 );
-                entityClassesByTableName.put(tableName, c);
+                entityClassesByTableName.put(tableName, tableDescriptor.entityType());
             }
 
             private YdbTableHint extractHint() {
                 try {
-                    Field ydbTableHintField = c.getDeclaredField("ydbTableHint");
+                    Field ydbTableHintField = tableDescriptor.entityType().getDeclaredField("ydbTableHint");
                     ydbTableHintField.setAccessible(true);
                     return (YdbTableHint) ydbTableHintField.get(null);
                 } catch (NoSuchFieldException | IllegalAccessException ignored) {
@@ -279,17 +269,17 @@ private YdbTableHint extractHint() {
 
             @Override
             public void drop() {
-                String tableName = schema.getName();
+                String tableName = tableDescriptor.tableName();
                 getSchemaOperations().dropTable(tableName);
                 entityClassesByTableName.remove(tableName);
             }
 
             @Override
             public boolean exists() {
-                String tableName = schema.getName();
+                String tableName = tableDescriptor.tableName();
                 boolean exists = getSchemaOperations().hasTable(tableName);
                 if (exists) {
-                    entityClassesByTableName.put(tableName, c);
+                    entityClassesByTableName.put(tableName, tableDescriptor.entityType());
                 } else {
                     entityClassesByTableName.remove(tableName);
                 }
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
index 5a9edd93..55a184e6 100644
--- 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
@@ -36,6 +36,7 @@
 import tech.ydb.yoj.repository.db.IsolationLevel;
 import tech.ydb.yoj.repository.db.RepositoryTransaction;
 import tech.ydb.yoj.repository.db.Table;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.TxOptions;
 import tech.ydb.yoj.repository.db.bulk.BulkParams;
 import tech.ydb.yoj.repository.db.cache.RepositoryCache;
@@ -120,6 +121,11 @@ public <T extends Entity<T>> Table<T> table(Class<T> c) {
         return new YdbTable<>(c, this);
     }
 
+    @Override
+    public <T extends Entity<T>> Table<T> table(TableDescriptor<T> tableDescriptor) {
+        return new YdbTable<>(tableDescriptor, this);
+    }
+
     @Override
     public void commit() {
         if (isBadSession) {
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/table/YdbTable.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/table/YdbTable.java
index ba253df8..393853cd 100644
--- a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/table/YdbTable.java
+++ b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/table/YdbTable.java
@@ -85,6 +85,13 @@ protected YdbTable(QueryExecutor executor) {
         this.tableDescriptor = TableDescriptor.from(schema);
     }
 
+    public YdbTable(TableDescriptor<T> tableDescriptor, QueryExecutor executor) {
+        this.type = tableDescriptor.entityType();
+        this.executor = new CheckingQueryExecutor(executor);
+        this.schema = EntitySchema.of(type);
+        this.tableDescriptor = tableDescriptor;
+    }
+
     @SuppressWarnings("unchecked")
     private Class<T> resolveEntityType() {
         return (Class<T>) (new TypeToken<T>(getClass()) {
diff --git a/repository/src/main/java/tech/ydb/yoj/repository/BaseDb.java b/repository/src/main/java/tech/ydb/yoj/repository/BaseDb.java
index 7a0eaf2b..ebeffbde 100644
--- a/repository/src/main/java/tech/ydb/yoj/repository/BaseDb.java
+++ b/repository/src/main/java/tech/ydb/yoj/repository/BaseDb.java
@@ -1,7 +1,9 @@
 package tech.ydb.yoj.repository;
 
 import tech.ydb.yoj.repository.db.Entity;
+import tech.ydb.yoj.repository.db.EntitySchema;
 import tech.ydb.yoj.repository.db.Table;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.Tx;
 import tech.ydb.yoj.util.lang.Proxies;
 
@@ -10,5 +12,9 @@ static <T> T current(Class<T> type) {
         return Proxies.proxy(type, () -> Tx.Current.get().getRepositoryTransaction());
     }
 
-    <T extends Entity<T>> Table<T> table(Class<T> c);
+    default <T extends Entity<T>> Table<T> table(Class<T> c) {
+        return table(TableDescriptor.from(EntitySchema.of(c)));
+    }
+
+    <T extends Entity<T>> Table<T> table(TableDescriptor<T> tableDescriptor);
 }
diff --git a/repository/src/main/java/tech/ydb/yoj/repository/db/Repository.java b/repository/src/main/java/tech/ydb/yoj/repository/db/Repository.java
index f7b44ca6..ecdc8aad 100644
--- a/repository/src/main/java/tech/ydb/yoj/repository/db/Repository.java
+++ b/repository/src/main/java/tech/ydb/yoj/repository/db/Repository.java
@@ -1,7 +1,5 @@
 package tech.ydb.yoj.repository.db;
 
-import java.util.Set;
-
 public interface Repository {
     default void createTablespace() {
     }
@@ -12,14 +10,11 @@ default void checkDataCompatibility() {
     default void checkSchemaCompatibility() {
     }
 
-    <T extends Entity<T>> SchemaOperations<T> schema(Class<T> c);
+    default <T extends Entity<T>> SchemaOperations<T> schema(Class<T> c) {
+        return schema(TableDescriptor.from(EntitySchema.of(c)));
+    }
 
-    /**
-     * @deprecated For testing purposes only. Will only <em>reliably</em> work for tables that were created or inspected
-     * using calls to {@link #schema(Class)}.
-     */
-    @Deprecated
-    Set<Class<? extends Entity<?>>> tables();
+    <T extends Entity<T>> SchemaOperations<T> schema(TableDescriptor<T> c);
 
     default RepositoryTransaction startTransaction() {
         return startTransaction(IsolationLevel.SERIALIZABLE_READ_WRITE);
diff --git a/repository/src/main/java/tech/ydb/yoj/repository/db/RepositoryTransaction.java b/repository/src/main/java/tech/ydb/yoj/repository/db/RepositoryTransaction.java
index 112f9c8c..4a5d9e8a 100644
--- a/repository/src/main/java/tech/ydb/yoj/repository/db/RepositoryTransaction.java
+++ b/repository/src/main/java/tech/ydb/yoj/repository/db/RepositoryTransaction.java
@@ -8,7 +8,11 @@
  * one call to either method) lest your transaction stays active on the DB server.
  */
 public interface RepositoryTransaction {
-    <T extends Entity<T>> Table<T> table(Class<T> c);
+    default <T extends Entity<T>> Table<T> table(Class<T> c) {
+        return table(TableDescriptor.from(EntitySchema.of(c)));
+    }
+
+    <T extends Entity<T>> Table<T> table(TableDescriptor<T> tableDescriptor);
 
     /**
      * Commits the transaction or throws exception. Note that this method is not expected to be called, if the last