如何用JPA Criteria Builder添加CockroachDB的AS OF SYSTEM TIME子句

问题描述 #

需要在动态构建的JPA Criteria Query中添加CockroachDB特有的AS OF SYSTEM TIME follower_read_timestamp()子句——查询会根据用户输入添加大量可选谓词，必须依赖Criteria Builder实现。当前遇到事务问题，官方技术支持建议添加该子句，同时希望了解如何在Criteria Builder中注入这类自定义SQL子句。

依赖版本：jakarta.persistence:jakarta.persistence-api:2.2.3（来自org.springframework.boot:spring-boot-starter-data-jpa:2.7.7的传递依赖）

原代码示例：

// Get criteria builder
CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();

// Create criteria query and root
CriteriaQuery<StorageEntity> criteriaQuery = criteriaBuilder.createQuery(StorageEntity.class);
Root<StorageEntity> root = criteriaQuery.from(StorageEntity.class);

List<Predicate> predicates = new ArrayList<>();

if (CollectionUtils.isNotEmpty(ids)) {
  predicates.add(criteriaBuilder.<List<UUID>>in(root.get("id")).value(ids));
}

/**
A lot of other dynamic predicates are added here which is why I'm using criteria builder
*/

if (CollectionUtils.isNotEmpty(otherIds)) {
  predicates.add(criteriaBuilder.<List<UUID>>in(root.get("otherIds")).value(otherIds));
}

// Get criteria query with filter predicates
criteriaQuery = criteriaQuery.where(predicates.toArray(new Predicate[] {}));

// Create query
TypedQuery<StorageEntity> query = entityManager.createQuery(criteriaQuery);

/**
Ideally here I would be able to add `AS OF SYSTEM TIME <timestamp>` clause cause statement
*/

// Execute query and get results
List<StorageEntity> storageEntities = query.getResultList();

解决方案 #

方案1：通过Hibernate Query API修改生成的SQL #

Spring Boot默认使用Hibernate作为JPA实现，可将TypedQuery转换为Hibernate原生Query对象，获取自动生成的SQL后手动拼接AS OF SYSTEM TIME子句，再通过原生查询执行：

// 保留原Criteria Builder构建逻辑不变
TypedQuery<StorageEntity> typedQuery = entityManager.createQuery(criteriaQuery);

// 转换为Hibernate Query以获取生成的SQL
org.hibernate.query.Query<StorageEntity> hibernateQuery = typedQuery.unwrap(org.hibernate.query.Query.class);
String originalSql = hibernateQuery.getQueryString();

// 拼接AS OF SYSTEM TIME子句（根据实际SQL结构调整替换规则）
// 示例：针对单表查询，在FROM后的表名后追加子句
String modifiedSql = originalSql.replaceFirst("FROM (\\w+)", "FROM $1 AS OF SYSTEM TIME follower_read_timestamp()");

// 创建原生查询并映射结果到实体，同时复制原查询的参数
List<StorageEntity> storageEntities = entityManager.createNativeQuery(modifiedSql, StorageEntity.class)
        .setParameters(hibernateQuery.getParameters())
        .getResultList();

注意：字符串替换逻辑需匹配实际生成的SQL结构，如果涉及多表关联，需要调整替换规则以保证子句添加位置正确。

方案2：自定义Hibernate SQL AST扩展（全局通用方案） #

若需要全局为所有查询添加该子句，可通过自定义Hibernate的SQL AST Visitor实现自动注入：

1. 实现自定义SQL AST Visitor：

public class CockroachAsOfSqlAstVisitor extends StandardSqlAstVisitor {
    public CockroachAsOfSqlAstVisitor(SqlAstTranslator<?> translator) {
        super(translator);
    }

    @Override
    public void visitTableReference(TableReference tableReference) {
        super.visitTableReference(tableReference);
        // 自动为每张表添加AS OF SYSTEM TIME子句
        getTranslator().appendSql(" AS OF SYSTEM TIME follower_read_timestamp()");
    }
}

2. 配置Hibernate使用自定义Visitor（通过Spring Boot配置文件）：

# 自定义SQL AST转换器，需继承CockroachDB原生转换器并替换Visitor
hibernate.query.sql.ast.translator=com.yourpackage.CustomCockroachSqlAstTranslator

注意：该方案需要了解Hibernate内部API，且需保证版本兼容性，Hibernate的SQL AST API可能随版本迭代变更。

方案3：命名查询（不推荐） #

若仅针对固定查询场景，可使用@NamedQuery定义包含AS OF SYSTEM TIME的SQL，但此方案无法利用Criteria Builder的动态谓词构建能力，不适合你的场景。

内容的提问来源于stack exchange，提问作者joshuatvernon