GH-173 Added custom YdbRepositoryFactory infrastructure

spring-data-jdbc-ydb/README.md

@@ -60,7 +60,7 @@ spring.datasource.driver-class-name=tech.ydb.jdbc.YdbDriver
 spring.datasource.url=jdbc:ydb:grpc://localhost:2136/local
 ```
 
-Java configuration for @YdbType annotation:
+Java configuration for `@YdbType` annotation:
 
 ```java
 @Import(AbstractYdbJdbcConfiguration.class)
@@ -76,4 +76,4 @@ See [connect to YDB](../README.md/#connect-to-ydb).
 
 ## Support and Contact
 
-For support, you can open issues in the repository issue tracker with tag `spring-data-jdbc`.
+For support, you can open issues in the repository issue tracker with tag `spring-data-jdbc`.

spring-data-jdbc-ydb/pom.xml

@@ -82,6 +82,7 @@
         <dependency>
             <groupId>org.springframework.data</groupId>
             <artifactId>spring-data-jdbc</artifactId>
+            <version>3.5.0-SNAPSHOT</version>
             <scope>provided</scope>
         </dependency>
         <dependency>
@@ -90,7 +91,6 @@
             <version>${ydb.jdbc.version}</version>
             <scope>provided</scope>
         </dependency>
-
         <dependency>
             <groupId>tech.ydb.test</groupId>
             <artifactId>ydb-junit5-support</artifactId>

spring-data-jdbc-ydb/src/main/java/tech/ydb/data/core/YdbQueryMethod.java

@@ -0,0 +1,36 @@
+package tech.ydb.data.core;
+
+import java.lang.reflect.Method;
+
+import org.springframework.data.jdbc.repository.query.JdbcQueryMethod;
+import org.springframework.data.mapping.context.MappingContext;
+import org.springframework.data.projection.ProjectionFactory;
+import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
+import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
+import org.springframework.data.repository.core.NamedQueries;
+import org.springframework.data.repository.core.RepositoryMetadata;
+import org.springframework.data.repository.query.Parameters;
+import org.springframework.data.repository.query.ParametersSource;
+
+/**
+ * Custom {@link JdbcQueryMethod} implementation specific to YDB.
+ *
+ * @author Mikhail Polivakha
+ */
+public class YdbQueryMethod extends JdbcQueryMethod {
+
+    public YdbQueryMethod(
+      Method method,
+      RepositoryMetadata metadata,
+      ProjectionFactory factory,
+      NamedQueries namedQueries,
+      MappingContext<? extends RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> mappingContext
+    ) {
+        super(method, metadata, factory, namedQueries, mappingContext);
+    }
+
+    @Override
+    protected Parameters<?, ?> createParameters(ParametersSource parametersSource) {
+        return new YdbQueryParameters(parametersSource);
+    }
+}

spring-data-jdbc-ydb/src/main/java/tech/ydb/data/core/YdbQueryParameters.java

@@ -0,0 +1,60 @@
+package tech.ydb.data.core;
+
+import java.lang.annotation.Annotation;
+import java.sql.SQLType;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.data.jdbc.repository.query.JdbcParameter;
+import org.springframework.data.jdbc.repository.query.JdbcParameters;
+import org.springframework.data.relational.repository.query.RelationalParameters;
+import org.springframework.data.repository.query.ParametersSource;
+import org.springframework.data.util.Lazy;
+import org.springframework.data.util.TypeInformation;
+
+import tech.ydb.data.core.convert.YQLType;
+import tech.ydb.data.core.convert.YdbType;
+import tech.ydb.table.values.PrimitiveType;
+
+/**
+ * Parameters of the YDB query. Custom implementation of {@link JdbcParameters}, emerged from requirement
+ * to support {@link YdbType} on method parameters:
+ * <p/>
+ * <pre class="code">
+ * &#064;Query("SELECT * FROM my_table WHERE type = :type")
+ * Optional<MyEntity> findByType(@Param("type") @YdbType("Text") String type);
+ * </pre>
+ *
+ * @author Mikhail Polivakha
+ */
+public class YdbQueryParameters extends JdbcParameters {
+
+    public YdbQueryParameters(ParametersSource parametersSource) {
+        super(parametersSource, methodParameter -> {
+
+            YdbType ydbType = methodParameter.getParameterAnnotation(YdbType.class);
+
+            if (ydbType != null) {
+                YQLType sqlType = new YQLType(PrimitiveType.valueOf(ydbType.value()));
+                return new YdbMethodParameter(
+                  methodParameter,
+                  parametersSource.getDomainTypeInformation(),
+                  sqlType,
+                  Lazy.of(() -> sqlType) // TODO: Do we support the array/collection-like types?
+                );
+            } else {
+                return new YdbMethodParameter(methodParameter, parametersSource.getDomainTypeInformation());
+            }
+        });
+    }
+
+    static class YdbMethodParameter extends JdbcParameter {
+
+        public YdbMethodParameter(MethodParameter parameter, TypeInformation<?> domainType) {
+            super(parameter, domainType);
+        }
+
+        public YdbMethodParameter(MethodParameter parameter, TypeInformation<?> domainType, SQLType sqlType, Lazy<SQLType> actualSqlType) {
+            super(parameter, domainType, sqlType, actualSqlType);
+        }
+    }
+}

spring-data-jdbc-ydb/src/main/java/tech/ydb/data/core/convert/YdbType.java

@@ -5,11 +5,14 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import tech.ydb.table.values.PrimitiveType;
+
 /**
  * @author Madiyar Nurgazin
+ * @author Mikhail Polivakha
  */
 @Retention(RetentionPolicy.RUNTIME)
-@Target(ElementType.FIELD)
+@Target({ElementType.FIELD, ElementType.PARAMETER})
 public @interface YdbType {
     String value();
 }

spring-data-jdbc-ydb/src/main/java/tech/ydb/data/repository/support/YdbRepositoryFactory.java

@@ -0,0 +1,131 @@
+package tech.ydb.data.repository.support;
+
+import java.lang.reflect.Method;
+import java.util.Optional;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.data.jdbc.core.convert.DataAccessStrategy;
+import org.springframework.data.jdbc.core.convert.JdbcConverter;
+import org.springframework.data.jdbc.repository.QueryMappingConfiguration;
+import org.springframework.data.jdbc.repository.query.JdbcQueryMethod;
+import org.springframework.data.jdbc.repository.query.PartTreeJdbcQuery;
+import org.springframework.data.jdbc.repository.query.RowMapperFactory;
+import org.springframework.data.jdbc.repository.query.StringBasedJdbcQuery;
+import org.springframework.data.jdbc.repository.support.BeanFactoryAwareRowMapperFactory;
+import org.springframework.data.jdbc.repository.support.JdbcQueryLookupStrategy;
+import org.springframework.data.jdbc.repository.support.JdbcRepositoryFactory;
+import org.springframework.data.mapping.callback.EntityCallbacks;
+import org.springframework.data.projection.ProjectionFactory;
+import org.springframework.data.relational.core.dialect.Dialect;
+import org.springframework.data.relational.core.mapping.RelationalMappingContext;
+import org.springframework.data.repository.core.NamedQueries;
+import org.springframework.data.repository.core.RepositoryMetadata;
+import org.springframework.data.repository.query.CachingValueExpressionDelegate;
+import org.springframework.data.repository.query.QueryLookupStrategy;
+import org.springframework.data.repository.query.RepositoryQuery;
+import org.springframework.data.repository.query.ValueExpressionDelegate;
+import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
+import org.springframework.lang.Nullable;
+
+import tech.ydb.data.core.YdbQueryMethod;
+
+/**
+ * Custom {@link JdbcRepositoryFactory repository factory} to allow for query tuning.
+ *
+ * @author Mikhail Polivakha
+ */
+public class YdbRepositoryFactory extends JdbcRepositoryFactory {
+
+    private final EntityCallbacks entityCallbacks;
+    private final QueryMappingConfiguration queryMappingConfiguration;
+
+    private static final Log LOG = LogFactory.getLog(JdbcQueryLookupStrategy.class);
+
+    /**
+     * Creates a new {@link JdbcRepositoryFactory} for the given {@link DataAccessStrategy},
+     * {@link RelationalMappingContext} and {@link ApplicationEventPublisher}.
+     *
+     * @param dataAccessStrategy must not be {@literal null}.
+     * @param context            must not be {@literal null}.
+     * @param converter          must not be {@literal null}.
+     * @param dialect            must not be {@literal null}.
+     * @param publisher          must not be {@literal null}.
+     * @param operations         must not be {@literal null}.
+     */
+    public YdbRepositoryFactory(
+      DataAccessStrategy dataAccessStrategy,
+      RelationalMappingContext context,
+      JdbcConverter converter,
+      Dialect dialect,
+      ApplicationEventPublisher publisher,
+      NamedParameterJdbcOperations operations,
+      EntityCallbacks entityCallbacks,
+      QueryMappingConfiguration queryMappingConfiguration
+    ) {
+        super(dataAccessStrategy, context, converter, dialect, publisher, operations);
+
+        this.entityCallbacks = entityCallbacks;
+        this.queryMappingConfiguration = queryMappingConfiguration;
+    }
+
+    @Override
+    protected Optional<QueryLookupStrategy> getQueryLookupStrategy(QueryLookupStrategy.Key key, ValueExpressionDelegate valueExpressionDelegate) {
+        return Optional.of(
+          new YdbQueryLookupStrategy(
+            publisher,
+            entityCallbacks,
+            context,
+            converter,
+            dialect,
+            beanFactory,
+            queryMappingConfiguration,
+            operations,
+            new CachingValueExpressionDelegate(valueExpressionDelegate)
+          )
+        );
+    }
+
+    static class YdbQueryLookupStrategy extends JdbcQueryLookupStrategy {
+
+        private final RowMapperFactory rowMapperFactory;
+
+        public YdbQueryLookupStrategy(
+          ApplicationEventPublisher publisher, @Nullable EntityCallbacks callbacks,
+          RelationalMappingContext context, JdbcConverter converter, Dialect dialect, BeanFactory beanFactory,
+          QueryMappingConfiguration queryMappingConfiguration, NamedParameterJdbcOperations operations,
+          ValueExpressionDelegate delegate
+        ) {
+            super(publisher, callbacks, context, converter, dialect, queryMappingConfiguration, operations, delegate);
+
+            this.rowMapperFactory = new BeanFactoryAwareRowMapperFactory(context, converter, queryMappingConfiguration, callbacks, publisher, beanFactory);
+        }
+
+        @Override
+        public RepositoryQuery resolveQuery(
+          Method method,
+          RepositoryMetadata repositoryMetadata,
+          ProjectionFactory projectionFactory,
+          NamedQueries namedQueries
+        ) {
+            JdbcQueryMethod queryMethod = new YdbQueryMethod(method, repositoryMetadata, projectionFactory, namedQueries, getMappingContext());
+
+            if (namedQueries.hasQuery(queryMethod.getNamedQueryName()) || queryMethod.hasAnnotatedQuery()) {
+
+                if (queryMethod.hasAnnotatedQuery() && queryMethod.hasAnnotatedQueryName()) {
+                    LOG.warn(String.format(
+                      "Query method %s is annotated with both, a query and a query name; Using the declared query", method));
+                }
+
+                String queryString = evaluateTableExpressions(repositoryMetadata, queryMethod.getRequiredQuery());
+
+                return new StringBasedJdbcQuery(queryString, queryMethod, getOperations(), rowMapperFactory, getConverter(),
+                  delegate);
+            } else {
+                return new PartTreeJdbcQuery(getMappingContext(), queryMethod, getDialect(), getConverter(), getOperations(), rowMapperFactory);
+            }
+        }
+    }
+}

spring-data-jdbc-ydb/src/main/java/tech/ydb/data/repository/support/YdbRepositoryFactoryBean.java

@@ -0,0 +1,38 @@
+package tech.ydb.data.repository.support;
+
+import java.io.Serializable;
+
+import org.springframework.data.jdbc.repository.support.JdbcRepositoryFactoryBean;
+import org.springframework.data.repository.Repository;
+import org.springframework.data.repository.core.support.RepositoryFactorySupport;
+
+/**
+ * Custom implementation, specific for YDB Spring Data repositories.
+ *
+ * @author Mikhail Polivakha
+ */
+public class YdbRepositoryFactoryBean<T extends Repository<S, ID>, S, ID extends Serializable> extends JdbcRepositoryFactoryBean<T, S, ID> {
+
+    /**
+     * Creates a new {@link JdbcRepositoryFactoryBean} for the given repository interface.
+     *
+     * @param repositoryInterface must not be {@literal null}.
+     */
+    public YdbRepositoryFactoryBean(Class<? extends T> repositoryInterface) {
+        super(repositoryInterface);
+    }
+
+    @Override
+    protected RepositoryFactorySupport doCreateRepositoryFactory() {
+        return new YdbRepositoryFactory(
+          dataAccessStrategy,
+          mappingContext,
+          converter,
+          dialect,
+          publisher,
+          operations,
+          entityCallbacks,
+          queryMappingConfiguration
+        );
+    }
+}

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/YdbBaseTest.java

@@ -1,5 +1,6 @@
 package tech.ydb.data;
 
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.springframework.boot.test.autoconfigure.data.jdbc.AutoConfigureDataJdbc;
 import org.springframework.boot.test.context.SpringBootTest;

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/YdbJdbcConfiguration.java

@@ -4,7 +4,9 @@
 import org.springframework.context.annotation.Import;
 import org.springframework.data.jdbc.repository.config.EnableJdbcAuditing;
 import org.springframework.data.jdbc.repository.config.EnableJdbcRepositories;
+
 import tech.ydb.data.repository.config.AbstractYdbJdbcConfiguration;
+import tech.ydb.data.repository.support.YdbRepositoryFactoryBean;
 
 /**
  * @author Madiyar Nurgazin
@@ -13,8 +15,11 @@
 @Configuration
 @EnableJdbcRepositories(
   considerNestedRepositories = true,
-  basePackages = "tech.ydb.data"
+  basePackages = "tech.ydb.data",
+  repositoryFactoryBeanClass = YdbRepositoryFactoryBean.class
 )
 @EnableJdbcAuditing
 @Import(AbstractYdbJdbcConfiguration.class)
-public class YdbJdbcConfiguration {}
+public class YdbJdbcConfiguration {
+
+}

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/books/repository/AuthorRepository.java → spring-data-jdbc-ydb/src/test/java/tech/ydb/data/all_types_table/repository/AuthorRepository.java

@@ -1,4 +1,4 @@
-package tech.ydb.data.books.repository;
+package tech.ydb.data.all_types_table.repository;
 
 import java.util.List;
 import org.springframework.data.jdbc.repository.query.Query;

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/books/repository/BookRepository.java → spring-data-jdbc-ydb/src/test/java/tech/ydb/data/all_types_table/repository/BookRepository.java

@@ -1,4 +1,4 @@
-package tech.ydb.data.books.repository;
+package tech.ydb.data.all_types_table.repository;
 
 import java.util.List;
 import java.util.Optional;

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/books/repository/ReviewRepository.java → spring-data-jdbc-ydb/src/test/java/tech/ydb/data/all_types_table/repository/ReviewRepository.java

@@ -1,4 +1,4 @@
-package tech.ydb.data.books.repository;
+package tech.ydb.data.all_types_table.repository;
 
 import java.util.List;
 import org.springframework.data.domain.Pageable;

spring-data-jdbc-ydb/src/test/java/tech/ydb/data/books/RepositoriesIntegrationTest.java

@@ -14,9 +14,9 @@
 import tech.ydb.data.books.entity.Book;
 import tech.ydb.data.books.entity.BookAuthor;
 import tech.ydb.data.books.entity.Review;
-import tech.ydb.data.books.repository.AuthorRepository;
-import tech.ydb.data.books.repository.BookRepository;
-import tech.ydb.data.books.repository.ReviewRepository;
+import tech.ydb.data.all_types_table.repository.AuthorRepository;
+import tech.ydb.data.all_types_table.repository.BookRepository;
+import tech.ydb.data.all_types_table.repository.ReviewRepository;
 
 /**
  * @author Madiyar Nurgazin