TableDescriptor in Statements (#106)

There is only one change - I create a TableDescriptor record and pass it to the only one сщтыекгсещк ща YqlStatement, what means that all statements have to get TableDescriptor in their constructor Co-authored-by: Alexander Lavrukov <[email protected]>
ydb-platform · Dec 6, 2024 · e0c3f8b · e0c3f8b
1 parent ee521b4
commit e0c3f8b
Show file tree

Hide file tree

Showing 25 changed files with 278 additions and 125 deletions.
diff --git a/...tory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/merge/ByEntityYqlQueriesMerger.java b/...tory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/merge/ByEntityYqlQueriesMerger.java
@@ -7,6 +7,7 @@
 import org.slf4j.LoggerFactory;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.cache.RepositoryCache;
 import tech.ydb.yoj.repository.db.exception.EntityAlreadyExistsException;
 import tech.ydb.yoj.repository.ydb.YdbRepository;
@@ -161,12 +162,20 @@ private MergingState doTransition(MergingState state, Statement.QueryType nextQu
 
     @SuppressWarnings("unchecked")
     private static YdbRepository.Query convertInsertToUpsert(YdbRepository.Query<?> query) {
-        return new YdbRepository.Query<>(new UpsertYqlStatement<>(EntitySchema.of(getEntityClass(query))), query.getValues().get(0));
+        var type = getEntityClass(query);
+        var schema = EntitySchema.of(type);
+        var tableDescriptor = TableDescriptor.from(schema);
+        var statement = new UpsertYqlStatement<>(tableDescriptor, schema);
+        return new YdbRepository.Query<>(statement, query.getValues().get(0));
     }
 
     @SuppressWarnings("unchecked")
     private static YdbRepository.Query convertInsertToDelete(YdbRepository.Query<?> query) {
-        return new YdbRepository.Query<>(new DeleteByIdStatement<>(EntitySchema.of(getEntityClass(query))), getEntityId(query));
+        var type = getEntityClass(query);
+        var schema = EntitySchema.of(type);
+        var tableDescriptor = TableDescriptor.from(schema);
+        var statement = new DeleteByIdStatement<>(tableDescriptor, schema);
+        return new YdbRepository.Query<>(statement, getEntityId(query));
     }
 
     private static Entity.Id getEntityId(YdbRepository.Query<?> query) {

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/CountAllStatement.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/CountAllStatement.java
@@ -3,6 +3,7 @@
 import tech.ydb.yoj.databind.schema.ObjectSchema;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlPredicate;
 import tech.ydb.yoj.repository.ydb.yql.YqlStatementPart;
 
@@ -15,8 +16,10 @@
 public class CountAllStatement<ENTITY extends Entity<ENTITY>> extends PredicateStatement<Collection<? extends YqlStatementPart<?>>, ENTITY, Count> {
     private final List<YqlStatementPart<?>> parts;
 
-    public CountAllStatement(EntitySchema<ENTITY> schema, List<YqlStatementPart<?>> parts) {
-        super(schema, ObjectSchema.of(Count.class), parts, YqlPredicate::from);
+    public CountAllStatement(
+            TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema, List<YqlStatementPart<?>> parts
+    ) {
+        super(tableDescriptor, schema, ObjectSchema.of(Count.class), parts, YqlPredicate::from);
         this.parts = parts;
     }
 

diff --git a/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/DeleteAllStatement.java b/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/DeleteAllStatement.java
@@ -2,10 +2,11 @@
 
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 public class DeleteAllStatement<PARAMS, ENTITY extends Entity<ENTITY>> extends YqlStatement<PARAMS, ENTITY, ENTITY> {
-    public DeleteAllStatement(EntitySchema<ENTITY> schema) {
-        super(schema, schema);
+    public DeleteAllStatement(TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema) {
+        super(tableDescriptor, schema, schema);
     }
 
     @Override

diff --git a/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/DeleteByIdStatement.java b/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/DeleteByIdStatement.java
@@ -2,13 +2,14 @@
 
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 import java.util.Map;
 import java.util.function.Function;
 
 public class DeleteByIdStatement<IN, T extends Entity<T>> extends MultipleVarsYqlStatement.Simple<IN, T> {
-    public DeleteByIdStatement(EntitySchema<T> schema) {
-        super(schema);
+    public DeleteByIdStatement(TableDescriptor<T> tableDescriptor, EntitySchema<T> schema) {
+        super(tableDescriptor, schema);
     }
 
     @Override

diff --git a/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindAllYqlStatement.java b/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindAllYqlStatement.java
@@ -1,14 +1,16 @@
 package tech.ydb.yoj.repository.ydb.statement;
 
-import lombok.NonNull;
 import tech.ydb.yoj.databind.schema.Schema;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 public class FindAllYqlStatement<PARAMS, ENTITY extends Entity<ENTITY>, RESULT> extends YqlStatement<PARAMS, ENTITY, RESULT> {
 
-    public FindAllYqlStatement(@NonNull EntitySchema<ENTITY> schema, @NonNull Schema<RESULT> resultSchema) {
-        super(schema, resultSchema);
+    public FindAllYqlStatement(
+            TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema, Schema<RESULT> resultSchema
+    ) {
+        super(tableDescriptor, schema, resultSchema);
     }
 
     @Override

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindInStatement.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindInStatement.java
@@ -16,6 +16,7 @@
 import tech.ydb.yoj.databind.schema.Schema.JavaField;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlListingQuery;
 import tech.ydb.yoj.repository.ydb.yql.YqlPredicate;
 import tech.ydb.yoj.repository.ydb.yql.YqlType;
@@ -128,6 +129,7 @@ public final class FindInStatement<IN, T extends Entity<T>, RESULT> extends Mult
      *                     in the {@code resultSchema}
      */
     public static <ID extends Entity.Id<T>, T extends Entity<T>, RESULT> FindInStatement<Set<ID>, T, RESULT> from(
+            TableDescriptor<T> tableDescriptor,
             EntitySchema<T> schema,
             Schema<RESULT> resultSchema,
             Iterable<ID> ids,
@@ -138,10 +140,13 @@ public static <ID extends Entity.Id<T>, T extends Entity<T>, RESULT> FindInState
         var keySchema = schema.getIdSchema();
         var keyFields = collectKeyFieldsFromIds(schema.getIdSchema(), ids);
 
-        return new FindInStatement<>(schema, resultSchema, keySchema, keyFields, null, filter, orderBy, limit);
+        return new FindInStatement<>(
+                tableDescriptor,schema, resultSchema, keySchema, keyFields, null, filter, orderBy, limit
+        );
     }
 
     public static <K, T extends Entity<T>, RESULT> FindInStatement<Set<K>, T, RESULT> from(
+            TableDescriptor<T> tableDescriptor,
             EntitySchema<T> schema,
             Schema<RESULT> resultSchema,
             String indexName,
@@ -153,10 +158,13 @@ public static <K, T extends Entity<T>, RESULT> FindInStatement<Set<K>, T, RESULT
         Schema<K> keySchema = getKeySchemaFromValues(keys);
         Set<String> keyFields = collectKeyFieldsFromKeys(schema, indexName, keySchema, keys);
 
-        return new FindInStatement<>(schema, resultSchema, keySchema, keyFields, indexName, filter, orderBy, limit);
+        return new FindInStatement<>(
+                tableDescriptor, schema, resultSchema, keySchema, keyFields, indexName, filter, orderBy, limit
+        );
     }
 
     private <PARAMS> FindInStatement(
+            TableDescriptor<T> tableDescriptor,
             EntitySchema<T> schema,
             Schema<RESULT> resultSchema,
             Schema<PARAMS> keySchema,
@@ -167,14 +175,18 @@ private <PARAMS> FindInStatement(
             @Nullable Integer limit
 
     ) {
-        super(schema, resultSchema);
+        super(tableDescriptor, schema, resultSchema);
 
         this.indexName = indexName;
         this.orderBy = orderBy;
         this.limit = limit;
         this.keySchema = keySchema;
         this.keyFields = keyFields;
-        this.predicate = (filter == null) ? null : new PredicateClause<>(schema, YqlListingQuery.toYqlPredicate(filter));
+        if (filter != null) {
+            this.predicate = new PredicateClause<>(tableDescriptor, schema, YqlListingQuery.toYqlPredicate(filter));
+        } else {
+            this.predicate = null;
+        }
 
         validateOrderByFields();
     }
@@ -437,8 +449,8 @@ private boolean hasPredicate() {
     private static class PredicateClause<T extends Entity<T>> extends PredicateStatement<Class<Void>, T, T> {
         private final YqlPredicate predicate;
 
-        public PredicateClause(EntitySchema<T> schema, YqlPredicate predicate) {
-            super(schema, schema, Void.class, __ -> predicate);
+        public PredicateClause(TableDescriptor<T> tableDescriptor, EntitySchema<T> schema, YqlPredicate predicate) {
+            super(tableDescriptor, schema, schema, Void.class, __ -> predicate);
 
             this.predicate = predicate;
         }

diff --git a/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindRangeStatement.java b/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindRangeStatement.java
@@ -7,6 +7,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 tech.ydb.yoj.repository.ydb.yql.YqlType;
 
 import java.util.List;
@@ -23,8 +24,13 @@ public class FindRangeStatement<ENTITY extends Entity<ENTITY>, ID extends Entity
     @Getter
     private final List<YqlStatementParam> params;
 
-    public FindRangeStatement(EntitySchema<ENTITY> schema, Schema<RESULT> outSchema, Range<ID> range) {
-        super(schema, outSchema);
+    public FindRangeStatement(
+            TableDescriptor<ENTITY> tableDescriptor,
+            EntitySchema<ENTITY> schema,
+            Schema<RESULT> outSchema,
+            Range<ID> range
+    ) {
+        super(tableDescriptor, schema, outSchema);
         this.params = Stream.of(RangeBound.values())
                 .flatMap(b -> toParams(b.map(range).keySet(), b))
                 .collect(toList());

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindStatement.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindStatement.java
@@ -4,6 +4,7 @@
 import tech.ydb.yoj.databind.schema.Schema;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlOrderBy;
 import tech.ydb.yoj.repository.ydb.yql.YqlPredicate;
 import tech.ydb.yoj.repository.ydb.yql.YqlStatementPart;
@@ -20,6 +21,7 @@ public class FindStatement<ENTITY extends Entity<ENTITY>, RESULT> extends Predic
     private final List<YqlStatementPart<?>> parts;
 
     public static <E extends Entity<E>, R> FindStatement<E, R> from(
+            @NonNull TableDescriptor<E> tableDescriptor,
             @NonNull EntitySchema<E> schema,
             @NonNull Schema<R> outSchema,
             @NonNull Collection<? extends YqlStatementPart<?>> parts,
@@ -32,16 +34,17 @@ public static <E extends Entity<E>, R> FindStatement<E, R> from(
             }
         }
 
-        return new FindStatement<>(schema, outSchema, partsList, distinct);
+        return new FindStatement<>(tableDescriptor, schema, outSchema, partsList, distinct);
     }
 
     private FindStatement(
+            @NonNull TableDescriptor<ENTITY> tableDescriptor,
             @NonNull EntitySchema<ENTITY> schema,
             @NonNull Schema<RESULT> outSchema,
             @NonNull List<YqlStatementPart<?>> parts,
             boolean distinct
     ) {
-        super(schema, outSchema, parts, YqlPredicate::from);
+        super(tableDescriptor, schema, outSchema, parts, YqlPredicate::from);
         this.distinct = distinct;
         this.parts = parts;
     }

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindYqlStatement.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/FindYqlStatement.java
@@ -1,9 +1,9 @@
 package tech.ydb.yoj.repository.ydb.statement;
 
-import lombok.NonNull;
 import tech.ydb.yoj.databind.schema.Schema;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.db.cache.RepositoryCache;
 import tech.ydb.yoj.repository.ydb.yql.YqlType;
 
@@ -13,8 +13,10 @@
 import static java.util.stream.Collectors.toList;
 
 public class FindYqlStatement<PARAMS, ENTITY extends Entity<ENTITY>, RESULT> extends YqlStatement<PARAMS, ENTITY, RESULT> {
-    public FindYqlStatement(@NonNull EntitySchema<ENTITY> schema, @NonNull Schema<RESULT> resultSchema) {
-        super(schema, resultSchema);
+    public FindYqlStatement(
+            TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema, Schema<RESULT> resultSchema
+    ) {
+        super(tableDescriptor, schema, resultSchema);
     }
 
     @Override

diff --git a/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/InsertYqlStatement.java b/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/InsertYqlStatement.java
@@ -2,13 +2,14 @@
 
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 import java.util.Map;
 import java.util.function.Function;
 
 public class InsertYqlStatement<PARAMS, ENTITY extends Entity<ENTITY>> extends MultipleVarsYqlStatement.Simple<PARAMS, ENTITY> {
-    public InsertYqlStatement(EntitySchema<ENTITY> schema) {
-        super(schema);
+    public InsertYqlStatement(TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema) {
+        super(tableDescriptor, schema);
     }
 
     @Override

diff --git a/...-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/MultipleVarsYqlStatement.java b/...-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/MultipleVarsYqlStatement.java
@@ -4,6 +4,7 @@
 import tech.ydb.yoj.databind.schema.Schema;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlType;
 
 import java.util.Collection;
@@ -21,8 +22,10 @@
 public abstract class MultipleVarsYqlStatement<PARAMS, ENTITY extends Entity<ENTITY>, RESULT> extends YqlStatement<PARAMS, ENTITY, RESULT> {
     public static final String listName = "$Input";
 
-    public MultipleVarsYqlStatement(EntitySchema<ENTITY> schema, Schema<RESULT> resultSchema) {
-        super(schema, resultSchema);
+    public MultipleVarsYqlStatement(
+            TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema, Schema<RESULT> resultSchema
+    ) {
+        super(tableDescriptor, schema, resultSchema);
     }
 
     @Override
@@ -92,8 +95,8 @@ protected String toDebugParams(PARAMS params) {
 
     public abstract static class Simple<PARAMS, ENTITY extends Entity<ENTITY>>
             extends MultipleVarsYqlStatement<PARAMS, ENTITY, ENTITY> {
-        public Simple(EntitySchema<ENTITY> schema) {
-            super(schema, schema);
+        public Simple(TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema) {
+            super(tableDescriptor, schema, schema);
         }
     }
 }
diff --git a/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/PredicateStatement.java b/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/PredicateStatement.java
@@ -10,6 +10,7 @@
 import tech.ydb.yoj.databind.schema.Schema.JavaField;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlCompositeType;
 import tech.ydb.yoj.repository.ydb.yql.YqlPredicate;
 import tech.ydb.yoj.repository.ydb.yql.YqlPredicateParam;
@@ -34,12 +35,13 @@ public abstract class PredicateStatement<PARAMS, ENTITY extends Entity<ENTITY>,
     private final Map<String, PredParam> predParams;
 
     public PredicateStatement(
+            @NonNull TableDescriptor<ENTITY> tableDescriptor,
             @NonNull EntitySchema<ENTITY> schema,
             @NonNull Schema<RESULT> outSchema,
             @NonNull PARAMS params,
             @NonNull Function<PARAMS, YqlPredicate> getPredicate
     ) {
-        super(schema, outSchema);
+        super(tableDescriptor, schema, outSchema);
 
         this.getPredicate = getPredicate;
         YqlPredicate pred = getPredicate.apply(params);

diff --git a/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpdateByIdStatement.java b/...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpdateByIdStatement.java
@@ -3,6 +3,7 @@
 import tech.ydb.proto.ValueProtos;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlType;
 
 import java.util.Collection;
@@ -24,8 +25,10 @@ public final class UpdateByIdStatement<ENTITY extends Entity<ENTITY>, ID extends
 
     private final Set<YqlStatementParam> idParams;
 
-    public UpdateByIdStatement(EntitySchema<ENTITY> schema, UpdateModel.ById<ID> model) {
-        super(schema, schema);
+    public UpdateByIdStatement(
+            TableDescriptor<ENTITY> tableDescriptor, EntitySchema<ENTITY> schema, UpdateModel.ById<ID> model
+    ) {
+        super(tableDescriptor, schema, schema);
         this.idParams = schema.flattenId().stream()
                 .map(c -> YqlStatementParam.required(YqlType.of(c), c.getName()))
                 .collect(toUnmodifiableSet());

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpdateInStatement.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpdateInStatement.java
@@ -8,6 +8,7 @@
 import tech.ydb.yoj.databind.schema.Schema.JavaField;
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 import tech.ydb.yoj.repository.ydb.yql.YqlType;
 
 import java.util.Collection;
@@ -33,11 +34,12 @@ public class UpdateInStatement<T extends Entity<T>, RESULT>
     private final Set<String> keyFields;
 
     public UpdateInStatement(
+            TableDescriptor<T> tableDescriptor,
             EntitySchema<T> schema,
             Schema<RESULT> resultSchema,
             UpdateInStatementInput<T> in
     ) {
-        super(schema, resultSchema);
+        super(tableDescriptor, schema, resultSchema);
 
         this.keyFields = collectKeyFields(in.ids);
         this.values = new HashMap<>(in.values.size());

diff --git a/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpsertYqlStatement.java b/...sitory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/statement/UpsertYqlStatement.java
@@ -2,13 +2,14 @@
 
 import tech.ydb.yoj.repository.db.Entity;
 import tech.ydb.yoj.repository.db.EntitySchema;
+import tech.ydb.yoj.repository.db.TableDescriptor;
 
 import java.util.Map;
 import java.util.function.Function;
 
 public class UpsertYqlStatement<IN, T extends Entity<T>> extends MultipleVarsYqlStatement.Simple<IN, T> {
-    public UpsertYqlStatement(EntitySchema<T> schema) {
-        super(schema);
+    public UpsertYqlStatement(TableDescriptor<T> tableDescriptor, EntitySchema<T> schema) {
+        super(tableDescriptor, schema);
     }
 
     @Override