Query multiple tables using one entity (Refactoring) (#96)

ydb-platform · Nov 27, 2024 · f8860f0 · f8860f0
1 parent 12db3e6
commit f8860f0
Show file tree

Hide file tree

Showing 12 changed files with 292 additions and 127 deletions.
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
@@ -0,0 +1,60 @@
+package tech.ydb.yoj.repository.ydb.statement;
+
+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.ydb.yql.YqlPredicate;
+import tech.ydb.yoj.repository.ydb.yql.YqlStatementPart;
+
+import java.util.Collection;
+import java.util.function.Function;
+
+import static java.util.Comparator.comparing;
+import static java.util.stream.Collectors.joining;
+
+class CountAllStatement<ENTITY extends Entity<ENTITY>> extends PredicateStatement<Collection<? extends YqlStatementPart<?>>, ENTITY, Count> {
+    private final Collection<? extends YqlStatementPart<?>> parts;
+
+    public CountAllStatement(
+            EntitySchema<ENTITY> schema,
+            Schema<Count> resultSchema,
+            Collection<? extends YqlStatementPart<?>> parts,
+            Function<Collection<? extends YqlStatementPart<?>>, YqlPredicate> predicateFrom
+    ) {
+        super(schema, resultSchema, parts, predicateFrom);
+        this.parts = parts;
+    }
+
+    public CountAllStatement(
+            EntitySchema<ENTITY> schema,
+            Schema<Count> resultSchema,
+            Collection<? extends YqlStatementPart<?>> parts,
+            Function<Collection<? extends YqlStatementPart<?>>, YqlPredicate> predicateFrom,
+            String tableName
+    ) {
+        super(schema, resultSchema, parts, predicateFrom, tableName);
+        this.parts = parts;
+    }
+
+    @Override
+    public String getQuery(String tablespace) {
+        return declarations()
+                + "SELECT COUNT(*) AS count"
+                + " FROM " + table(tablespace)
+                + " " + mergeParts(parts.stream())
+                .sorted(comparing(YqlStatementPart::getPriority))
+                .map(sp -> sp.toFullYql(schema))
+                .map(this::resolveParamNames)
+                .collect(joining(" "));
+    }
+
+    @Override
+    public QueryType getQueryType() {
+        return QueryType.SELECT;
+    }
+
+    @Override
+    public String toDebugString(Collection<? extends YqlStatementPart<?>> yqlStatementParts) {
+        return "count(" + 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
@@ -0,0 +1,30 @@
+package tech.ydb.yoj.repository.ydb.statement;
+
+import lombok.NonNull;
+import tech.ydb.yoj.repository.db.Entity;
+import tech.ydb.yoj.repository.db.EntitySchema;
+
+class DeleteAllStatement<PARAMS, ENTITY extends Entity<ENTITY>> extends YqlStatement<PARAMS, ENTITY, ENTITY> {
+    public DeleteAllStatement(@NonNull Class<ENTITY> type) {
+        super(EntitySchema.of(type), EntitySchema.of(type));
+    }
+
+    public DeleteAllStatement(@NonNull EntitySchema<ENTITY> schema, String tableName) {
+        super(schema, schema, tableName);
+    }
+
+    @Override
+    public String getQuery(String tablespace) {
+        return "DELETE FROM " + table(tablespace);
+    }
+
+    @Override
+    public QueryType getQueryType() {
+        return QueryType.DELETE_ALL;
+    }
+
+    @Override
+    public String toDebugString(PARAMS params) {
+        return "deleteAll(" + schema.getName() + ")";
+    }
+}
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
@@ -10,6 +10,10 @@ public class DeleteByIdStatement<IN, T extends Entity<T>> extends MultipleVarsYq
         super(type);
     }
 
+    DeleteByIdStatement(Class<T> type, String tableName) {
+        super(type, tableName);
+    }
+
     @Override
     public QueryType getQueryType() {
         return QueryType.DELETE;

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
@@ -135,7 +135,19 @@ protected FindInStatement(
             @Nullable OrderExpression<T> orderBy,
             @Nullable Integer limit
     ) {
-        super(schema, resultSchema);
+        this(schema, resultSchema, ids, filter, orderBy, limit, schema.getName());
+    }
+
+    protected FindInStatement(
+            EntitySchema<T> schema,
+            Schema<RESULT> resultSchema,
+            Iterable<? extends Entity.Id<T>> ids,
+            @Nullable FilterExpression<T> filter,
+            @Nullable OrderExpression<T> orderBy,
+            @Nullable Integer limit,
+            String tableName
+    ) {
+        super(schema, resultSchema, tableName);
 
         this.orderBy = orderBy;
         this.limit = limit;
@@ -160,7 +172,20 @@ protected <V> FindInStatement(
             @Nullable OrderExpression<T> orderBy,
             @Nullable Integer limit
     ) {
-        super(schema, resultSchema);
+        this(schema, resultSchema, indexName, keys, filter, orderBy, limit, schema.getName());
+    }
+
+    protected <V> FindInStatement(
+            EntitySchema<T> schema,
+            Schema<RESULT> resultSchema,
+            String indexName,
+            Iterable<V> keys,
+            @Nullable FilterExpression<T> filter,
+            @Nullable OrderExpression<T> orderBy,
+            @Nullable Integer limit,
+            String tableName
+    ) {
+        super(schema, resultSchema, tableName);
 
         this.indexName = indexName;
         this.orderBy = orderBy;

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
@@ -24,7 +24,11 @@ public class FindRangeStatement<ENTITY extends Entity<ENTITY>, ID extends Entity
     private final List<YqlStatementParam> params;
 
     public FindRangeStatement(EntitySchema<ENTITY> schema, Schema<RESULT> outSchema, Range<ID> range) {
-        super(schema, outSchema);
+        this(schema, outSchema, range, schema.getName());
+    }
+
+    public FindRangeStatement(EntitySchema<ENTITY> schema, Schema<RESULT> outSchema, Range<ID> range, String tableName) {
+        super(schema, outSchema, tableName);
         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
@@ -0,0 +1,52 @@
+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.ydb.yql.YqlPredicate;
+import tech.ydb.yoj.repository.ydb.yql.YqlStatementPart;
+
+import java.util.Collection;
+import java.util.function.Function;
+
+import static java.util.Comparator.comparing;
+import static java.util.stream.Collectors.joining;
+
+public class FindStatement<ENTITY extends Entity<ENTITY>, RESULT> extends PredicateStatement<Collection<? extends YqlStatementPart<?>>, ENTITY, RESULT> {
+    private final boolean distinct;
+    private final Collection<? extends YqlStatementPart<?>> parts;
+
+    public FindStatement(
+            @NonNull EntitySchema<ENTITY> schema,
+            @NonNull Schema<RESULT> outSchema,
+            @NonNull Collection<? extends YqlStatementPart<?>> parts,
+            @NonNull Function<Collection<? extends YqlStatementPart<?>>, YqlPredicate> predicateFrom,
+            boolean distinct,
+            String tableName) {
+        super(schema, outSchema, parts, predicateFrom, tableName);
+        this.distinct = distinct;
+        this.parts = parts;
+    }
+
+    public String getQuery(String tablespace) {
+        return declarations()
+                + "SELECT " + (distinct ? "DISTINCT " : "") + outNames()
+                + " FROM " + table(tablespace)
+                + " " + mergeParts(parts.stream())
+                .sorted(comparing(YqlStatementPart::getPriority))
+                .map(sp -> sp.toFullYql(schema))
+                .map(this::resolveParamNames)
+                .collect(joining(" "));
+    }
+
+    @Override
+    public Statement.QueryType getQueryType() {
+        return Statement.QueryType.SELECT;
+    }
+
+    @Override
+    public String toDebugString(Collection<? extends YqlStatementPart<?>> yqlStatementParts) {
+        return "find(" + yqlStatementParts + ")";
+    }
+}
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
@@ -0,0 +1,67 @@
+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.cache.RepositoryCache;
+import tech.ydb.yoj.repository.ydb.yql.YqlType;
+
+import java.util.Collections;
+import java.util.List;
+
+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(@NonNull EntitySchema<ENTITY> schema, @NonNull Schema<RESULT> resultSchema, String tableName) {
+        super(schema, resultSchema, tableName);
+    }
+
+    @Override
+    public List<YqlStatementParam> getParams() {
+        return schema.flattenId().stream()
+                .map(c -> YqlStatementParam.required(YqlType.of(c), c.getName()))
+                .collect(toList());
+    }
+
+    @Override
+    public String getQuery(String tablespace) {
+        return declarations()
+                + "SELECT " + outNames()
+                + " FROM " + table(tablespace)
+                + " WHERE " + nameEqVars();
+    }
+
+    @Override
+    public List<RESULT> readFromCache(PARAMS params, RepositoryCache cache) {
+        RepositoryCache.Key key = new RepositoryCache.Key(resultSchema.getType(), params);
+        if (!cache.contains(key)) {
+            return null;
+        }
+
+        //noinspection unchecked
+        return cache.get(key)
+                .map(o -> Collections.singletonList((RESULT) o))
+                .orElse(Collections.emptyList());
+    }
+
+    @Override
+    public void storeToCache(PARAMS params, List<RESULT> result, RepositoryCache cache) {
+        RepositoryCache.Key key = new RepositoryCache.Key(resultSchema.getType(), params);
+        cache.put(key, result.stream().findFirst().orElse(null));
+    }
+
+    @Override
+    public QueryType getQueryType() {
+        return QueryType.SELECT;
+    }
+
+    @Override
+    public String toDebugString(PARAMS params) {
+        return "find(" + params + ")";
+    }
+}
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
@@ -10,6 +10,10 @@ public InsertYqlStatement(Class<ENTITY> type) {
         super(type);
     }
 
+    public InsertYqlStatement(Class<ENTITY> type, String tableName) {
+        super(type, tableName);
+    }
+
     @Override
     public QueryType getQueryType() {
         return QueryType.INSERT;

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
@@ -39,7 +39,17 @@ public PredicateStatement(
             @NonNull PARAMS params,
             @NonNull Function<PARAMS, YqlPredicate> getPredicate
     ) {
-        super(schema, outSchema);
+        this(schema, outSchema, params, getPredicate, schema.getName());
+    }
+
+    public PredicateStatement(
+            @NonNull EntitySchema<ENTITY> schema,
+            @NonNull Schema<RESULT> outSchema,
+            @NonNull PARAMS params,
+            @NonNull Function<PARAMS, YqlPredicate> getPredicate,
+            String tableName
+    ) {
+        super(schema, outSchema, tableName);
 
         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
@@ -2,6 +2,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.ydb.yql.YqlType;
 
 import java.util.Collection;
@@ -18,13 +19,17 @@
 import static java.util.stream.Stream.concat;
 
 public final class UpdateByIdStatement<ENTITY extends Entity<ENTITY>, ID extends Entity.Id<ENTITY>>
-        extends YqlStatement.Simple<UpdateModel.ById<ID>, ENTITY> {
+        extends YqlStatement<UpdateModel.ById<ID>, ENTITY, ENTITY> {
     private final Map<String, UpdateSetParam> setParams;
 
     private final Set<YqlStatementParam> idParams;
 
     public UpdateByIdStatement(Class<ENTITY> type, UpdateModel.ById<ID> model) {
-        super(type);
+        this(type, model, EntitySchema.of(type).getName());
+    }
+
+    public UpdateByIdStatement(Class<ENTITY> type, UpdateModel.ById<ID> model, String tableName) {
+        super(EntitySchema.of(type), EntitySchema.of(type), tableName);
         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
@@ -37,7 +37,16 @@ public UpdateInStatement(
             Schema<RESULT> resultSchema,
             UpdateInStatementInput<T> in
     ) {
-        super(schema, resultSchema);
+        this(schema, resultSchema, in, schema.getName());
+    }
+
+    public UpdateInStatement(
+            EntitySchema<T> schema,
+            Schema<RESULT> resultSchema,
+            UpdateInStatementInput<T> in,
+            String tableName
+    ) {
+        super(schema, resultSchema, tableName);
 
         this.keyFields = collectKeyFields(in.ids);
         this.values = new HashMap<>(in.values.size());