feat: added scan query hint

ydb-platform · Nov 3, 2024 · 0a9241b · 0a9241b
1 parent f834908
commit 0a9241b
Show file tree

Hide file tree

Showing 6 changed files with 190 additions and 8 deletions.
diff --git a/hibernate-dialect/pom.xml b/hibernate-dialect/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>tech.ydb.dialects</groupId>
     <artifactId>hibernate-ydb-dialect</artifactId>
-    <version>1.0.0</version>
+    <version>1.0.1</version>
 
     <packaging>jar</packaging>
 

diff --git a/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/YdbDialect.java b/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/YdbDialect.java
@@ -56,6 +56,8 @@
 import tech.ydb.hibernate.dialect.exporter.EmptyExporter;
 import tech.ydb.hibernate.dialect.exporter.YdbIndexExporter;
 import tech.ydb.hibernate.dialect.hint.IndexQueryHintHandler;
+import tech.ydb.hibernate.dialect.hint.QueryHintHandler;
+import tech.ydb.hibernate.dialect.hint.ScanQueryHintHandler;
 import tech.ydb.hibernate.dialect.translator.YdbSqlAstTranslatorFactory;
 import tech.ydb.hibernate.dialect.types.InstantJavaType;
 import tech.ydb.hibernate.dialect.types.InstantJdbcType;
