diff --git a/.github/workflows/ci-trino-adapter.yaml b/.github/workflows/ci-trino-adapter.yaml new file mode 100644 index 00000000..1b251633 --- /dev/null +++ b/.github/workflows/ci-trino-adapter.yaml @@ -0,0 +1,43 @@ +name: YDB Trino Adapter CI with Maven + +on: + push: + paths: + - 'ydb-trino-adapter/**' + - '.github/workflows/ci-trino-adapter.yaml' + branches: + - main + pull_request: + paths: + - 'ydb-trino-adapter/**' + - '.github/workflows/ci-trino-adapter.yaml' + +env: + MAVEN_ARGS: --batch-mode --update-snapshots -Dstyle.color=always + +jobs: + build: + name: YDB Trino Adapter + runs-on: ubuntu-latest + + strategy: + matrix: + java: [ '25' ] + + steps: + - uses: actions/checkout@v5 + + - name: Set up JDK ${{matrix.java}} + uses: actions/setup-java@v5 + with: + java-version: ${{matrix.java}} + distribution: 'temurin' + cache: maven + + - name: Download Trino Adapter dependencies + working-directory: ./ydb-trino-adapter + run: mvn $MAVEN_ARGS dependency:go-offline + + - name: Build and test Trino Adapter + working-directory: ./ydb-trino-adapter + run: mvn $MAVEN_ARGS clean test diff --git a/ydb-trino-adapter/README.md b/ydb-trino-adapter/README.md index 290c6dd2..3342b263 100644 --- a/ydb-trino-adapter/README.md +++ b/ydb-trino-adapter/README.md @@ -1,57 +1,27 @@ # YDB Trino Adapter -Minimal Trino JDBC connector for YDB. Built on `trino-base-jdbc` and uses the official YDB JDBC driver. +План развития и покрытие Trino connector tests: [ROADMAP.md](ROADMAP.md). -## Build +# Инструкция по сборке ```bash -mvn -f ydb-trino-adapter/pom.xml -DskipTests package -``` - -## Local example with YDB and Trino -`examples/` contains a minimal Docker Compose setup that runs both Trino (479) and YDB locally. +mvn -f pom.xml -DskipTests package +mvn -f pom.xml -DskipTests dependency:copy-dependencies -DincludeScope=runtime -1) Build the plugin and assemble the plugin directory: +mkdir -p docker/trino/plugin +cp target/ydb-trino-0.1.0.jar docker/trino/plugin +cp target/dependency/*.jar docker/trino/plugin -```bash -mvn -f ydb-trino-adapter/pom.xml -DskipTests package -mvn -f ydb-trino-adapter/pom.xml -DskipTests dependency:copy-dependencies -DincludeScope=runtime - -mkdir -p ydb-trino-adapter/examples/trino/plugin -cp ydb-trino-adapter/target/ydb-trino-adapter-0.1.0.jar ydb-trino-adapter/examples/trino/plugin/ -cp ydb-trino-adapter/target/dependency/*.jar ydb-trino-adapter/examples/trino/plugin/ +cd docker +docker-compose down +docker-compose up -d ``` -2) Start Trino + YDB: +## Запуск Trino CLI ```bash -cd ydb-trino-adapter/examples -docker compose up -d -``` -3) Open Trino UI: `http://localhost:8080` - -4) Connect to Trino: - -```bash -docker exec -it ydb-trino-trino trino -``` - -Check schemas and tables: - -```sql -SHOW SCHEMAS FROM ydb; -SHOW TABLES FROM ydb.default; -``` - -## Connector configuration - -Catalog file: `examples/trino/etc/catalog/ydb.properties`: - -``` -connector.name=ydb -connection-url=jdbc:ydb:grpc://ydb-local:2136/local -``` -If you need a token or service account, add parameters to the JDBC URL as supported by the YDB JDBC driver. +docker exec -it ydb-trino trino +``` \ No newline at end of file diff --git a/ydb-trino-adapter/ROADMAP.md b/ydb-trino-adapter/ROADMAP.md new file mode 100644 index 00000000..bd7eb67a --- /dev/null +++ b/ydb-trino-adapter/ROADMAP.md @@ -0,0 +1,134 @@ +# YDB Trino Adapter — roadmap + +Статус относительно Trino `BaseConnectorTest` / `BaseConnectorSmokeTest` +(ветка `ydb-trino-17-07`, ~314 тестов: ~187 green, ~127 skipped). + +Skipped почти всегда означают `hasBehavior(...)=false` или явный `NOT_SUPPORTED` +в `YdbClient`, а не случайный флаaky. Пустые `@Override` (тест «проходит», но ничего +не проверяет) отмечены отдельно — это тоже долг. + +## Уже есть + +- CREATE / DROP TABLE, INSERT (non-transactional), SELECT +- Predicate pushdown, в том числе по `varchar` (`FULL_PUSHDOWN`, bind через JDBC → `Text`/`Utf8`) +- LIMIT / TopN pushdown (включая TopN по varchar) +- Базовый набор скалярных типов: bool, int*, float/double, decimal, date, timestamp, varchar/Text +- Агрегации и часть expression rewrite (arithmetic, `IN`, string ops, …) + +## Phase 0 — честный тест-долг (быстро) + +Сейчас часть тестов green из‑за пустого override. Либо починить, либо явно +документировать/оставить skip с причиной. + +| Тест / тема | Проблема | +|---|---| +| Long table / column names | Лимиты имён YDB | +| Negative dates / year-of-era | Нет отрицательных дат в YQL | +| `testCharVarcharComparison` | CHAR без pad → не семантика Trino CHAR | +| `testVarcharCastToDateInPredicate` | Cast/pushdown не поддержан | +| Row-level UPDATE declaration / `testRowLevelUpdate` | Planner падает до `NOT_SUPPORTED` | +| `testInsertForDefaultColumn` | Нет default columns | + +**Критерий готовности:** нет «пустых» overrides без комментария «unsupported by design». + +## Phase 1 — DDL, которое уже есть в YQL (высокий ROI) + +YDB умеет `ALTER TABLE ... ADD/DROP COLUMN`, `SET/DROP NOT NULL`. В адаптере это +сейчас выключено. + +1. **DROP COLUMN** — снять `SUPPORTS_DROP_COLUMN=false`, реализовать в `YdbClient` +2. **ADD COLUMN** (без comment / position) — базовая добавка nullable-колонок +3. **DROP / SET NOT NULL** — если поведение совпадёт с ожиданиями Trino-тестов + +Ожидаемый эффект: разблокировка пачки `testDrop*Column`, `testAddAndDropColumnName`, +части not-null тестов (~10–20 кейсов). + +**Вне scope phase 1:** rename column, `SET DATA TYPE`, column/table comments, +`ADD COLUMN ... NOT NULL` с backfill-семантикой Trino, `WITH POSITION`. + +## Phase 2 — UPDATE / DELETE (без MERGE) + +В YQL есть `UPDATE` / `DELETE`. Коннектор сейчас бросает `MODIFYING_ROWS_MESSAGE`. + +1. Простой `DELETE` / `UPDATE` с pushdown предикатов (в т.ч. varchar через `?`) +2. Сложные предикаты из BaseConnectorTest (LIKE, subquery, semi-join) — по мере готовности +3. **MERGE не целиться** — в YDB нет Trino-MERGE; оставляем `SUPPORTS_MERGE=false` + (или позже эмулировать через UPSERT, отдельным решением) + +Ожидаемый эффект: до ~30–40 тестов из группы delete/update (без merge-сюиты). + +**Риски:** семантика транзакций, `testRollback` / `testInsertInTransaction`, +row-level update planner quirks, written stats. + +## Phase 3 — schema as path (дизайн) + +Trino `CREATE SCHEMA` ≠ SQL schema в YDB. Схемы естественно мапятся на директории +в path БД. + +Варианты: + +- оставить один schema `ydb` (как сейчас) — просто и предсказуемо; +- мапить `schema` → subdirectory + реализовать create/drop/rename directory. + +Без явного дизайн-решения флаги `SUPPORTS_CREATE_SCHEMA` / rename / cascade +не включать. Cascade + views/MV — отдельно. + +## Phase 4 — типы контейнеров + +| Trino | YDB | Статус | +|---|---|---| +| `ARRAY` | `List` | не замаплено → skip insert/array/field-in-array | +| `MAP` | `Dict` | не замаплено | +| `ROW` | `Struct` | не замаплено → skip row-field + projection pushdown по nested | + +Нужны read/write mappings, predicate/projection pushdown, тесты data-mapping. +Крупный объём, лучше отдельными PR по типу. + +## Phase 5 — Views / comments / прочее + +- **Views:** в YDB есть ограниченная поддержка; Trino VIEW + metadata-тесты — + отдельный трек. Materialized / federated MV — низкий приоритет. +- **COMMENT ON TABLE/COLUMN:** семантика Trino не совпадает 1:1 с YDB table/column + properties — не блокирует core DML/DDL. +- **RENAME COLUMN / SET COLUMN TYPE:** проверить актуальные возможности YQL; + сейчас считаем unsupported. +- **RENAME TABLE across schemas:** зависит от phase 3 (path move). + +## Порядок PR (предложение) + +```text +0. ROADMAP + подчистить пустые overrides (docs/honesty) +1. DROP COLUMN (+ минимальные ADD COLUMN) +2. UPDATE / DELETE (простые предикаты) → расширять предикаты +3. Design note: schema-as-path (да/нет) → реализация или явный отказ +4. List / Dict / Struct mappings по одному типу +5. Views / comments по необходимости продукта +``` + +## Как мерить прогресс + +После каждого PR: + +```bash +cd ydb-trino-adapter +mvn test +# смотреть surefire: Tests run / Skipped / Failures +``` + +Целевые ориентиры (грубо): + +| Milestone | Skipped (ориентир) | +|---|---| +| Сейчас | ~127 | +| После phase 1 | ~110 | +| После phase 2 | ~70–80 | +| После phase 4 | заметно ниже за счёт array/map/row | + +Точные числа зависят от того, сколько тестов завязано на комбинации флагов +(например MERGE останется большим блоком skip). + +## Не делать + +- Включать `hasBehavior=true` без реализации в `YdbClient` / QueryBuilder +- Ослаблять CI workflow, чтобы «позеленеть» +- Эмулировать MERGE «лишь бы тесты» без явной семантики и документации diff --git a/ydb-trino-adapter/examples/docker-compose.yml b/ydb-trino-adapter/examples/docker-compose.yml index c47b22a8..c0bcdc18 100644 --- a/ydb-trino-adapter/examples/docker-compose.yml +++ b/ydb-trino-adapter/examples/docker-compose.yml @@ -6,6 +6,9 @@ services: container_name: ydb-trino-trino ports: - "8080:8080" + - "5005:5005" + environment: + - JAVA_OPTS=-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005 volumes: - ./trino/etc:/etc/trino - ./trino/plugin:/usr/lib/trino/plugin/ydb @@ -26,4 +29,4 @@ services: ports: - "2135:2135" - "2136:2136" - - "8765:8765" + - "8765:8765" \ No newline at end of file diff --git a/ydb-trino-adapter/pom.xml b/ydb-trino-adapter/pom.xml index cf435027..e99f4f71 100644 --- a/ydb-trino-adapter/pom.xml +++ b/ydb-trino-adapter/pom.xml @@ -5,37 +5,14 @@ 4.0.0 tech.ydb.dialects - ydb-trino-adapter + ydb-trino 0.1.0 YDB Trino Adapter Trino JDBC connector for YDB - https://github.com/ydb-platform/ydb-java-dialects jar - - - Kirill Kurdyukov - kurdyukov-kir@ydb.tech - YDB - https://ydb.tech/ - - - - - https://github.com/ydb-platform/ydb-java-dialects - scm:git:https://github.com/ydb-platform/ydb-java-dialects.git - scm:git:https://github.com/ydb-platform/ydb-java-dialects.git - - - - - Apache License, Version 2.0 - https://www.apache.org/licenses/LICENSE-2.0 - - - UTF-8 25 @@ -92,12 +69,42 @@ ydb-junit5-support ${ydb.sdk.version} test + + + org.apache.logging.log4j + log4j-slf4j-impl + + + + + org.testcontainers + testcontainers + 1.20.0 + test + + + org.apache.logging.log4j + log4j-slf4j-impl + + + + + org.assertj + assertj-core + 3.25.3 + test io.trino trino-testing ${trino.version} test + + + org.apache.logging.log4j + log4j-slf4j-impl + + io.trino @@ -110,6 +117,12 @@ trino-main ${trino.version} test + + + org.apache.logging.log4j + log4j-slf4j-impl + + io.trino @@ -127,16 +140,20 @@ + + org.apache.maven.plugins + maven-compiler-plugin + 3.13.0 + + true + lines,vars,source + + org.apache.maven.plugins maven-surefire-plugin 3.1.0 - - true - localhost:2136 - /local - --add-modules jdk.incubator.vector diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteDivideModulus.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteDivideModulus.java new file mode 100644 index 00000000..9cc514e8 --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteDivideModulus.java @@ -0,0 +1,63 @@ +package tech.ydb.trino; + +import io.trino.matching.Captures; +import io.trino.matching.Pattern; +import io.trino.plugin.base.expression.ConnectorExpressionRule; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.spi.expression.Call; + +import java.util.Optional; + +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName; +import static io.trino.spi.expression.StandardFunctions.DIVIDE_FUNCTION_NAME; +import static io.trino.spi.expression.StandardFunctions.MODULUS_FUNCTION_NAME; +import static java.lang.String.format; + +/** + * Rewrite $divide(a, b), $modulus(a, b) as a / b, a % b, + * but only if b is a non-zero constant. + *

