From d60388b23d9994c3b6dfb42612a423fb3e0edcc5 Mon Sep 17 00:00:00 2001
From: Kirill Kurdyukov <kurdyukov-kir@ydb.tech>
Date: Mon, 19 Aug 2024 19:37:04 +0300
Subject: [PATCH] feat: view index on select from

---
 .../tech/ydb/hibernate/dialect/YdbDialect.java   | 16 ++++++++++++++++
 .../student/StudentsRepositoryTest.java          | 15 ++++++++++++++-
 .../tech/ydb/hibernate/student/entity/Group.java |  3 +--
 3 files changed, 31 insertions(+), 3 deletions(-)

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
index b3f549a..30f494b 100644
--- a/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/YdbDialect.java
+++ b/hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/YdbDialect.java
@@ -1,6 +1,8 @@
 package tech.ydb.hibernate.dialect;
 
 import java.time.LocalDateTime;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import org.hibernate.boot.model.FunctionContributions;
 import org.hibernate.boot.model.TypeContributions;
 import org.hibernate.dialect.Dialect;
@@ -62,6 +64,7 @@ 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 Pattern QUERY_PATTERN = Pattern.compile("^\\s*(select.+?from\\s+\\w+)(.+where.+)$", Pattern.CASE_INSENSITIVE);
 
     public YdbDialect(DialectResolutionInfo dialectResolutionInfo) {
         super(dialectResolutionInfo);
@@ -118,6 +121,19 @@ public void initializeFunctionRegistry(FunctionContributions functionContributio
         );
     }
 
+    @Override
+    public String getQueryHintString(String query, String hints) {
+        Matcher matcher = QUERY_PATTERN.matcher(query);
+        if (matcher.matches() && matcher.groupCount() > 1) {
+            String startToken = matcher.group(1);
+            String endToken = matcher.group(2);
+
+            return startToken + " view " + hints + " " + endToken;
+        } else {
+            return query;
+        }
+    }
+
     @Override
     public LimitHandler getLimitHandler() {
         return LimitOffsetLimitHandler.INSTANCE;
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
index de17063..ddc4591 100644
--- a/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/StudentsRepositoryTest.java
+++ b/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/StudentsRepositoryTest.java
@@ -143,7 +143,6 @@ void studentsAndCourses_Lazy_ManyToManyTest() {
         );
     }
 
-
     @Test
     void studentsByGroupName_Lazy_OneToManyTest() {
         inTransaction(
@@ -160,6 +159,20 @@ void studentsByGroupName_Lazy_OneToManyTest() {
         );
     }
 
+    @Test
+    void groupByGroupName_ViewIndex() {
+        inTransaction(
+                session -> {
+                    Group group = session
+                            .createQuery("FROM Group g WHERE g.name = 'M3439'", Group.class)
+                            .addQueryHint("group_name_index")
+                            .getSingleResult();
+
+                    assertEquals("M3439", group.getName());
+                }
+        );
+    }
+
     @Test
     void studentsByGroupName_Eager_OneToManyTest() {
         inTransaction(
diff --git a/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/entity/Group.java b/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/entity/Group.java
index cda1dd8..95b7fce 100644
--- a/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/entity/Group.java
+++ b/hibernate-dialect/src/test/java/tech/ydb/hibernate/student/entity/Group.java
@@ -7,11 +7,10 @@
 import jakarta.persistence.NamedQuery;
 import jakarta.persistence.OneToMany;
 import jakarta.persistence.Table;
+import java.util.List;
 import lombok.Getter;
 import lombok.Setter;
 
-import java.util.List;
-
 /**
  * @author Kirill Kurdyukov
  */