@@ -72,6 +74,10 @@ public class YdbDialect extends Dialect {
 
     private static final Exporter<ForeignKey> FOREIGN_KEY_EMPTY_EXPORTER = new EmptyExporter<>();
     private static final Exporter<Constraint> UNIQUE_KEY_EMPTY_EXPORTER = new EmptyExporter<>();
+    private static final List<QueryHintHandler> QUERY_HINT_HANDLERS = List.of(
+            IndexQueryHintHandler.INSTANCE,
+            ScanQueryHintHandler.INSTANCE
+    );
 
     public YdbDialect(DialectResolutionInfo dialectResolutionInfo) {
         super(dialectResolutionInfo);
@@ -144,11 +150,29 @@ public void initializeFunctionRegistry(FunctionContributions functionContributio
     @Override
     public String addSqlHintOrComment(String sql, QueryOptions queryOptions, boolean commentsEnabled) {
         if (queryOptions.getDatabaseHints() != null) {
-            sql = IndexQueryHintHandler.addQueryHints(sql, queryOptions.getDatabaseHints());
+            for (var queryHintHandler : QUERY_HINT_HANDLERS) {
+                sql = queryHintHandler.addQueryHints(sql, queryOptions.getDatabaseHints());
+            }
         }
 
-        if (queryOptions.getComment() != null && IndexQueryHintHandler.commentIsHint(queryOptions.getComment())) {
-            return IndexQueryHintHandler.addQueryHints(sql, List.of(queryOptions.getComment()));
+        if (queryOptions.getComment() != null) {
+            boolean commentIsHint = false;
+
+            var hints = queryOptions.getComment().split(",");
+
+            for (var queryHintHandler : QUERY_HINT_HANDLERS) {
+                for (var hint : hints) {
+                    hint = hint.trim();
+                    if (queryHintHandler.commentIsHint(hint)) {
+                        commentIsHint = true;
+                        sql = queryHintHandler.addQueryHints(sql, List.of(hint));
+                    }
+                }
+            }
+
+            if (commentIsHint) {
+                return sql;
+            }
         }
 
         if (commentsEnabled && queryOptions.getComment() != null) {

diff --git a/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/IndexQueryHintHandler.java b/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/IndexQueryHintHandler.java
@@ -8,17 +8,24 @@
 /**
  * @author Kirill Kurdyukov
  */
-public class IndexQueryHintHandler {
+public class IndexQueryHintHandler implements QueryHintHandler {
+    public static final IndexQueryHintHandler INSTANCE = new IndexQueryHintHandler();
+
     private static final Pattern SELECT_FROM_WHERE_QUERY_PATTERN = Pattern
             .compile("^\\s*(select.+?from\\s+\\w+)(.+where.+)$", Pattern.CASE_INSENSITIVE);
 
-    public static final String HINT_USE_INDEX = "use_index:";
+    private static final String HINT_USE_INDEX = "use_index:";
+
+    private IndexQueryHintHandler() {
+    }
 
-    public static boolean commentIsHint(String comment) {
+    @Override
+    public boolean commentIsHint(String comment) {
         return comment.startsWith(HINT_USE_INDEX);
     }
 
-    public static String addQueryHints(String query, List<String> hints) {
+    @Override
+    public String addQueryHints(String query, List<String> hints) {
         if (hints.isEmpty()) {
             return query;
         }

diff --git a/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/QueryHintHandler.java b/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/QueryHintHandler.java
@@ -0,0 +1,13 @@
+package tech.ydb.hibernate.dialect.hint;
+
+import java.util.List;
+
+/**
+ * @author Kirill Kurdyukov
+ */
+public interface QueryHintHandler {
+
+    String addQueryHints(String query, List<String> hints);
+
+    boolean commentIsHint(String comment);
+}
diff --git a/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/ScanQueryHintHandler.java b/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/ScanQueryHintHandler.java
@@ -0,0 +1,42 @@
+package tech.ydb.hibernate.dialect.hint;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author Kirill Kurdyukov
+ */
+public class ScanQueryHintHandler implements QueryHintHandler {
+    public static final ScanQueryHintHandler INSTANCE = new ScanQueryHintHandler();
+
+    private static final String HINT_USE_SCAN = "use_scan";
+
+    private ScanQueryHintHandler() {
+
+    }
+
+    @Override
+    public String addQueryHints(String query, List<String> hints) {
+        if (hints.isEmpty()) {
+            return query;
+        }
+
+        var useScan = new ArrayList<String>();
+        hints.forEach(hint -> {
+            if (hint.startsWith(HINT_USE_SCAN)) {
+                useScan.add(hint.substring(HINT_USE_SCAN.length()));
+            }
+        });
+
+        if (useScan.size() == 1) {
+            return "scan " + query;
+        }
+
+        return query;
+    }
+
+    @Override
+    public boolean commentIsHint(String comment) {
+        return comment.startsWith(HINT_USE_SCAN);
+    }
+}
diff --git a/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/StudentsRepositoryTest.java b/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/StudentsRepositoryTest.java
@@ -260,4 +260,100 @@ void sumAvgByStudentIdTest() {
                 }
         );
     }
+
+    @Test
+    void useScanQueryHintTest() {
+        /*
+        scan select
+            c1_0.CourseId,
+            c1_0.CourseName
+        from
+            Courses c1_0
+        order by
+            c1_0.CourseId
+         */
+        inTransaction(
+                session -> {
+                    var courses = session.createQuery("FROM Course c ORDER BY c.id", Course.class)
+                            .addQueryHint("use_scan")
+                            .getResultList();
+
+                    checkCourses(courses);
+                }
+        );
+
+        /*
+        scan select
+            c1_0.CourseId,
+            c1_0.CourseName
+        from
+            Courses c1_0
+        order by
+            c1_0.CourseId
+         */
+        inTransaction(
+                session -> {
+                    var courses = session.createQuery("FROM Course c ORDER BY c.id", Course.class)
+                            .setHint(HibernateHints.HINT_COMMENT, "use_scan")
+                            .getResultList();
+
+                    checkCourses(courses);
+                }
+        );
+    }
+
+    private static void checkCourses(List<Course> courses) {
+        assertEquals(6, courses.size());
+        assertEquals("Базы данных", courses.get(0).getName());
+        assertEquals("Управление проектами", courses.get(1).getName());
+        assertEquals("ППО", courses.get(2).getName());
+        assertEquals("Теория информации", courses.get(3).getName());
+        assertEquals("Математический анализ", courses.get(4).getName());
+        assertEquals("Технологии Java", courses.get(5).getName());
+    }
+
+    @Test
+    void useIndexAndUseScanHintsTogetherTest() {
+        /*
+        scan select
+            g1_0.GroupId,
+            g1_0.GroupName
+        from
+            Groups view group_name_index g1_0
+        where
+            g1_0.GroupName='M3439'
+         */
+        inTransaction(
+                session -> {
+                    Group group = session
+                            .createQuery("FROM Group g WHERE g.name = 'M3439'", Group.class)
+                            .addQueryHint("use_index:group_name_index") // Hibernate
+                            .addQueryHint("use_scan")
+                            .getSingleResult();
+
+                    assertEquals("M3439", group.getName());
+                }
+        );
+
+
+        /*
+        scan select
+            g1_0.GroupId,
+            g1_0.GroupName
+        from
+            Groups view group_name_index g1_0
+        where
+            g1_0.GroupName='M3439'
+         */
+        inTransaction(
+                session -> {
+                    Group group = session
+                            .createQuery("FROM Group g WHERE g.name = 'M3439'", Group.class)
+                            .setHint(HibernateHints.HINT_COMMENT, "use_index:group_name_index, use_scan") // JPA
+                            .getSingleResult();
+
+                    assertEquals("M3439", group.getName());
+                }
+        );
+    }
 }