+ * This is because YDB, unlike Trino, suppresses arithmetic errors (including division by zero), and any + * non-constant expression may evaluate to zero and cause a different result if pushed down to YDB. + */ +public class RewriteDivideModulus implements ConnectorExpressionRule { + private final Pattern PATTERN; + + public RewriteDivideModulus() { + this.PATTERN = call() + .with(functionName().matching(name -> name.equals(DIVIDE_FUNCTION_NAME) || name.equals(MODULUS_FUNCTION_NAME))) + .with(argumentCount().matching(count -> count == 2)) + .matching((Call call, RewriteContext _) -> + call.getArguments().stream().noneMatch(arg -> arg instanceof Call)); + } + + @Override + public Pattern getPattern() { + return PATTERN; + } + + @Override + public Optional rewrite(Call call, Captures captures, RewriteContext context) { + String operator = getOperator(call.getFunctionName()); + return RewriteUtils.rewriteBinaryExpression( + call, + context, + () -> call.getArguments().get(1) instanceof io.trino.spi.expression.Constant rightConstant && + rightConstant.getValue() instanceof Number number && + number.longValue() != 0, + (left, right) -> format("(%s) %s (%s)", left, operator, right) + ); + } + + private String getOperator(io.trino.spi.expression.FunctionName functionName) { + if (functionName.equals(DIVIDE_FUNCTION_NAME)) { + return "/"; + } + if (functionName.equals(MODULUS_FUNCTION_NAME)) { + return "%"; + } + return null; + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteNullIf.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteNullIf.java new file mode 100644 index 00000000..30181aea --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteNullIf.java @@ -0,0 +1,45 @@ +package tech.ydb.trino; + +import io.trino.matching.Captures; +import io.trino.matching.Pattern; +import io.trino.plugin.base.expression.ConnectorExpressionRule; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.spi.expression.Call; +import io.trino.spi.expression.StandardFunctions; + +import java.util.Optional; + +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName; +import static java.lang.String.format; + +/** + * Rewrite $nullif(a, b), as: + *
+ * CASE WHEN a = b THEN NULL ELSE a END. + */ +public class RewriteNullIf implements ConnectorExpressionRule { + private final Pattern PATTERN; + + public RewriteNullIf() { + this.PATTERN = call() + .with(functionName().matching(name -> name.equals(StandardFunctions.NULLIF_FUNCTION_NAME))) + .with(argumentCount().matching(count -> count == 2)); + } + + @Override + public Pattern getPattern() { + return PATTERN; + } + + @Override + public Optional rewrite(Call call, Captures captures, RewriteContext context) { + return RewriteUtils.rewriteBinaryExpression( + call, + context, + () -> true, + (left, right) -> format("CASE WHEN %s = %s THEN NULL ELSE %s END", left, right, left) + ); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteStringPosition.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteStringPosition.java new file mode 100644 index 00000000..b39bf5ed --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteStringPosition.java @@ -0,0 +1,86 @@ +package tech.ydb.trino; + +import com.google.common.collect.ImmutableList; +import io.trino.matching.Capture; +import io.trino.matching.Captures; +import io.trino.matching.Pattern; +import io.trino.plugin.base.projection.ProjectFunctionRule; +import io.trino.plugin.jdbc.JdbcExpression; +import io.trino.plugin.jdbc.QueryParameter; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.expression.Call; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.type.BigintType; +import org.jspecify.annotations.NonNull; + +import java.util.Optional; + +import static io.trino.matching.Capture.newCapture; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argument; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.expression; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName; +import static java.lang.String.format; + +/** + * Rewrite strpos(str, sub), as: + *
+ * CASE WHEN Find(str, sub) IS NULL THEN 0 ELSE Find(str, sub) + 1 END. + *

+ * The + 1 is due to the fact that Trino uses 0-based indexing while YDB uses 1-based indexing. + */ +public class RewriteStringPosition implements ProjectFunctionRule { + private static final Capture STRING = newCapture(); + private static final Capture SUBSTRING = newCapture(); + private static final String STRPOS_OPERATOR = "strpos"; + + private static final Pattern PATTERN = call() + .with(functionName().matching(name -> STRPOS_OPERATOR.equals(name.getName()))) + .with(argumentCount().equalTo(2)) + .with(argument(0).matching(expression().capturedAs(STRING))) + .with(argument(1).matching(expression().capturedAs(SUBSTRING))); + + @Override + public Pattern getPattern() { + return PATTERN; + } + + @Override + public Optional rewrite( + ConnectorTableHandle handle, + ConnectorExpression projectionExpression, + Captures captures, + RewriteContext context + ) { + ConnectorExpression stringExpr = captures.get(STRING); + ConnectorExpression substringExpr = captures.get(SUBSTRING); + + Optional rewrittenString = context.rewriteExpression(stringExpr); + Optional rewrittenSubstring = context.rewriteExpression(substringExpr); + + if (rewrittenString.isEmpty() || rewrittenSubstring.isEmpty()) { + return Optional.empty(); + } + + String strSql = rewrittenString.get().expression(); + String subSql = rewrittenSubstring.get().expression(); + + ImmutableList.Builder<@NonNull QueryParameter> parameters = ImmutableList.builder(); + parameters.addAll(rewrittenString.get().parameters()); + // Add substring parameters twice - for both occurrences in CASE expression + parameters.addAll(rewrittenSubstring.get().parameters()); + parameters.addAll(rewrittenSubstring.get().parameters()); + + String findExpr = format("Find(%s, %s)", strSql, subSql); + String expression = format("CASE WHEN (%s) IS NULL THEN 0 ELSE (%s) + 1 END", findExpr, findExpr); + + JdbcExpression result = new JdbcExpression( + expression, + parameters.build(), + YdbTypeUtils.toTypeHandle(BigintType.BIGINT).orElseThrow() + ); + return Optional.of(result); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUnaryStringOperations.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUnaryStringOperations.java new file mode 100644 index 00000000..5f4536e6 --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUnaryStringOperations.java @@ -0,0 +1,82 @@ +package tech.ydb.trino; + +import com.google.common.collect.ImmutableList; +import io.trino.matching.Capture; +import io.trino.matching.Captures; +import io.trino.matching.Pattern; +import io.trino.plugin.base.projection.ProjectFunctionRule; +import io.trino.plugin.jdbc.JdbcExpression; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.expression.Call; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.expression.FunctionName; +import io.trino.spi.type.VarcharType; + +import java.util.Objects; +import java.util.Optional; + +import static io.trino.matching.Capture.newCapture; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argument; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.expression; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName; +import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type; + +public class RewriteUnaryStringOperations implements ProjectFunctionRule { + private static final Capture VALUE = newCapture(); + + private static final Pattern PATTERN = call() + .with(functionName().matching(name -> name.equals(new FunctionName("trim")) || + name.equals(new FunctionName("upper")) || + name.equals(new FunctionName("lower")))) + .with(type().matching(type -> type instanceof VarcharType)) + .with(argumentCount().equalTo(1)) + .with(argument(0).matching(expression().capturedAs(VALUE))); + + @Override + public Pattern getPattern() { + return PATTERN; + } + + @Override + public Optional rewrite(ConnectorTableHandle handle, ConnectorExpression projectionExpression, Captures captures, RewriteContext context) { + JdbcTypeHandle varcharTypeHandle = YdbTypeUtils.toTypeHandle(VarcharType.VARCHAR).orElse(null); + if (Objects.isNull(varcharTypeHandle)) { + return Optional.empty(); + } + + Call call = (Call) projectionExpression; + + ConnectorExpression valueExpr; + if (call.getArguments().size() == 1) { + valueExpr = call.getArguments().getFirst(); + } else { + return Optional.empty(); + } + + Optional rewrittenValue = context.rewriteExpression(valueExpr); + if (rewrittenValue.isEmpty()) { + return Optional.empty(); + } + + String expression = mapToYql(call.getFunctionName().getName()).formatted(rewrittenValue.get().expression()); + + return Optional.of(new JdbcExpression( + expression, + ImmutableList.copyOf(rewrittenValue.get().parameters()), + varcharTypeHandle + )); + } + + private static String mapToYql(String functionName) { + return switch (functionName) { + case "upper" -> "Unicode::ToUpper(%s)"; + case "lower" -> "Unicode::ToLower(%s)"; + case "trim" -> "String::Strip(%s)"; + default -> throw new IllegalArgumentException("Unexpected function name: " + functionName); + }; + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUtils.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUtils.java new file mode 100644 index 00000000..9b68a86c --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/RewriteUtils.java @@ -0,0 +1,40 @@ +package tech.ydb.trino; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.base.expression.ConnectorExpressionRule.RewriteContext; +import io.trino.plugin.jdbc.QueryParameter; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.spi.expression.Call; +import io.trino.spi.expression.ConnectorExpression; +import org.jspecify.annotations.NonNull; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.function.BiFunction; +import java.util.function.Supplier; + +public class RewriteUtils { + public static Optional rewriteBinaryExpression( + Call call, + RewriteContext context, + Supplier condition, + // leftSql, rightSql -> resultSql + BiFunction queryCombiner + ) { + if (!condition.get()) { + return Optional.empty(); + } + List sqls = new ArrayList<>(); + ImmutableList.Builder<@NonNull QueryParameter> parameters = ImmutableList.builder(); + for (ConnectorExpression connectorExpression : call.getArguments()) { + Optional expression = context.defaultRewrite(connectorExpression); + if (expression.isEmpty()) { + return Optional.empty(); + } + parameters.addAll(expression.get().parameters()); + sqls.add(expression.get().expression()); + } + return Optional.of(new ParameterizedExpression(queryCombiner.apply(sqls.get(0), sqls.get(1)), parameters.build())); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java index 6218efbe..aa478e82 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java @@ -1,141 +1,326 @@ package tech.ydb.trino; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; +import io.opentelemetry.api.internal.StringUtils; +import io.trino.plugin.base.aggregation.AggregateFunctionRewriter; +import io.trino.plugin.base.aggregation.AggregateFunctionRule; +import io.trino.plugin.base.expression.ConnectorExpressionRewriter; import io.trino.plugin.base.mapping.IdentifierMapping; -import io.trino.plugin.jdbc.BaseJdbcClient; -import io.trino.plugin.jdbc.BaseJdbcConfig; -import io.trino.plugin.jdbc.ColumnMapping; -import io.trino.plugin.jdbc.ConnectionFactory; -import io.trino.plugin.jdbc.JdbcTypeHandle; -import static io.trino.plugin.jdbc.PredicatePushdownController.DISABLE_PUSHDOWN; -import io.trino.plugin.jdbc.QueryBuilder; -import io.trino.plugin.jdbc.RemoteTableName; -import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping; -import static io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.booleanWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.charReadFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.charWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.dateWriteFunctionUsingLocalDate; -import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping; -import static io.trino.plugin.jdbc.StandardColumnMappings.doubleWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.integerColumnMapping; -import static io.trino.plugin.jdbc.StandardColumnMappings.integerWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.longDecimalWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.realColumnMapping; -import static io.trino.plugin.jdbc.StandardColumnMappings.realWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.shortDecimalWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.smallintColumnMapping; -import static io.trino.plugin.jdbc.StandardColumnMappings.smallintWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.timestampWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.tinyintWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryWriteFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.varcharReadFunction; -import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction; -import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling; -import static io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR; -import io.trino.plugin.jdbc.WriteMapping; +import io.trino.plugin.base.projection.ProjectFunctionRewriter; +import io.trino.plugin.base.projection.ProjectFunctionRule; +import io.trino.plugin.jdbc.*; +import io.trino.plugin.jdbc.aggregation.ImplementAvgDecimal; +import io.trino.plugin.jdbc.aggregation.ImplementAvgFloatingPoint; +import io.trino.plugin.jdbc.aggregation.ImplementCount; +import io.trino.plugin.jdbc.aggregation.ImplementCountAll; +import io.trino.plugin.jdbc.aggregation.ImplementCountDistinct; +import io.trino.plugin.jdbc.aggregation.ImplementMinMax; +import io.trino.plugin.jdbc.aggregation.ImplementSum; +import io.trino.plugin.jdbc.expression.JdbcConnectorExpressionRewriterBuilder; +import io.trino.plugin.jdbc.expression.ParameterizedExpression; +import io.trino.plugin.jdbc.expression.RewriteIn; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import io.trino.spi.connector.*; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.type.*; import io.trino.spi.TrinoException; -import io.trino.spi.connector.ConnectorSession; + +import jakarta.annotation.Nullable; +import org.jspecify.annotations.NonNull; + +import java.sql.*; +import java.util.List; +import java.util.*; +import java.util.Locale; +import java.util.function.BiFunction; +import java.util.stream.Stream; + +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static io.trino.plugin.jdbc.PredicatePushdownController.DISABLE_PUSHDOWN; +import static io.trino.plugin.jdbc.PredicatePushdownController.FULL_PUSHDOWN; +import static io.trino.plugin.jdbc.StandardColumnMappings.*; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.connector.ConnectorMetadata.MODIFYING_ROWS_MESSAGE; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; -import io.trino.spi.type.CharType; -import static io.trino.spi.type.CharType.createCharType; import static io.trino.spi.type.DateType.DATE; -import io.trino.spi.type.DecimalType; +import static io.trino.spi.type.DecimalType.createDecimalType; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.RealType.REAL; import static io.trino.spi.type.SmallintType.SMALLINT; import static io.trino.spi.type.TimestampType.TIMESTAMP_MICROS; -import static io.trino.spi.type.TimestampType.TIMESTAMP_SECONDS; import static io.trino.spi.type.TinyintType.TINYINT; -import io.trino.spi.type.Type; -import io.trino.spi.type.VarbinaryType; -import io.trino.spi.type.VarcharType; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static io.trino.spi.type.VarcharType.createVarcharType; +import static java.lang.Math.max; import static java.lang.String.format; -import java.sql.Connection; -import java.sql.DatabaseMetaData; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.sql.Types; -import java.util.Collection; -import java.util.List; -import java.util.Optional; +import static java.util.stream.Collectors.joining; + +public class YdbClient extends BaseJdbcClient { + private static final String YDB_SCHEMA = "ydb"; + + private final ConnectorExpressionRewriter connectorExpressionRewriter; + private final AggregateFunctionRewriter aggregateFunctionRewriter; + private final ProjectFunctionRewriter projectFunctionRewriter; -public class YdbClient - extends BaseJdbcClient { @Inject public YdbClient( BaseJdbcConfig config, ConnectionFactory connectionFactory, QueryBuilder queryBuilder, IdentifierMapping identifierMapping, - RemoteQueryModifier remoteQueryModifier) { - super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, remoteQueryModifier, true); + RemoteQueryModifier remoteQueryModifier + ) { + super( + "`", + connectionFactory, + queryBuilder, + config.getJdbcTypesMappedToVarchar(), + identifierMapping, + remoteQueryModifier, + true + ); + + this.connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder() + .addStandardRules(this::quoted) + .add(new RewriteIn()) + .add(new RewriteDivideModulus()) + .add(new RewriteNullIf()) + .withTypeClass("integer_type", ImmutableSet.of("tinyint", "smallint", "integer", "bigint")) + .withTypeClass("numeric_type", ImmutableSet.of("tinyint", "smallint", "integer", "bigint", "decimal", "real", "double")) + .withTypeClass("comparable_type", ImmutableSet.of( + "tinyint", "smallint", "integer", "bigint", "decimal", "real", "double", "varchar", "char", "date", "timestamp")) + .map("$equal(left, right)").to("left = right") + .map("$not_equal(left, right)").to("left <> right") + .map("$add(left: integer_type, right: integer_type)").to("left + right") + .map("$subtract(left: integer_type, right: integer_type)").to("left - right") + .map("$multiply(left: integer_type, right: integer_type)").to("left * right") + .map("$negate(value: integer_type)").to("-value") + .map("$less_than(left: comparable_type, right: comparable_type)").to("left < right") + .map("$less_than_or_equal(left: comparable_type, right: comparable_type)").to("left <= right") + .map("$greater_than(left: comparable_type, right: comparable_type)").to("left > right") + .map("$greater_than_or_equal(left: comparable_type, right: comparable_type)").to("left >= right") + .map("$is_null(value)").to("value IS NULL") + .map("$not($is_null(value))").to("value IS NOT NULL") + .map("$concat(left: varchar, right: varchar)").to("left || right") + .build(); + + this.projectFunctionRewriter = new ProjectFunctionRewriter<>( + this.connectorExpressionRewriter, + ImmutableSet.>builder() + .add(new RewriteUnaryStringOperations()) + .add(new RewriteStringPosition()) + .build()); + + JdbcTypeHandle bigintTypeHandle = YdbTypeUtils.toTypeHandle(BIGINT).orElseThrow(); + this.aggregateFunctionRewriter = new AggregateFunctionRewriter<>( + this.connectorExpressionRewriter, + ImmutableSet.>builder() + .add(new ImplementCountAll(bigintTypeHandle)) + .add(new ImplementMinMax(true)) + .add(new ImplementCount(bigintTypeHandle)) + .add(new ImplementCountDistinct(bigintTypeHandle, true)) + .add(new ImplementSum(YdbTypeUtils::toTypeHandle)) + .add(new ImplementAvgFloatingPoint()) + .add(new ImplementAvgDecimal()) + .build()); + } + + @Override + public Optional implementAggregation( + ConnectorSession session, + AggregateFunction aggregate, + Map assignments + ) { + return aggregateFunctionRewriter.rewrite(session, aggregate, assignments); + } + + @Override + public Optional convertPredicate( + ConnectorSession session, + ConnectorExpression expression, + Map assignments + ) { + return connectorExpressionRewriter.rewrite(session, expression, assignments); + } + + @Override + public Optional convertProjection( + ConnectorSession session, + JdbcTableHandle handle, + ConnectorExpression expression, + Map assignments + ) { + JdbcTypeHandle typeHandle = YdbTypeUtils.toTypeHandle(expression.getType()).orElse(null); + if (Objects.isNull(typeHandle)) { + return Optional.empty(); + } + Optional result = connectorExpressionRewriter.rewrite(session, expression, assignments); + return result.map(parameterizedExpression -> new JdbcExpression( + parameterizedExpression.expression(), + parameterizedExpression.parameters(), + typeHandle)).or(() -> projectFunctionRewriter.rewrite(session, handle, expression, assignments)); } @Override public Collection listSchemas(Connection connection) { - return List.of("ydb"); + return ImmutableSet.of(YDB_SCHEMA); + } + + @Override + protected String escapeObjectNameForMetadataQuery(String name, String escape) { + return name; } @Override - public ResultSet getTables(Connection connection, Optional schemaName, Optional tableName) - throws SQLException { - return connection.getMetaData() - .getTables(null, null, tableName.orElse(null), null); + public List getTableNames(ConnectorSession session, Optional schema) { + try (Connection connection = connectionFactory.openConnection(session)) { + try (ResultSet resultSet = getTables(connection, Optional.empty(), Optional.empty())) { + ImmutableList.Builder<@NonNull SchemaTableName> list = ImmutableList.builder(); + while (resultSet.next()) { + String tableName = resultSet.getString("TABLE_NAME"); + list.add(new SchemaTableName(YDB_SCHEMA, tableName)); + } + return list.build(); + } + } catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } } @Override - protected ResultSet getColumns(RemoteTableName remoteTableName, DatabaseMetaData metadata) - throws SQLException { - return metadata.getColumns(null, null, remoteTableName.getTableName(), null); + public Optional getTableHandle( + ConnectorSession session, + SchemaTableName schemaTableName + ) { + try (Connection connection = connectionFactory.openConnection(session)) { + RemoteTableName remoteTableName = toRemoteTableName(schemaTableName); + try (ResultSet columns = getColumns(remoteTableName, connection.getMetaData())) { + if (!columns.next()) { + return Optional.empty(); + } + } + return Optional.of(new JdbcTableHandle( + new SchemaTableName(YDB_SCHEMA, schemaTableName.getTableName()), + remoteTableName, + Optional.empty()) + ); + } catch (SQLException e) { + return Optional.empty(); + } } @Override - public Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) { + public Optional toColumnMapping( + ConnectorSession session, + Connection connection, + JdbcTypeHandle typeHandle + ) { Optional mapping = getForcedMappingToVarchar(typeHandle); if (mapping.isPresent()) { return mapping; } - switch (typeHandle.jdbcType()) { - case Types.SMALLINT: - return Optional.of(smallintColumnMapping()); - case Types.INTEGER: - return Optional.of(integerColumnMapping()); + // YDB JDBC reports text columns as Bytes/String/Utf8/Text under various JDBC type codes. + // Always map them to unbounded varchar (YDB Text has no useful fixed length). + String jdbcTypeName = typeHandle.jdbcTypeName().orElse("").toLowerCase(Locale.ROOT); + if (jdbcTypeName.equals("bytes") + || jdbcTypeName.equals("string") + || jdbcTypeName.equals("utf8") + || jdbcTypeName.equals("text")) { + return Optional.of(unboundedVarcharColumnMapping()); + } - case Types.BIGINT: - return Optional.of(bigintColumnMapping()); + Optional columnMapping = switch (typeHandle.jdbcType()) { + case Types.BIT, Types.BOOLEAN -> Optional.of(booleanColumnMapping()); + case Types.TINYINT, Types.SMALLINT -> Optional.of(smallintColumnMapping()); + case Types.INTEGER -> Optional.of(integerColumnMapping()); + case Types.BIGINT -> Optional.of(bigintColumnMapping()); + case Types.REAL -> Optional.of(realColumnMapping()); + case Types.FLOAT, Types.DOUBLE -> Optional.of(doubleColumnMapping()); + case Types.DECIMAL -> { + // We need this hack because JDBC client for some reason does not return the scale + // in .requiredDecimalDigits() and we have to parse it manually. + String typeName = typeHandle.jdbcTypeName().get(); + int start = typeName.indexOf('('); + int end = typeName.indexOf(')'); + String[] parts = typeName.substring(start + 1, end).split(","); + int precision = Integer.parseInt(parts[0].trim()); + int scale = Integer.parseInt(parts[1].trim()); - case Types.REAL: - return Optional.of(realColumnMapping()); + yield Optional.of(ColumnMapping.mapping( + createDecimalType(precision, max(scale, 0)), + decimalColumnMapping(createDecimalType(precision, max(scale, 0))).getReadFunction(), + decimalColumnMapping(createDecimalType(precision, max(scale, 0))).getWriteFunction(), + DISABLE_PUSHDOWN)); + } + case Types.CHAR, Types.NCHAR -> { + String typeName = typeHandle.jdbcTypeName().orElseThrow(); + int length = typeName.toLowerCase().startsWith("char(") + ? Integer.parseInt(typeName.substring(5, typeName.length() - 1)) + : typeHandle.columnSize().orElse(VarcharType.MAX_LENGTH); + yield Optional.of(varcharColumnMapping(length)); + } + case Types.VARCHAR, Types.LONGVARCHAR, Types.NVARCHAR -> { + String typeName = typeHandle.jdbcTypeName().orElseThrow(); + int length = typeName.toLowerCase().startsWith("varchar(") + ? Integer.parseInt(typeName.substring(8, typeName.length() - 1)) + : typeHandle.columnSize().orElse(VarcharType.MAX_LENGTH); + yield Optional.of(varcharColumnMapping(length)); + } + case Types.DATE -> Optional.of(dateColumnMapping()); + case Types.TIMESTAMP -> Optional.of(timestampColumnMapping()); + default -> Optional.empty(); + }; - case Types.DOUBLE: - return Optional.of(doubleColumnMapping()); + if (columnMapping.isPresent()) { + return columnMapping; + } - case Types.CHAR: - return Optional.of(charColumnMapping(typeHandle.requiredColumnSize())); + return mapToUnboundedVarchar(typeHandle); + } - case Types.VARCHAR: - return Optional.of(varcharColumnMapping(typeHandle.requiredColumnSize())); - } + private static ColumnMapping unboundedVarcharColumnMapping() { + VarcharType varcharType = createUnboundedVarcharType(); + return ColumnMapping.sliceMapping( + varcharType, + varcharReadFunction(varcharType), + varcharWriteFunction(), + FULL_PUSHDOWN); + } - if (getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR) { - return mapToUnboundedVarchar(typeHandle); - } + private static ColumnMapping varcharColumnMapping(int varcharLength) { + VarcharType varcharType = varcharLength <= VarcharType.MAX_LENGTH + ? createVarcharType(varcharLength) + : createUnboundedVarcharType(); + return ColumnMapping.sliceMapping( + varcharType, + varcharReadFunction(varcharType), + varcharWriteFunction(), + FULL_PUSHDOWN); + } - return Optional.empty(); + private static ColumnMapping dateColumnMapping() { + return ColumnMapping.longMapping( + DATE, + dateReadFunctionUsingLocalDate(), + dateWriteFunctionUsingLocalDate()); + } + + private static ColumnMapping timestampColumnMapping() { + return ColumnMapping.longMapping( + TIMESTAMP_MICROS, + timestampReadFunction(TIMESTAMP_MICROS), + timestampWriteFunction(TIMESTAMP_MICROS)); } @Override public WriteMapping toWriteMapping(ConnectorSession session, Type type) { if (type == BOOLEAN) { - return WriteMapping.booleanMapping("Bool", booleanWriteFunction()); + return WriteMapping.booleanMapping("Bool", BooleanWriteFunction.of(Types.BOOLEAN, PreparedStatement::setBoolean)); } if (type == TINYINT) { return WriteMapping.longMapping("Int8", tinyintWriteFunction()); @@ -157,49 +342,236 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type) { } if (type instanceof DecimalType decimalType) { String dataType = format("Decimal(%s, %s)", decimalType.getPrecision(), decimalType.getScale()); - if (decimalType.isShort()) { - return WriteMapping.longMapping(dataType, shortDecimalWriteFunction(decimalType)); - } - return WriteMapping.objectMapping(dataType, longDecimalWriteFunction(decimalType)); + return decimalType.isShort() + ? WriteMapping.longMapping(dataType, shortDecimalWriteFunction(decimalType)) + : WriteMapping.objectMapping(dataType, longDecimalWriteFunction(decimalType)); } - if (type instanceof CharType || type instanceof VarcharType) { - // The String type replaces the types VARCHAR, BLOB, CLOB, and others from other DBMSs. - return WriteMapping.sliceMapping("String", varcharWriteFunction()); + if (type instanceof VarcharType) { + return WriteMapping.sliceMapping("Text", varcharWriteFunction()); } - if (type instanceof VarbinaryType) { - // Strings of an arbitrary length. The length is not limited - return WriteMapping.sliceMapping("String", varbinaryWriteFunction()); + if (type instanceof CharType) { + return WriteMapping.sliceMapping("Text", charWriteFunction()); } if (type == DATE) { return WriteMapping.longMapping("Date", dateWriteFunctionUsingLocalDate()); } - if (type == TIMESTAMP_SECONDS) { - return WriteMapping.longMapping("DateTime", timestampWriteFunction(TIMESTAMP_MICROS)); + if (type == TIMESTAMP_MICROS) { + return WriteMapping.longMapping("Timestamp", timestampWriteFunction(TIMESTAMP_MICROS)); } throw new TrinoException(NOT_SUPPORTED, "Unsupported column type: " + type); } - private static ColumnMapping charColumnMapping(int charLength) { - if (charLength > CharType.MAX_LENGTH) { - return varcharColumnMapping(charLength); + @Override + public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, List sortOrder) { + return true; + } + + @Override + @SuppressWarnings("all") + protected Optional topNFunction() { + // We need the hack below because Trino and YDB handle nulls differently when sorting. + return Optional.of((query, sortItems, limit) -> { + String orderBy = sortItems.stream() + .flatMap(sortItem -> { + String columnName = quoted(sortItem.column().getColumnName()); + SortOrder sortOrder = sortItem.sortOrder(); + // NULLS FIRST: CASE WHEN col IS NULL THEN 0 ELSE 1 END ASC + // NULLS LAST: CASE WHEN col IS NULL THEN 1 ELSE 0 END ASC + if (sortOrder.isNullsFirst()) { + // Add null-sorting key first, then the actual column + String nullSort = format("CASE WHEN %s IS NULL THEN 0 ELSE 1 END ASC", columnName); + String valueSort = format("%s %s", columnName, sortOrder.isAscending() ? "ASC" : "DESC"); + return Stream.of(nullSort, valueSort); + } else { + // NULLS LAST + String nullSort = format("CASE WHEN %s IS NULL THEN 1 ELSE 0 END ASC", columnName); + String valueSort = format("%s %s", columnName, sortOrder.isAscending() ? "ASC" : "DESC"); + return Stream.of(nullSort, valueSort); + } + }) + .collect(joining(", ")); + return format("%s ORDER BY %s LIMIT %d", query, orderBy, limit); + }); + } + + @Override + public boolean isTopNGuaranteed(ConnectorSession session) { + return true; + } + + @Override + protected Optional> limitFunction() { + return Optional.of((sql, limit) -> sql + " LIMIT " + limit); + } + + @Override + public boolean isLimitGuaranteed(ConnectorSession session) { + return true; + } + + private RemoteTableName toRemoteTableName(SchemaTableName schemaTableName) { + return new RemoteTableName(Optional.empty(), Optional.empty(), schemaTableName.getTableName()); + } + + @Override + protected String quoted(@Nullable String catalog, @Nullable String schema, String table) { + // YDB doesn't use catalog & schema in table names, only the table path + return quoted(table); + } + + @Override + protected void execute(ConnectorSession session, Connection connection, String query) throws SQLException { + YdbRetryUtils.withRetry(() -> super.execute(session, connection, query)); + } + + @Override + public boolean supportsRetries() { + // Disable Trino-retries to avoid temporary tables. + return false; + } + + @Override + public void createSchema(ConnectorSession session, String schemaName) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating schemas"); + } + + @Override + public void dropSchema(ConnectorSession session, String schemaName, boolean cascade) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping schemas"); + } + + @Override + public void renameSchema(ConnectorSession session, String schemaName, String newSchemaName) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support renaming schemas"); + } + + @Override + public void dropColumn(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping columns"); + } + + @Override + public void setColumnType(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column, Type type) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting column types"); + } + + @Override + public void dropNotNullConstraint(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping a not null constraint"); + } + + @Override + public void truncateTable(ConnectorSession session, JdbcTableHandle handle) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support truncating tables"); + } + + @Override + public void renameTable(ConnectorSession session, JdbcTableHandle handle, SchemaTableName newTableName) { + SchemaTableName currentName = handle.asPlainTable().getSchemaTableName(); + if (!currentName.getSchemaName().equalsIgnoreCase(newTableName.getSchemaName())) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support renaming tables across schemas"); } - CharType charType = createCharType(charLength); - return ColumnMapping.sliceMapping( - charType, - charReadFunction(charType), - charWriteFunction(), - DISABLE_PUSHDOWN); + super.renameTable(session, handle, newTableName); } - private static ColumnMapping varcharColumnMapping(int varcharLength) { - VarcharType varcharType = varcharLength <= VarcharType.MAX_LENGTH - ? createVarcharType(varcharLength) - : createUnboundedVarcharType(); - return ColumnMapping.sliceMapping( - varcharType, - varcharReadFunction(varcharType), - varcharWriteFunction(), - DISABLE_PUSHDOWN); + @Override + protected void renameTable( + ConnectorSession session, + Connection connection, + String catalogName, + String remoteSchemaName, + String remoteTableName, + String newRemoteSchemaName, + String newRemoteTableName + ) throws SQLException { + // YDB rename is table-path only; ignore catalog/schema in the SQL. + execute(session, connection, format( + "ALTER TABLE %s RENAME TO %s", + quoted(remoteTableName), + quoted(newRemoteTableName))); + } + + @Override + protected String getColumnDefinitionSql(ConnectorSession session, ColumnMetadata column, String columnName) { + // YDB restriction, does not support column comments. + if (!StringUtils.isNullOrEmpty(column.getComment())) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with column comment"); + } + + StringBuilder sb = new StringBuilder() + .append(quoted(columnName)) + .append(" ") + .append(toWriteMapping(session, column.getType()).getDataType()); + + if (!column.isNullable()) { + sb.append(" NOT NULL"); + } + if (column.getDefaultValue().isPresent()) { + sb.append(" DEFAULT ").append(column.getDefaultValue().get()); + } + + return sb.toString(); + } + + @Override + protected void addColumn( + ConnectorSession session, + Connection connection, + RemoteTableName table, + ColumnMetadata column + ) throws SQLException { + if (!column.isNullable()) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support adding not null columns"); + } + String columnName = column.getName(); + String remoteColumnName = getIdentifierMapping().toRemoteColumnName(getRemoteIdentifiers(connection), columnName); + String sql = format( + "ALTER TABLE %s ADD %s", + quoted(table), + getColumnDefinitionSql(session, column, remoteColumnName)); + execute(session, connection, sql); + } + + @Override + protected void renameColumn( + ConnectorSession session, + Connection connection, + RemoteTableName remoteTableName, + String remoteColumnName, + String newRemoteColumnName + ) throws SQLException { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support renaming columns"); + } + + // TODO maybe support in the future :) + @Override + public JdbcMergeTableHandle beginMerge( + ConnectorSession session, + JdbcTableHandle handle, + Map> updateColumnHandles, + List rollbackActions, + RetryMode retryMode + ) { + throw new TrinoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + + // TODO maybe support in the future :) + @Override + public OptionalLong delete( + ConnectorSession session, + JdbcTableHandle handle + ) { + throw new TrinoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + + @Override + public OptionalLong update(ConnectorSession session, JdbcTableHandle handle) { + throw new TrinoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + + @Override + protected List createTableSqls(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) { + return super.createTableSqls(remoteTableName, columns, tableMetadata); } -} +} \ No newline at end of file diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java index 6af159e2..53b2da8c 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java @@ -5,30 +5,54 @@ import com.google.inject.Provides; import com.google.inject.Scopes; import com.google.inject.Singleton; +import io.trino.plugin.base.mapping.IdentifierMapping; import io.trino.plugin.jdbc.BaseJdbcConfig; import io.trino.plugin.jdbc.ConnectionFactory; -import io.trino.plugin.jdbc.DecimalModule; import io.trino.plugin.jdbc.DriverConnectionFactory; import io.trino.plugin.jdbc.ForBaseJdbc; import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcMetadataFactory; +import io.trino.plugin.jdbc.QueryBuilder; import io.trino.plugin.jdbc.credential.CredentialProvider; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; import tech.ydb.jdbc.YdbDriver; -public class YdbClientModule - implements Module -{ +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; + +public class YdbClientModule implements Module { + @Override - public void configure(Binder binder) - { - binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(YdbClient.class).in(Scopes.SINGLETON); - binder.install(new DecimalModule()); + public void configure(Binder binder) { + newOptionalBinder(binder, JdbcMetadataFactory.class) + .setBinding() + .to(YdbMetadataFactory.class) + .in(Scopes.SINGLETON); + + binder.bind(YdbConnector.class).in(Scopes.SINGLETON); + } + + @Provides + @Singleton + @ForBaseJdbc + public JdbcClient provideJdbcClient( + BaseJdbcConfig config, + ConnectionFactory connectionFactory, + QueryBuilder queryBuilder, + IdentifierMapping identifierMapping, + RemoteQueryModifier remoteQueryModifier) { + return new YdbClient(config, connectionFactory, queryBuilder, identifierMapping, remoteQueryModifier); } @Provides @Singleton @ForBaseJdbc - public static ConnectionFactory createConnectionFactory(BaseJdbcConfig config, CredentialProvider credentialProvider) - { - return DriverConnectionFactory.builder(new YdbDriver(), config.getConnectionUrl(), credentialProvider).build(); + public static ConnectionFactory createConnectionFactory( + BaseJdbcConfig config, + CredentialProvider credentialProvider) { + return DriverConnectionFactory.builder( + new YdbDriver(), + config.getConnectionUrl(), + credentialProvider) + .build(); } } diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnector.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnector.java new file mode 100644 index 00000000..ad99f3b9 --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnector.java @@ -0,0 +1,44 @@ +package tech.ydb.trino; + +import com.google.inject.Inject; +import io.airlift.bootstrap.LifeCycleManager; +import io.trino.plugin.base.session.SessionPropertiesProvider; +import io.trino.plugin.jdbc.JdbcConnector; +import io.trino.plugin.jdbc.JdbcTransactionManager; +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.spi.connector.*; +import io.trino.spi.function.table.ConnectorTableFunction; +import io.trino.spi.procedure.Procedure; + +import java.util.Optional; +import java.util.Set; + +import static com.google.common.collect.Sets.immutableEnumSet; +import static io.trino.spi.connector.ConnectorCapabilities.NOT_NULL_COLUMN_CONSTRAINT; + +@SuppressWarnings("all") +public class YdbConnector extends JdbcConnector { + @Inject + public YdbConnector( + LifeCycleManager lifeCycleManager, + ConnectorSplitManager jdbcSplitManager, + ConnectorPageSourceProvider jdbcPageSourceProvider, + ConnectorPageSinkProvider jdbcPageSinkProvider, + Optional accessControl, + Set procedures, + Set connectorTableFunctions, + Set sessionProperties, + Set tableProperties, + JdbcTransactionManager transactionManager) + { + super(lifeCycleManager, jdbcSplitManager, jdbcPageSourceProvider, jdbcPageSinkProvider, + accessControl, procedures, connectorTableFunctions, sessionProperties, + tableProperties, transactionManager); + } + + @Override + public Set getCapabilities() { + // DEFAULT values are not reliably exposed via information_schema yet. + return immutableEnumSet(NOT_NULL_COLUMN_CONSTRAINT); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnectorFactory.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnectorFactory.java new file mode 100644 index 00000000..b7a68ac5 --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbConnectorFactory.java @@ -0,0 +1,38 @@ +package tech.ydb.trino; + +import com.google.inject.Injector; +import com.google.inject.Module; +import io.airlift.bootstrap.Bootstrap; +import io.trino.plugin.base.ConnectorContextModule; +import io.trino.plugin.jdbc.JdbcModule; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorContext; +import io.trino.spi.connector.ConnectorFactory; + +import java.util.Map; +import java.util.function.Supplier; + +public record YdbConnectorFactory(String name, Supplier module) implements ConnectorFactory { + + @Override + public String getName() { + return name; + } + + @Override + public Connector create(String catalogName, Map requiredConfig, ConnectorContext context) { + Bootstrap app = new Bootstrap( + "io.trino.bootstrap.catalog." + catalogName, + new ConnectorContextModule(catalogName, context), + new JdbcModule(), + module.get()); + + Injector injector = app + .doNotInitializeLogging() + .disableSystemProperties() + .setRequiredConfigurationProperties(requiredConfig) + .initialize(); + + return injector.getInstance(YdbConnector.class); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadata.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadata.java new file mode 100644 index 00000000..4389043a --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadata.java @@ -0,0 +1,18 @@ +package tech.ydb.trino; + +import io.trino.plugin.jdbc.DefaultJdbcMetadata; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcQueryEventListener; +import io.trino.plugin.jdbc.TimestampTimeZoneDomain; + +import java.util.Set; + +public class YdbMetadata extends DefaultJdbcMetadata { + public YdbMetadata( + JdbcClient jdbcClient, + TimestampTimeZoneDomain timestampTimeZoneDomain, + Set jdbcQueryEventListeners + ) { + super(jdbcClient, timestampTimeZoneDomain, false, jdbcQueryEventListeners); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadataFactory.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadataFactory.java new file mode 100644 index 00000000..faa9960d --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMetadataFactory.java @@ -0,0 +1,33 @@ +package tech.ydb.trino; + +import com.google.inject.Inject; +import io.trino.plugin.base.cache.identity.IdentityCacheMapping; +import io.trino.plugin.jdbc.DefaultJdbcMetadataFactory; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcMetadata; +import io.trino.plugin.jdbc.JdbcQueryEventListener; +import io.trino.plugin.jdbc.TimestampTimeZoneDomain; + +import java.util.Set; + +public class YdbMetadataFactory extends DefaultJdbcMetadataFactory { + private final TimestampTimeZoneDomain timestampTimeZoneDomain; + private final Set jdbcQueryEventListeners; + + @Inject + public YdbMetadataFactory( + JdbcClient jdbcClient, + TimestampTimeZoneDomain timestampTimeZoneDomain, + Set jdbcQueryEventListeners, + IdentityCacheMapping identityCacheMapping + ) { + super(jdbcClient, timestampTimeZoneDomain, jdbcQueryEventListeners, identityCacheMapping); + this.timestampTimeZoneDomain = timestampTimeZoneDomain; + this.jdbcQueryEventListeners = jdbcQueryEventListeners; + } + + @Override + protected JdbcMetadata create(JdbcClient transactionCachingJdbcClient) { + return new YdbMetadata(transactionCachingJdbcClient, timestampTimeZoneDomain, jdbcQueryEventListeners); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbPlugin.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbPlugin.java index 0fd0a057..5ef33be0 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbPlugin.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbPlugin.java @@ -1,12 +1,28 @@ package tech.ydb.trino; -import io.trino.plugin.jdbc.JdbcPlugin; - -public class YdbPlugin - extends JdbcPlugin -{ - public YdbPlugin() - { - super("ydb", YdbClientModule::new); +import com.google.common.collect.ImmutableList; +import com.google.inject.Module; +import io.trino.plugin.jdbc.credential.CredentialProviderModule; +import io.trino.spi.Plugin; +import io.trino.spi.connector.ConnectorFactory; + +import static io.airlift.configuration.ConfigurationAwareModule.combine; + +public record YdbPlugin(Module module) implements Plugin { + private static final String NAME = "ydb"; + + public YdbPlugin() { + this(new YdbClientModule()); + } + + @Override + public Iterable getConnectorFactories() { + return ImmutableList.of(new YdbConnectorFactory( + NAME, + () -> combine( + new CredentialProviderModule(), + module + ) + )); } } diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbRetryUtils.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbRetryUtils.java new file mode 100644 index 00000000..010a7ecb --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbRetryUtils.java @@ -0,0 +1,49 @@ +package tech.ydb.trino; + +import io.trino.spi.TrinoException; + +import java.sql.SQLException; +import java.time.Duration; + +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; + +public final class YdbRetryUtils { + private static final int MAX_RETRIES = 4; + private static final int BASE_DELAY_MILLIS = 20; + private YdbRetryUtils() { + + } + + @FunctionalInterface + public interface SqlRunnable { + void run() throws SQLException; + } + + public static void withRetry(SqlRunnable action) throws SQLException { + SQLException lastException = null; + for (int attempt = 0; attempt < MAX_RETRIES; attempt++) { + try { + action.run(); + return; + } catch (SQLException e) { + lastException = e; + String message = e.getMessage(); + + if (message != null) { + Duration delay = Duration.ofMillis(BASE_DELAY_MILLIS).multipliedBy(1L << attempt); + + try { + Thread.sleep(delay.toMillis()); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new TrinoException(JDBC_ERROR, "InterruptedException", ie); + } + } else { + throw e; + } + } + } + + throw new TrinoException(JDBC_ERROR, lastException); + } +} diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTypeUtils.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTypeUtils.java new file mode 100644 index 00000000..344d9f4b --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTypeUtils.java @@ -0,0 +1,42 @@ +package tech.ydb.trino; + +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Type; + +import java.sql.Types; +import java.util.Optional; + +public final class YdbTypeUtils { + + private YdbTypeUtils() { + + } + + public static Optional toTypeHandle(Type type) { + return switch (type) { + case io.trino.spi.type.BooleanType _ -> + Optional.of(new JdbcTypeHandle(Types.BOOLEAN, Optional.of("Bool"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.TinyintType _ -> + Optional.of(new JdbcTypeHandle(Types.TINYINT, Optional.of("Int8"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.SmallintType _ -> + Optional.of(new JdbcTypeHandle(Types.SMALLINT, Optional.of("Int16"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.IntegerType _ -> + Optional.of(new JdbcTypeHandle(Types.INTEGER, Optional.of("Int32"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.BigintType _ -> + Optional.of(new JdbcTypeHandle(Types.BIGINT, Optional.of("Int64"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.RealType _ -> + Optional.of(new JdbcTypeHandle(Types.REAL, Optional.of("Float"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.DoubleType _ -> + Optional.of(new JdbcTypeHandle(Types.DOUBLE, Optional.of("Double"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.DateType _ -> + Optional.of(new JdbcTypeHandle(Types.DATE, Optional.of("Date"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case io.trino.spi.type.TimestampType _ -> + Optional.of(new JdbcTypeHandle(Types.TIMESTAMP, Optional.of("Timestamp"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + case DecimalType decimalType -> + Optional.of(new JdbcTypeHandle(Types.DECIMAL, Optional.of("Decimal"), Optional.of(decimalType.getPrecision()), Optional.of(decimalType.getScale()), Optional.empty(), Optional.empty())); + default -> + Optional.of(new JdbcTypeHandle(Types.VARCHAR, Optional.of("String"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())); + }; + } +} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorSmokeTest.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorSmokeTest.java new file mode 100644 index 00000000..f9cf8b27 --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorSmokeTest.java @@ -0,0 +1,85 @@ +package tech.ydb.trino; + +import io.trino.testing.BaseConnectorSmokeTest; +import io.trino.testing.QueryRunner; +import io.trino.testing.TestingConnectorBehavior; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import tech.ydb.test.junit5.YdbHelperExtension; + +public class TestYdbConnectorSmokeTest extends BaseConnectorSmokeTest { + @RegisterExtension + static final YdbHelperExtension ydb = new YdbHelperExtension(); + + @Override + protected QueryRunner createQueryRunner() throws Exception { + return YdbQueryRunner.builder(ydb) + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); + } + + @Override + protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { + return switch (connectorBehavior) { + case SUPPORTS_MERGE, + SUPPORTS_UPDATE, + SUPPORTS_DELETE, + SUPPORTS_CREATE_VIEW, + SUPPORTS_CREATE_SCHEMA, + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_SET_COLUMN_TYPE, + SUPPORTS_DROP_COLUMN, + SUPPORTS_ROW_TYPE, + SUPPORTS_RENAME_COLUMN, + SUPPORTS_ROW_LEVEL_UPDATE, + SUPPORTS_TRUNCATE, + SUPPORTS_COMMENT_ON_COLUMN, + SUPPORTS_COMMENT_ON_TABLE, + SUPPORTS_DROP_SCHEMA_CASCADE, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT, + SUPPORTS_CREATE_TABLE_WITH_TABLE_COMMENT, + SUPPORTS_ADD_COLUMN_WITH_COMMENT, + SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_CREATE_FEDERATED_MATERIALIZED_VIEW, + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_ARRAY, + SUPPORTS_MAP_TYPE, + SUPPORTS_DEFAULT_COLUMN_VALUE, + SUPPORTS_SET_DEFAULT_COLUMN_VALUE, + SUPPORTS_DROP_DEFAULT_COLUMN_VALUE, + SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, + SUPPORTS_DROP_NOT_NULL_CONSTRAINT -> false; + case SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR -> true; + default -> super.hasBehavior(connectorBehavior); + }; + } + + @Test + @Override + public void verifySupportsRowLevelUpdateDeclaration() { + // Planner fails with IllegalArgumentException before connector NOT_SUPPORTED path + } + + @Test + @Override + public void testRowLevelUpdate() { + // Planner fails with IllegalArgumentException before connector NOT_SUPPORTED path + } + + @Test + @Override + public void testShowCreateTable() { + String catalog = getSession().getCatalog().orElseThrow(); + String schema = getSession().getSchema().orElseThrow(); + Assertions.assertThat(computeScalar("SHOW CREATE TABLE region")).isEqualTo(String.format( + "CREATE TABLE %s.%s.region (\n" + + " regionkey bigint,\n" + + " name varchar,\n" + + " comment varchar\n" + + ")", + catalog, + schema)); + } +} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorTest.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorTest.java new file mode 100644 index 00000000..e3f84051 --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbConnectorTest.java @@ -0,0 +1,172 @@ +package tech.ydb.trino; + +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import io.trino.testing.*; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import tech.ydb.test.junit5.YdbHelperExtension; + +import java.util.Optional; + +public class TestYdbConnectorTest extends BaseConnectorTest { + + @RegisterExtension + static final YdbHelperExtension ydb = new YdbHelperExtension(); + + @Override + protected QueryRunner createQueryRunner() throws Exception { + return YdbQueryRunner.builder(ydb) + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); + } + + @Override + protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { + return switch (connectorBehavior) { + case SUPPORTS_MERGE, + SUPPORTS_UPDATE, + SUPPORTS_DELETE, + SUPPORTS_CREATE_VIEW, + SUPPORTS_CREATE_SCHEMA, + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_SET_COLUMN_TYPE, + SUPPORTS_DROP_COLUMN, + SUPPORTS_ROW_TYPE, + SUPPORTS_RENAME_COLUMN, + SUPPORTS_ROW_LEVEL_UPDATE, + SUPPORTS_TRUNCATE, + SUPPORTS_COMMENT_ON_COLUMN, + SUPPORTS_COMMENT_ON_TABLE, + SUPPORTS_DROP_SCHEMA_CASCADE, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT, + SUPPORTS_CREATE_TABLE_WITH_TABLE_COMMENT, + SUPPORTS_ADD_COLUMN_WITH_COMMENT, + SUPPORTS_ADD_COLUMN_WITH_POSITION, + SUPPORTS_CREATE_FEDERATED_MATERIALIZED_VIEW, + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_ARRAY, + SUPPORTS_MAP_TYPE, + SUPPORTS_DEFAULT_COLUMN_VALUE, + SUPPORTS_SET_DEFAULT_COLUMN_VALUE, + SUPPORTS_DROP_DEFAULT_COLUMN_VALUE, + SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, + SUPPORTS_DROP_NOT_NULL_CONSTRAINT -> false; + case SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR -> true; + default -> super.hasBehavior(connectorBehavior); + }; + } + + @Test + @Override + public void testCreateTableWithLongTableName() { + // YDB не поддерживает длинные названия таблиц + } + + @Test + @Override + public void testRenameTableToLongTableName() { + // YDB не поддерживает длинные названия таблиц + } + + @Test + @Override + public void testAlterTableAddLongColumnName() { + // YDB не поддерживает длинные названия колонок + } + + @Test + @Override + public void testCreateTableWithLongColumnName() { + // YDB не поддерживает длинные названия колонок + } + + @Test + @Override + public void testInsertNegativeDate() { + // YDB не поддерживает, negative daysSinceEpoch + } + + @Test + @Override + public void testDateYearOfEraPredicate() { + // YDB не поддерживает, negative daysSinceEpoch + } + + @Test + @Override + public void testCreateTableAsSelectNegativeDate() { + // YDB не поддерживает, negative daysSinceEpoch + } + + @Test + @Override + public void testCharVarcharComparison() { + // CHAR хранится как String без паддинга + } + + @Test + @Override + public void testVarcharCastToDateInPredicate() { + // YDB не поддерживает такой pushdown/cast + } + + @Test + @Override + public void verifySupportsRowLevelUpdateDeclaration() { + // Planner fails with IllegalArgumentException before connector NOT_SUPPORTED path + } + + @Test + @Override + public void testInsertForDefaultColumn() { + // Requires createTableWithDefaultColumns() which is connector-specific and not supported yet + } + + @Override + protected String errorMessageForInsertIntoNotNullColumn(String columnName) { + return "(?s).*(" + + "NULL value not allowed for NOT NULL column: " + columnName + + "|Cannot set NULL to not nullable column: " + columnName + + "|Missing value for not null column: " + columnName + + "|Missing not null column in input: " + columnName + + ").*"; + } + + @Override + protected boolean isColumnNameRejected(Exception exception, String columnName, boolean delimited) { + return requiresDelimiting(columnName); + } + + @Override + protected Optional filterDataMappingSmokeTestData(BaseConnectorTest.DataMappingTestSetup dataMappingTestSetup) { + if (dataMappingTestSetup.getTrinoTypeName().equals("date")) { + return Optional.of(new DataMappingTestSetup( + dataMappingTestSetup.getTrinoTypeName(), + "DATE '2006-06-06'", + "DATE '2026-06-06'" + )); + } else if (dataMappingTestSetup.getTrinoTypeName().startsWith("time") || dataMappingTestSetup.getTrinoTypeName().equals("varbinary")) { + // Нет time и varbinary в YQL + return Optional.empty(); + } + return Optional.of(dataMappingTestSetup); + } + + @Override + protected MaterializedResult getDescribeOrdersResult() { + // В YQL строки произвольной длины + return MaterializedResult.resultBuilder(this.getSession(), new Type[]{VarcharType.VARCHAR, VarcharType.VARCHAR, VarcharType.VARCHAR, VarcharType.VARCHAR}).row(new Object[]{"orderkey", "bigint", "", ""}).row(new Object[]{"custkey", "bigint", "", ""}).row(new Object[]{"orderstatus", "varchar", "", ""}).row(new Object[]{"totalprice", "double", "", ""}).row(new Object[]{"orderdate", "date", "", ""}).row(new Object[]{"orderpriority", "varchar", "", ""}).row(new Object[]{"clerk", "varchar", "", ""}).row(new Object[]{"shippriority", "integer", "", ""}).row(new Object[]{"comment", "varchar", "", ""}).build(); + } + + @Test + @Override + public void testShowCreateTable() { + // В YQL строки произвольной длины + String catalog = this.getSession().getCatalog().orElseThrow(); + String schema = this.getSession().getSchema().orElseThrow(); + Assertions.assertThat(this.computeScalar("SHOW CREATE TABLE orders")).isEqualTo(String.format("CREATE TABLE %s.%s.orders (\n orderkey bigint,\n custkey bigint,\n orderstatus varchar,\n totalprice double,\n orderdate date,\n orderpriority varchar,\n clerk varchar,\n shippriority integer,\n comment varchar\n)", catalog, schema)); + } +} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java new file mode 100644 index 00000000..d97e6149 --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java @@ -0,0 +1,65 @@ +package tech.ydb.trino; + +import io.trino.plugin.base.mapping.IdentifierMapping; +import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.SchemaTableName; + +import java.util.List; + +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; + +public class TestingYdbJdbcClient extends YdbClient { + private static final String YDB_HIDDEN_PK_COLUMN = "pk"; + + public TestingYdbJdbcClient( + BaseJdbcConfig config, + ConnectionFactory connectionFactory, + QueryBuilder queryBuilder, + IdentifierMapping identifierMapping, + RemoteQueryModifier remoteQueryModifier) { + super(config, connectionFactory, queryBuilder, identifierMapping, remoteQueryModifier); + } + + @Override + public List getColumns( + ConnectorSession session, + SchemaTableName schemaTableName, + RemoteTableName remoteTableName) { + return super.getColumns(session, schemaTableName, remoteTableName) + .stream() + .filter(column -> !column.getColumnName().equals(YDB_HIDDEN_PK_COLUMN)) + .toList(); + } + + @Override + protected List createTableSqls(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) { + if (tableMetadata.getComment().isPresent()) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with table comment"); + } + String tableName = quoted(remoteTableName); + String columnsDeclaration = String.join(", ", columns); + + boolean hasPrimaryKey = columnsDeclaration.toUpperCase().contains("PRIMARY KEY"); + boolean hasHiddenPkColumn = columns.stream() + .anyMatch(col -> col.startsWith(quoted(YDB_HIDDEN_PK_COLUMN)) || col.startsWith(YDB_HIDDEN_PK_COLUMN)); + + String sql; + if (hasPrimaryKey) { + sql = String.format("CREATE TABLE %s (%s)", tableName, columnsDeclaration); + } else if (hasHiddenPkColumn) { + sql = String.format("CREATE TABLE %s (%s, PRIMARY KEY (%s))", tableName, columnsDeclaration, quoted(YDB_HIDDEN_PK_COLUMN)); + } else { + String hiddenPkColumn = quoted(YDB_HIDDEN_PK_COLUMN) + " Serial"; + sql = String.format("CREATE TABLE %s (%s, %s, PRIMARY KEY (%s))", tableName, columnsDeclaration, hiddenPkColumn, quoted(YDB_HIDDEN_PK_COLUMN)); + } + return List.of(sql); + } +} \ No newline at end of file diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java new file mode 100644 index 00000000..7a4b6570 --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java @@ -0,0 +1,58 @@ +package tech.ydb.trino; + +import com.google.inject.Binder; +import com.google.inject.Module; +import com.google.inject.Provides; +import com.google.inject.Scopes; +import com.google.inject.Singleton; +import io.trino.plugin.base.mapping.IdentifierMapping; +import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.DriverConnectionFactory; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcMetadataFactory; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.credential.CredentialProvider; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import tech.ydb.jdbc.YdbDriver; + +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; + +public class TestingYdbJdbcModule implements Module { + + @Override + public void configure(Binder binder) { + newOptionalBinder(binder, JdbcMetadataFactory.class) + .setBinding() + .to(YdbMetadataFactory.class) + .in(Scopes.SINGLETON); + + binder.bind(YdbConnector.class).in(Scopes.SINGLETON); + } + + @Provides + @Singleton + @ForBaseJdbc + public JdbcClient provideJdbcClient( + BaseJdbcConfig config, + ConnectionFactory connectionFactory, + QueryBuilder queryBuilder, + IdentifierMapping identifierMapping, + RemoteQueryModifier remoteQueryModifier) { + return new TestingYdbJdbcClient(config, connectionFactory, queryBuilder, identifierMapping, remoteQueryModifier); + } + + @Provides + @Singleton + @ForBaseJdbc + public static ConnectionFactory createConnectionFactory( + BaseJdbcConfig config, + CredentialProvider credentialProvider) { + return DriverConnectionFactory.builder( + new YdbDriver(), + config.getConnectionUrl(), + credentialProvider) + .build(); + } +} \ No newline at end of file diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbIntegrationTest.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbIntegrationTest.java deleted file mode 100644 index a4ea1876..00000000 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbIntegrationTest.java +++ /dev/null @@ -1,56 +0,0 @@ -package tech.ydb.trino; - -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.QueryRunner; -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.Statement; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; -import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; -import org.junit.jupiter.api.extension.RegisterExtension; -import tech.ydb.test.junit5.YdbHelperExtension; - -@TestInstance(PER_CLASS) -public class YdbIntegrationTest - extends AbstractTestQueryFramework -{ - @RegisterExtension - private static final YdbHelperExtension ydb = new YdbHelperExtension(); - - @Override - protected QueryRunner createQueryRunner() throws Exception - { - QueryRunner runner = YdbQueryRunner.create(ydb); - prepareTestData(); - return runner; - } - - private void prepareTestData() throws Exception - { - try (Connection connection = DriverManager.getConnection(YdbQueryRunner.jdbcUrl(ydb)); - Statement statement = connection.createStatement()) { - statement.execute("DROP TABLE IF EXISTS test_table"); - statement.execute("CREATE TABLE test_table (" + - "id Uint64 NOT NULL, " + - "name Utf8, " + - "created_at Timestamp, " + - "PRIMARY KEY (id)" + - ");"); - statement.execute("UPSERT INTO test_table (id, name, created_at) VALUES (1, 'alice', CurrentUtcTimestamp())"); - statement.execute("UPSERT INTO test_table (id, name, created_at) VALUES (2, 'bob', CurrentUtcTimestamp())"); - } - } - - @Test - public void testSelectAndPredicate() - { - assertQuery("SELECT name FROM test_table WHERE id = 1", "VALUES 'alice'"); - } - - @Test - public void testAggregation() - { - assertQuery("SELECT count(*), sum(id) FROM test_table", "VALUES (2, 3)"); - } -} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbPluginTest.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbPluginTest.java deleted file mode 100644 index 6888f13a..00000000 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbPluginTest.java +++ /dev/null @@ -1,24 +0,0 @@ -package tech.ydb.trino; - -import io.trino.spi.connector.ConnectorFactory; -import org.junit.jupiter.api.Test; - -import java.util.Iterator; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -public class YdbPluginTest -{ - @Test - void exposesSingleConnectorFactory() - { - YdbPlugin plugin = new YdbPlugin(); - Iterator factories = plugin.getConnectorFactories().iterator(); - assertTrue(factories.hasNext(), "Expected a connector factory"); - ConnectorFactory factory = factories.next(); - assertEquals("ydb", factory.getName()); - assertFalse(factories.hasNext(), "Expected only one connector factory"); - } -} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java index 322388a9..f99ee3b6 100644 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java @@ -1,33 +1,112 @@ package tech.ydb.trino; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.trino.Session; import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.MaterializedResult; import io.trino.testing.QueryRunner; +import io.trino.tpch.TpchColumn; +import io.trino.tpch.TpchTable; +import org.intellij.lang.annotations.Language; import tech.ydb.test.junit5.YdbHelperExtension; +import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static io.trino.testing.TestingSession.testSessionBuilder; -public final class YdbQueryRunner -{ - private YdbQueryRunner() {} +public final class YdbQueryRunner { + public static final String TPCH_SCHEMA = "ydb"; + public static final String YDB_HIDDEN_PK_COLUMN = "pk"; - public static QueryRunner create(YdbHelperExtension ydb) - throws Exception - { - Session session = testSessionBuilder().setCatalog("ydb").setSchema("ydb").build(); + private YdbQueryRunner() {} - DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).build(); - queryRunner.installPlugin(new YdbPlugin()); - queryRunner.createCatalog("ydb", "ydb", Map.of( - "connection-url", jdbcUrl(ydb))); + public static Builder builder(YdbHelperExtension ydb) { + String jdbcUrl = buildJdbcUrl(ydb); + return new Builder() + // Avoid temporary-table CTAS during INSERT; YDB does not support CREATE TABLE AS SELECT. + .addConnectorProperty("insert.non-transactional-insert.enabled", "true") + .addConnectorProperty("connection-url", jdbcUrl); + } - return queryRunner; + private static String buildJdbcUrl(YdbHelperExtension ydb) { + StringBuilder url = new StringBuilder("jdbc:ydb:"); + url.append(ydb.useTls() ? "grpcs://" : "grpc://"); + url.append(ydb.endpoint()); + url.append(ydb.database()); + url.append("?useQueryService=true&sessionPoolMaxSize=10"); + if (ydb.authToken() != null) { + url.append("&token=").append(ydb.authToken()); + } + return url.toString(); } - public static String jdbcUrl(YdbHelperExtension ydb) - { - return "jdbc:ydb:" + (ydb.useTls() ? "grpcs://" : "grpc://") + ydb.endpoint() + ydb.database(); + public static class Builder extends DistributedQueryRunner.Builder { + private final Map connectorProperties = new HashMap<>(); + private List> initialTables = ImmutableList.of(); + + private Builder() { + super(testSessionBuilder() + .setCatalog("ydb") + .setSchema(TPCH_SCHEMA) + .build()); + } + + public Builder addConnectorProperty(String key, String value) { + connectorProperties.put(key, value); + return this; + } + + public Builder setInitialTables(List> tables) { + this.initialTables = tables; + return this; + } + + @Override + public DistributedQueryRunner build() throws Exception { + DistributedQueryRunner queryRunner = super.build(); + try { + queryRunner.installPlugin(new io.trino.plugin.tpch.TpchPlugin()); + queryRunner.createCatalog("tpch", "tpch"); + + queryRunner.installPlugin(new YdbPlugin(new TestingYdbJdbcModule())); + queryRunner.createCatalog("ydb", "ydb", ImmutableMap.copyOf(connectorProperties)); + + for (TpchTable table : initialTables) { + dropTable(queryRunner, table); + } + + for (TpchTable table : initialTables) { + createTableWithPk(queryRunner, table); + } + + return queryRunner; + } catch (Throwable e) { + queryRunner.close(); + throw e; + } + } + + private static void createTableWithPk(DistributedQueryRunner queryRunner, TpchTable table) { + String tableName = table.getTableName(); + + String columnNames = table.getColumns().stream() + .map(col -> col.getColumnName().substring(2)) + .collect(Collectors.joining(", ")); + + String createSql = String.format( + "CREATE TABLE %s AS SELECT %s, row_number() OVER () AS %s FROM tpch.tiny.%s", + tableName, columnNames, YDB_HIDDEN_PK_COLUMN, tableName); + + queryRunner.execute(createSql); + } + + private static void dropTable(DistributedQueryRunner queryRunner, TpchTable table) { + String tableName = table.getTableName(); + queryRunner.execute("DROP TABLE IF EXISTS " + tableName); + } } } diff --git a/ydb-trino-adapter/start.sh b/ydb-trino-adapter/start.sh new file mode 100755 index 00000000..aa02dc75 --- /dev/null +++ b/ydb-trino-adapter/start.sh @@ -0,0 +1,10 @@ +mvn -f pom.xml -DskipTests package +mvn -f pom.xml -DskipTests dependency:copy-dependencies -DincludeScope=runtime + +mkdir -p docker/trino/plugin +cp target/ydb-trino-0.1.0.jar docker/trino/plugin +cp target/dependency/*.jar docker/trino/plugin + +cd docker +docker-compose down +docker-compose up -d \ No newline at end of file