diff --git a/core/docker/arenadata/coordinator/etc/catalog/adb.properties b/core/docker/arenadata/coordinator/etc/catalog/adb.properties new file mode 100644 index 000000000000..748a182054f7 --- /dev/null +++ b/core/docker/arenadata/coordinator/etc/catalog/adb.properties @@ -0,0 +1,4 @@ +connector.name=adb +connection-url=jdbc:postgresql://adb:6000/postgres +connection-user=gpadmin +connection-password=gpadmin diff --git a/core/docker/arenadata/coordinator/etc/log.properties b/core/docker/arenadata/coordinator/etc/log.properties index abee45ebcde7..9b0788d5a139 100644 --- a/core/docker/arenadata/coordinator/etc/log.properties +++ b/core/docker/arenadata/coordinator/etc/log.properties @@ -1,2 +1,2 @@ # Enable verbose logging from Trino -#io.trino=DEBUG +io.trino=DEBUG diff --git a/core/docker/arenadata/docker-compose-with-adb.yml b/core/docker/arenadata/docker-compose-with-adb.yml new file mode 100644 index 000000000000..bd6571f689ad --- /dev/null +++ b/core/docker/arenadata/docker-compose-with-adb.yml @@ -0,0 +1,184 @@ +version: "3.8" +services: + postgres: + image: "docker.io/library/postgres:12.0" + restart: unless-stopped + container_name: postgres-db + hostname: postgres-db + environment: + POSTGRES_DB: '${POSTGRES_DB}' + POSTGRES_USER: '${POSTGRES_USER}' + POSTGRES_PASSWORD: '${POSTGRES_PASSWORD}' + ports: + - '5432:5432' + healthcheck: + test: psql -d ${POSTGRES_DB} -U ${POSTGRES_USER} -Atc 'SELECT 1;' + interval: 30s + timeout: 15s + retries: 3 + networks: + - trino-with-adb + + coordinator: + image: "trino:${TRINO_VERSION}" + ports: + - "8080:8080" + - "8990:8990" + container_name: "coordinator" + volumes: + - ./coordinator/etc:/etc/trino:rw + depends_on: + - postgres + networks: + - trino-with-adb + + worker0: + image: "trino:${TRINO_VERSION}" + container_name: "worker0" + ports: + - "8081:8080" + - "8991:8990" + volumes: + - ./worker/etc:/etc/trino:rw + depends_on: + - postgres + networks: + - "trino-with-adb" + + worker1: + image: "trino:${TRINO_VERSION}" + container_name: "worker1" + ports: + - "8082:8080" + - "8992:8990" + volumes: + - ./worker/etc:/etc/trino:rw + depends_on: + - postgres + networks: + - trino-with-adb + + adb: + image: "hub.adsw.io/adbcc/gpdb6-adbm-distributed:it" + working_dir: /home/gpadmin + hostname: mdw + ports: + - "6000:6000" + - "38005:35285" + environment: + - ADB_CLUSTER_NAME=adb + - ADBM_EUREKA_CLIENT_HOSTNAME=mdw + - ADBM_EUREKA_CLIENT_IP_ADDRESS=mdw + - ADBM_EUREKA_CLIENT_PREFER_IP_ADDRESS=true + - ADBM_EUREKA_CLIENT_SERV_URL_DEF_ZONE=http://registry:8761/eureka + - AGENT_JAVA_OPTS=-XX:+HeapDumpOnOutOfMemoryError -XX:+UseStringDeduplication -XX:+OptimizeStringConcat -XX:+UseLWPSynchronization -XX:HeapDumpPath=/opt/adbm-agent/tmp/java_pid.hprof -Xlog:gc*,safepoint:gc.log:time,uptime:filecount=100,filesize=50M -Djava.library.path=/opt/adbm-agent/tmp -Djava.io.tmpdir=/opt/adbm-agent/tmp -Dlogging.config=/opt/adbm-agent/adbm-agent-logback.xml -Xmx2g -Xms256m -Dspring.config.location=file:///opt/adbm-agent/adbm-agent-application.yml + - DOCKER_GP_CLUSTER_HOSTS=mdw,smdw,sdw1,sdw2 + - DOCKER_GP_MASTER_SERVER=mdw + - DOCKER_GP_STANDBY_SERVER=smdw + - DOCKER_GP_SEGMENT_SERVERS=sdw1,sdw2 + - DOCKER_GP_PRIMARY_SEGMENTS_PER_HOST=2 + - DOCKER_GP_WITH_MIRROR=true + depends_on: + - sdw1 + - sdw2 + - smdw + volumes: + - backup-data:/tmp/pgbackrest/:rw + privileged: true + extra_hosts: + - "host.docker.internal:host-gateway" + sysctls: + kernel.sem: 500 1024000 200 4096 + networks: + - trino-with-adb + + sdw1: + image: "hub.adsw.io/adbcc/gpdb6-adbm-distributed:it" + privileged: true + hostname: sdw1 + ports: + - "38006:35285" + environment: + - ADB_CLUSTER_NAME=adb + - ADBM_EUREKA_CLIENT_HOSTNAME=sdw1 + - ADBM_EUREKA_CLIENT_IP_ADDRESS=sdw1 + - ADBM_EUREKA_CLIENT_PREFER_IP_ADDRESS=true + - ADBM_EUREKA_CLIENT_SERV_URL_DEF_ZONE=http://registry:8761/eureka + - AGENT_JAVA_OPTS=-XX:+HeapDumpOnOutOfMemoryError -XX:+UseStringDeduplication -XX:+OptimizeStringConcat -XX:+UseLWPSynchronization -XX:HeapDumpPath=/opt/adbm-agent/tmp/java_pid.hprof -Xlog:gc*,safepoint:gc.log:time,uptime:filecount=100,filesize=50M -Djava.library.path=/opt/adbm-agent/tmp -Djava.io.tmpdir=/opt/adbm-agent/tmp -Dlogging.config=/opt/adbm-agent/adbm-agent-logback.xml -Xmx2g -Xms256m -Dspring.config.location=file:///opt/adbm-agent/adbm-agent-application.yml + - DOCKER_GP_CLUSTER_HOSTS=mdw,smdw,sdw1,sdw2 + - DOCKER_GP_MASTER_SERVER=mdw + - DOCKER_GP_STANDBY_SERVER=smdw + - DOCKER_GP_SEGMENT_SERVERS=sdw1,sdw2 + - DOCKER_GP_PRIMARY_SEGMENTS_PER_HOST=2 + - DOCKER_GP_WITH_MIRROR=true + volumes: + - backup-data:/tmp/pgbackrest/:rw + extra_hosts: + - "host.docker.internal:host-gateway" + sysctls: + kernel.sem: 500 1024000 200 4096 + networks: + - trino-with-adb + + sdw2: + image: "hub.adsw.io/adbcc/gpdb6-adbm-distributed:it" + privileged: true + hostname: sdw2 + ports: + - "38007:35285" + environment: + - ADB_CLUSTER_NAME=adb + - ADBM_EUREKA_CLIENT_HOSTNAME=sdw2 + - ADBM_EUREKA_CLIENT_IP_ADDRESS=sdw2 + - ADBM_EUREKA_CLIENT_PREFER_IP_ADDRESS=true + - ADBM_EUREKA_CLIENT_SERV_URL_DEF_ZONE=http://registry:8761/eureka + - AGENT_JAVA_OPTS=-XX:+HeapDumpOnOutOfMemoryError -XX:+UseStringDeduplication -XX:+OptimizeStringConcat -XX:+UseLWPSynchronization -XX:HeapDumpPath=/opt/adbm-agent/tmp/java_pid.hprof -Xlog:gc*,safepoint:gc.log:time,uptime:filecount=100,filesize=50M -Djava.library.path=/opt/adbm-agent/tmp -Djava.io.tmpdir=/opt/adbm-agent/tmp -Dlogging.config=/opt/adbm-agent/adbm-agent-logback.xml -Xmx2g -Xms256m -Dspring.config.location=file:///opt/adbm-agent/adbm-agent-application.yml + - DOCKER_GP_CLUSTER_HOSTS=mdw,smdw,sdw1,sdw2 + - DOCKER_GP_MASTER_SERVER=mdw + - DOCKER_GP_STANDBY_SERVER=smdw + - DOCKER_GP_SEGMENT_SERVERS=sdw1,sdw2 + - DOCKER_GP_PRIMARY_SEGMENTS_PER_HOST=2 + - DOCKER_GP_WITH_MIRROR=true + volumes: + - backup-data:/tmp/pgbackrest/:rw + extra_hosts: + - "host.docker.internal:host-gateway" + sysctls: + kernel.sem: 500 1024000 200 4096 + networks: + - trino-with-adb + + smdw: + image: "hub.adsw.io/adbcc/gpdb6-adbm-distributed:it" + privileged: true + hostname: smdw + ports: + - "5433:6000" + - "38008:35285" + environment: + - ADB_CLUSTER_NAME=adb + - ADBM_EUREKA_CLIENT_HOSTNAME=smdw + - ADBM_EUREKA_CLIENT_IP_ADDRESS=smdw + - ADBM_EUREKA_CLIENT_PREFER_IP_ADDRESS=true + - ADBM_EUREKA_CLIENT_SERV_URL_DEF_ZONE=http://registry:8761/eureka + - AGENT_JAVA_OPTS=-XX:+HeapDumpOnOutOfMemoryError -XX:+UseStringDeduplication -XX:+OptimizeStringConcat -XX:+UseLWPSynchronization -XX:HeapDumpPath=/opt/adbm-agent/tmp/java_pid.hprof -Xlog:gc*,safepoint:gc.log:time,uptime:filecount=100,filesize=50M -Djava.library.path=/opt/adbm-agent/tmp -Djava.io.tmpdir=/opt/adbm-agent/tmp -Dlogging.config=/opt/adbm-agent/adbm-agent-logback.xml -Xmx2g -Xms256m -Dspring.config.location=file:///opt/adbm-agent/adbm-agent-application.yml + - DOCKER_GP_CLUSTER_HOSTS=mdw,smdw,sdw1,sdw2 + - DOCKER_GP_MASTER_SERVER=mdw + - DOCKER_GP_STANDBY_SERVER=smdw + - DOCKER_GP_SEGMENT_SERVERS=sdw1,sdw2 + - DOCKER_GP_PRIMARY_SEGMENTS_PER_HOST=2 + - DOCKER_GP_WITH_MIRROR=true + volumes: + - backup-data:/tmp/pgbackrest/:rw + extra_hosts: + - "host.docker.internal:host-gateway" + sysctls: + kernel.sem: 500 1024000 200 4096 + networks: + - trino-with-adb + +networks: + trino-with-adb: +volumes: + backup-data: + driver: local diff --git a/core/docker/arenadata/worker/etc/catalog/adb.properties b/core/docker/arenadata/worker/etc/catalog/adb.properties new file mode 100644 index 000000000000..748a182054f7 --- /dev/null +++ b/core/docker/arenadata/worker/etc/catalog/adb.properties @@ -0,0 +1,4 @@ +connector.name=adb +connection-url=jdbc:postgresql://adb:6000/postgres +connection-user=gpadmin +connection-password=gpadmin diff --git a/core/docker/arenadata/worker/etc/log.properties b/core/docker/arenadata/worker/etc/log.properties index abee45ebcde7..9b0788d5a139 100644 --- a/core/docker/arenadata/worker/etc/log.properties +++ b/core/docker/arenadata/worker/etc/log.properties @@ -1,2 +1,2 @@ # Enable verbose logging from Trino -#io.trino=DEBUG +io.trino=DEBUG diff --git a/core/trino-server/src/main/provisio/trino.xml b/core/trino-server/src/main/provisio/trino.xml index 993ef4e912d5..35c855fd7dd9 100644 --- a/core/trino-server/src/main/provisio/trino.xml +++ b/core/trino-server/src/main/provisio/trino.xml @@ -33,6 +33,12 @@ + + + + + + diff --git a/plugin/trino-adb/pom.xml b/plugin/trino-adb/pom.xml new file mode 100644 index 000000000000..0fc84cadc2b1 --- /dev/null +++ b/plugin/trino-adb/pom.xml @@ -0,0 +1,247 @@ + + + 4.0.0 + + io.trino + trino-root + 455-SNAPSHOT + ../../pom.xml + + + trino-adb + trino-plugin + Trino - Adb connector + + + + com.google.guava + guava + + + + com.google.inject + guice + + + + com.opencsv + opencsv + + + commons-logging + commons-logging + + + + + + io.airlift + concurrent + + + + io.airlift + configuration + + + io.airlift + log + + + + + + io.airlift + event + compile + + + + io.airlift + http-server + compile + + + + io.airlift + jaxrs + + + io.airlift + event + + + io.airlift + http-server + + + io.airlift + node + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-context + + + jakarta.servlet + jakarta.servlet-api + + + + + + io.airlift + json + + + + io.airlift + log + compile + + + + io.airlift + node + compile + + + + io.airlift + units + compile + + + io.trino + trino-base-jdbc + + + io.airlift + concurrent + + + io.airlift + log-manager + + + io.airlift + units + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-context + + + io.trino + trino-plugin-toolkit + + + joda-time + joda-time + + + org.weakref + jmxutils + + + + + + io.trino + trino-plugin-toolkit + compile + + + + jakarta.servlet + jakarta.servlet-api + compile + + + + jakarta.validation + jakarta.validation-api + + + + jakarta.ws.rs + jakarta.ws.rs-api + + + + joda-time + joda-time + compile + + + + org.jdbi + jdbi3-core + + + + org.postgresql + postgresql + + + + org.weakref + jmxutils + compile + + + + com.fasterxml.jackson.core + jackson-annotations + provided + + + + io.airlift + slice + provided + + + + io.opentelemetry + opentelemetry-api + provided + + + + io.opentelemetry + opentelemetry-context + provided + + + + io.trino + trino-spi + provided + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-context + + + + + + io.airlift + log-manager + runtime + + + diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPlugin.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPlugin.java new file mode 100644 index 000000000000..2cadd0eafbf0 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPlugin.java @@ -0,0 +1,28 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb; + +import io.trino.plugin.adb.connector.AdbClientModule; +import io.trino.plugin.jdbc.JdbcPlugin; + +public class AdbPlugin + extends JdbcPlugin +{ + private static final String PLUGIN_NAME = "adb"; + + public AdbPlugin() + { + super(PLUGIN_NAME, new AdbClientModule()); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPluginConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPluginConfig.java new file mode 100644 index 000000000000..199292e192ba --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/AdbPluginConfig.java @@ -0,0 +1,142 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb; + +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import io.airlift.units.MinDataSize; +import io.trino.plugin.adb.connector.protocol.TransferDataProtocol; +import jakarta.validation.constraints.Min; +import jakarta.validation.constraints.NotNull; + +public class AdbPluginConfig +{ + public static final String IDENTIFIER_QUOTE = "\""; + private AdbPluginConfig.ArrayMapping arrayMapping = AdbPluginConfig.ArrayMapping.DISABLED; + private int maxScanParallelism = 1; + private boolean includeSystemTables; + private Integer fetchSize; + private DataSize writeBufferSize = DataSize.of(64L, DataSize.Unit.MEGABYTE); + private DataSize readBufferSize = DataSize.of(64L, DataSize.Unit.MEGABYTE); + private final TransferDataProtocol dataProtocol = TransferDataProtocol.GPFDIST; + private Duration gpfdistRetryTimeout; + + public TransferDataProtocol getDataProtocol() + { + return dataProtocol; + } + + public Integer getFetchSize() + { + return fetchSize; + } + + @Config("adb.fetch-size") + public AdbPluginConfig setFetchSize(int fetchSize) + { + this.fetchSize = fetchSize; + return this; + } + + @NotNull + public AdbPluginConfig.ArrayMapping getArrayMapping() + { + return arrayMapping; + } + + @Config("adb.array-mapping") + public AdbPluginConfig setArrayMapping(AdbPluginConfig.ArrayMapping arrayMapping) + { + this.arrayMapping = arrayMapping; + return this; + } + + public boolean isIncludeSystemTables() + { + return this.includeSystemTables; + } + + @Config("adb.include-system-tables") + public AdbPluginConfig setIncludeSystemTables(boolean includeSystemTables) + { + this.includeSystemTables = includeSystemTables; + return this; + } + + @Min(1L) + public int getMaxScanParallelism() + { + return maxScanParallelism; + } + + @Config("adb.max-scan-parallelism") + @ConfigDescription("Maximum degree of parallelism when scanning tables. Defaults to 1.") + public AdbPluginConfig setMaxScanParallelism(int maxScanParallelism) + { + this.maxScanParallelism = maxScanParallelism; + return this; + } + + @MinDataSize("1kB") + @NotNull + public DataSize getWriteBufferSize() + { + return writeBufferSize; + } + + @Config("adb.connector.write-buffer-size") + @ConfigDescription("Maximum amount of memory that could be allocated per sink when executing write queries. Defaults to 64MB") + public AdbPluginConfig setWriteBufferSize(DataSize writeBufferSize) + { + this.writeBufferSize = writeBufferSize; + return this; + } + + @MinDataSize("1kB") + @NotNull + public DataSize getReadBufferSize() + { + return readBufferSize; + } + + @Config("adb.connector.read-buffer-size") + @ConfigDescription("Maximum amount of memory that could be allocated per record cursor when executing read queries. Defaults to 64MB") + public AdbPluginConfig setReadBufferSize(DataSize readBufferSize) + { + this.readBufferSize = readBufferSize; + return this; + } + + public Duration getGpfdistRetryTimeout() + { + return this.gpfdistRetryTimeout; + } + + @Config("adb.gpfdist.retry-timeout") + @ConfigDescription("Value of adb gpfdist_retry_timeout property. Defaults to null (use adb defaults)") + public AdbPluginConfig setGpfdistRetryTimeout(Duration gpfdistRetryTimeout) + { + this.gpfdistRetryTimeout = gpfdistRetryTimeout; + return this; + } + + public static enum ArrayMapping + { + DISABLED, + AS_ARRAY, + AS_JSON; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/TypeUtil.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/TypeUtil.java new file mode 100644 index 000000000000..00c55fc0c18f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/TypeUtil.java @@ -0,0 +1,253 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb; + +import com.google.common.primitives.Shorts; +import com.google.common.primitives.SignedBytes; +import io.airlift.slice.Slice; +import io.trino.plugin.adb.connector.datatype.mapper.DataTypeMapper; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.CharType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Int128; +import io.trino.spi.type.LongTimestampWithTimeZone; +import io.trino.spi.type.TimestampType; +import io.trino.spi.type.TimestampWithTimeZoneType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import org.joda.time.DateTimeZone; +import org.postgresql.util.PGobject; + +import java.lang.reflect.Array; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.math.MathContext; +import java.sql.Date; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.format.DateTimeFormatter; +import java.time.format.DateTimeFormatterBuilder; +import java.time.format.SignStyle; +import java.time.temporal.ChronoField; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.jdbc.StandardColumnMappings.fromTrinoTimestamp; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc; +import static io.trino.spi.type.DateType.DATE; +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.TimeZoneKey.UTC_KEY; +import static io.trino.spi.type.Timestamps.MILLISECONDS_PER_SECOND; +import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND; +import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND; +import static io.trino.spi.type.TinyintType.TINYINT; +import static io.trino.spi.type.TypeUtils.readNativeValue; +import static java.lang.Float.intBitsToFloat; +import static java.lang.Math.floorDiv; +import static java.lang.Math.floorMod; +import static java.lang.Math.toIntExact; +import static java.util.concurrent.TimeUnit.DAYS; +import static org.joda.time.DateTimeZone.UTC; + +public final class TypeUtil +{ + public static final DateTimeFormatter TIMESTAMP_TYPE_FORMATTER = new DateTimeFormatterBuilder() + .appendValue(ChronoField.YEAR_OF_ERA, 4, 9, SignStyle.NORMAL) + .parseLenient() + .appendPattern("-MM-dd HH:mm[:ss[.SSSSSS]][XXX][ G]") + .toFormatter(); + public static final DateTimeFormatter DATE_TYPE_FORMATTER = new DateTimeFormatterBuilder() + .appendValue(ChronoField.YEAR_OF_ERA, 4, 9, SignStyle.NORMAL) + .appendPattern("-MM-dd[ G]") + .toFormatter(); + public static final DateTimeFormatter TIME_TYPE_FORMATTER = DateTimeFormatter.ofPattern("HH:mm:ss.SSSSSS"); + + private TypeUtil() + { + } + + public static PGobject toPgTimestamp(LocalDateTime localDateTime) + throws SQLException + { + PGobject pgObject = new PGobject(); + pgObject.setType("timestamp"); + pgObject.setValue(localDateTime.toString()); + return pgObject; + } + + public static int arrayDepth(Object jdbcArray) + { + checkArgument(jdbcArray.getClass().isArray(), "jdbcArray is not an array"); + int depth = 0; + while (jdbcArray != null && jdbcArray.getClass().isArray()) { + depth++; + if (Array.getLength(jdbcArray) == 0) { + return depth; + } + jdbcArray = Array.get(jdbcArray, 0); + } + return depth; + } + + public static Object[] getJdbcObjectArray(ConnectorSession session, Type elementType, Block block) + throws SQLException + { + int positionCount = block.getPositionCount(); + Object[] valuesArray = new Object[positionCount]; + int subArrayLength = 1; + for (int i = 0; i < positionCount; i++) { + Object objectValue = trinoNativeToJdbcObject(session, elementType, readNativeValue(elementType, block, i)); + valuesArray[i] = objectValue; + if (objectValue != null && objectValue.getClass().isArray()) { + subArrayLength = Math.max(subArrayLength, Array.getLength(objectValue)); + } + } + // PostgreSql requires subarrays with matching dimensions, including arrays of null + if (elementType instanceof ArrayType) { + handleArrayNulls(valuesArray, subArrayLength); + } + return valuesArray; + } + + public static String getArrayElementPgTypeName(ConnectorSession session, DataTypeMapper typeMapper, Type elementType) + { + if (DOUBLE.equals(elementType)) { + return "float8"; + } + + if (REAL.equals(elementType)) { + return "float4"; + } + + // TypeInfoCache#getPGArrayType apparently doesn't allow for specifying variable-length limits. + // Since PostgreSQL char[] can only hold single byte elements, map all CharType to varchar + if (elementType instanceof VarcharType || elementType instanceof CharType) { + return "varchar"; + } + + // TypeInfoCache#getPGArrayType apparently doesn't allow for specifying variable-length limits. + // Map all timestamptz(x) types to unparametrized one which defaults to highest precision + if (elementType instanceof TimestampWithTimeZoneType) { + return "timestamptz"; + } + + // TypeInfoCache#getPGArrayType doesn't allow for specifying variable-length limits. + // Map all timestamp(x) types to unparametrized one which defaults to highest precision + if (elementType instanceof TimestampType) { + return "timestamp"; + } + + if (elementType instanceof DecimalType) { + return "decimal"; + } + + if (elementType instanceof ArrayType) { + return getArrayElementPgTypeName(session, typeMapper, ((ArrayType) elementType).getElementType()); + } + + return typeMapper.toWriteMapping(session, elementType).getDataType(); + } + + private static Object trinoNativeToJdbcObject(ConnectorSession session, Type trinoType, Object trinoNative) + throws SQLException + { + if (trinoNative == null) { + return null; + } + + if (DOUBLE.equals(trinoType) || BOOLEAN.equals(trinoType) || BIGINT.equals(trinoType)) { + return trinoNative; + } + + if (trinoType instanceof DecimalType decimalType) { + if (decimalType.isShort()) { + BigInteger unscaledValue = BigInteger.valueOf((long) trinoNative); + return new BigDecimal(unscaledValue, decimalType.getScale(), new MathContext(decimalType.getPrecision())); + } + BigInteger unscaledValue = ((Int128) trinoNative).toBigInteger(); + return new BigDecimal(unscaledValue, decimalType.getScale(), new MathContext(decimalType.getPrecision())); + } + + if (REAL.equals(trinoType)) { + return intBitsToFloat(toIntExact((long) trinoNative)); + } + + if (TINYINT.equals(trinoType)) { + return SignedBytes.checkedCast((long) trinoNative); + } + + if (SMALLINT.equals(trinoType)) { + return Shorts.checkedCast((long) trinoNative); + } + + if (INTEGER.equals(trinoType)) { + return toIntExact((long) trinoNative); + } + + if (DATE.equals(trinoType)) { + // convert to midnight in default time zone + long millis = DAYS.toMillis((long) trinoNative); + return new Date(UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millis)); + } + + if (trinoType instanceof TimestampType timestampType && timestampType.isShort()) { + return toPgTimestamp(fromTrinoTimestamp((long) trinoNative)); + } + + if (trinoType instanceof TimestampWithTimeZoneType timestampWithTimeZoneType) { + // PostgreSQL does not store zone, only the point in time + int precision = timestampWithTimeZoneType.getPrecision(); + if (precision <= TimestampWithTimeZoneType.MAX_SHORT_PRECISION) { + long millisUtc = unpackMillisUtc((long) trinoNative); + return new Timestamp(millisUtc); + } + LongTimestampWithTimeZone value = (LongTimestampWithTimeZone) trinoNative; + long epochSeconds = floorDiv(value.getEpochMillis(), MILLISECONDS_PER_SECOND); + long nanosOfSecond = floorMod(value.getEpochMillis(), MILLISECONDS_PER_SECOND) * NANOSECONDS_PER_MILLISECOND + + value.getPicosOfMilli() / PICOSECONDS_PER_NANOSECOND; + return OffsetDateTime.ofInstant(Instant.ofEpochSecond(epochSeconds, nanosOfSecond), UTC_KEY.getZoneId()); + } + + if (trinoType instanceof VarcharType || trinoType instanceof CharType) { + return ((Slice) trinoNative).toStringUtf8(); + } + + if (trinoType instanceof ArrayType arrayType) { + // process subarray of multi-dimensional array + return getJdbcObjectArray(session, arrayType.getElementType(), (Block) trinoNative); + } + + throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + trinoType); + } + + private static void handleArrayNulls(Object[] valuesArray, int length) + { + for (int i = 0; i < valuesArray.length; i++) { + if (valuesArray[i] == null) { + valuesArray[i] = new Object[length]; + } + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbClientModule.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbClientModule.java new file mode 100644 index 000000000000..5c055518324e --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbClientModule.java @@ -0,0 +1,78 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import com.google.inject.Binder; +import com.google.inject.Scopes; +import com.google.inject.multibindings.Multibinder; +import com.google.inject.multibindings.OptionalBinder; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.configuration.ConfigBinder; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.datatype.mapper.DataTypeMapper; +import io.trino.plugin.adb.connector.datatype.mapper.DataTypeMapperImpl; +import io.trino.plugin.adb.connector.encode.DataFormatModule; +import io.trino.plugin.adb.connector.metadata.AdbMetadataDao; +import io.trino.plugin.adb.connector.metadata.impl.AdbMetadataDaoImpl; +import io.trino.plugin.adb.connector.protocol.TransferDataProtocol; +import io.trino.plugin.adb.connector.protocol.gpfdist.GpfdistModule; +import io.trino.plugin.adb.connector.table.AdbCreateTableStorageConfig; +import io.trino.plugin.adb.connector.table.AdbTableProperties; +import io.trino.plugin.adb.connector.table.SplitSourceManager; +import io.trino.plugin.adb.connector.table.SplitSourceManagerImpl; +import io.trino.plugin.adb.connector.table.StatisticsManager; +import io.trino.plugin.adb.connector.table.StatisticsManagerImpl; +import io.trino.plugin.jdbc.DecimalModule; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcJoinPushdownSupportModule; +import io.trino.plugin.jdbc.JdbcModule; +import io.trino.plugin.jdbc.JdbcStatisticsConfig; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RemoteQueryCancellationModule; +import io.trino.plugin.jdbc.ptf.Query; +import io.trino.spi.function.table.ConnectorTableFunction; + +public class AdbClientModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + binder.bind(AdbMetadataDao.class).to(AdbMetadataDaoImpl.class).in(Scopes.SINGLETON); + binder.bind(DataTypeMapper.class).to(DataTypeMapperImpl.class).in(Scopes.SINGLETON); + binder.bind(StatisticsManager.class).to(StatisticsManagerImpl.class).in(Scopes.SINGLETON); + binder.bind(SplitSourceManager.class).to(SplitSourceManagerImpl.class).in(Scopes.SINGLETON); + binder.bind(JdbcClient.class).annotatedWith(ForBaseJdbc.class).to(AdbSqlClient.class).in(Scopes.SINGLETON); + ConfigBinder.configBinder(binder).bindConfig(AdbCreateTableStorageConfig.class); + ConfigBinder.configBinder(binder).bindConfig(JdbcStatisticsConfig.class); + JdbcModule.bindSessionPropertiesProvider(binder, AdbSessionProperties.class); + JdbcModule.bindTablePropertiesProvider(binder, AdbTableProperties.class); + OptionalBinder.newOptionalBinder(binder, QueryBuilder.class).setBinding().to(CollationAwareQueryBuilder.class).in(Scopes.SINGLETON); + Multibinder.newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON); + AdbPluginConfig pluginConfig = this.buildConfigObject(AdbPluginConfig.class); + + install(new DataFormatModule()); + install(new DecimalModule()); + install(new JdbcJoinPushdownSupportModule()); + install(new RemoteQueryCancellationModule()); + + if (pluginConfig.getDataProtocol() == TransferDataProtocol.GPFDIST) { + install(new GpfdistModule()); + } + else { + throw new UnsupportedOperationException("Unsupported data protocol: " + pluginConfig.getDataProtocol()); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbColumnMapping.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbColumnMapping.java new file mode 100644 index 000000000000..1ef610415be7 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbColumnMapping.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.jdbc.ColumnMapping; + +public record AdbColumnMapping(ColumnMapping columnMapping, ColumnDataType columnDataType) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbJdbcSplit.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbJdbcSplit.java new file mode 100644 index 000000000000..b81ae2abd06b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbJdbcSplit.java @@ -0,0 +1,64 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.airlift.slice.SizeOf; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcDynamicFilterAwareSplit; +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.spi.predicate.TupleDomain; + +import java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class AdbJdbcSplit + extends JdbcSplit + implements JdbcDynamicFilterAwareSplit +{ + private static final int INSTANCE_SIZE = SizeOf.instanceSize(AdbJdbcSplit.class); + private final List distribution; + + @JsonCreator + public AdbJdbcSplit( + @JsonProperty("distribution") List distribution, + @JsonProperty("additionalPredicate") Optional additionalPredicate, + @JsonProperty("dynamicFilter") TupleDomain dynamicFilter) + { + super(additionalPredicate, dynamicFilter); + this.distribution = requireNonNull(distribution, "distribution is null"); + } + + public AdbJdbcSplit withDynamicFilter(TupleDomain dynamicFilter) + { + return new AdbJdbcSplit(this.distribution, this.getAdditionalPredicate(), dynamicFilter); + } + + @JsonProperty + public List getDistribution() + { + return this.distribution; + } + + public long getRetainedSizeInBytes() + { + return (long) INSTANCE_SIZE + + SizeOf.sizeOf(this.getAdditionalPredicate(), SizeOf::estimatedSizeOf) + + this.getDynamicFilter().getRetainedSizeInBytes(JdbcColumnHandle::getRetainedSizeInBytes) + + SizeOf.estimatedSizeOf(this.distribution, SizeOf::estimatedSizeOf); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSessionProperties.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSessionProperties.java new file mode 100644 index 000000000000..2660b80833ba --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSessionProperties.java @@ -0,0 +1,69 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; +import io.airlift.units.Duration; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.base.session.PropertyMetadataUtil; +import io.trino.plugin.base.session.SessionPropertiesProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.session.PropertyMetadata; + +import java.util.List; +import java.util.Optional; + +public class AdbSessionProperties + implements SessionPropertiesProvider +{ + private final List> sessionProperties; + + @Inject + public AdbSessionProperties(AdbPluginConfig config) + { + this.sessionProperties = ImmutableList.of( + PropertyMetadata.enumProperty("array_mapping", "Handling of PostgreSql arrays", AdbPluginConfig.ArrayMapping.class, config.getArrayMapping(), false), + PropertyMetadata.integerProperty( + "max_scan_parallelism", "Maximum degree of parallelism when scanning tables. Defaults to 1.", config.getMaxScanParallelism(), false), + PropertyMetadataUtil.durationProperty( + "gpfdist_retry_timeout", "Value of adb gpfdist_retry_timeout property", config.getGpfdistRetryTimeout(), false)); + } + + @Override + public List> getSessionProperties() + { + return sessionProperties; + } + + public static AdbPluginConfig.ArrayMapping getArrayMapping(ConnectorSession session) + { + return session.getProperty("array_mapping", AdbPluginConfig.ArrayMapping.class); + } + + public static boolean isEnableStringPushdownWithCollate(ConnectorSession session) + { + return session.getProperty("enable_string_pushdown_with_collate", Boolean.class); + } + + public static int getMaxScanParallelism(ConnectorSession session) + { + return session.getProperty("max_scan_parallelism", Integer.class); + } + + public static Optional getGpfdistRetryTimeout(ConnectorSession session) + { + return Optional.ofNullable(session.getProperty("gpfdist_retry_timeout", Duration.class)); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSqlClient.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSqlClient.java new file mode 100644 index 000000000000..7fe71c2602ec --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/AdbSqlClient.java @@ -0,0 +1,916 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.inject.Inject; +import io.airlift.log.Logger; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.datatype.mapper.DataTypeMapper; +import io.trino.plugin.adb.connector.metadata.AdbMetadataDao; +import io.trino.plugin.adb.connector.table.AdbTableDistributed; +import io.trino.plugin.adb.connector.table.AdbTableProperties; +import io.trino.plugin.adb.connector.table.AdbTableStorageCompressType; +import io.trino.plugin.adb.connector.table.AdbTableStorageOrientation; +import io.trino.plugin.adb.connector.table.SplitSourceManager; +import io.trino.plugin.adb.connector.table.StatisticsManager; +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.base.mapping.RemoteIdentifiers; +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.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcErrorCode; +import io.trino.plugin.jdbc.JdbcExpression; +import io.trino.plugin.jdbc.JdbcJoinCondition; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.JdbcSortItem; +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.plugin.jdbc.QueryBuilder; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.UnsupportedTypeHandling; +import io.trino.plugin.jdbc.WriteMapping; +import io.trino.plugin.jdbc.aggregation.ImplementAvgDecimal; +import io.trino.plugin.jdbc.aggregation.ImplementAvgFloatingPoint; +import io.trino.plugin.jdbc.aggregation.ImplementCorr; +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.ImplementCovariancePop; +import io.trino.plugin.jdbc.aggregation.ImplementCovarianceSamp; +import io.trino.plugin.jdbc.aggregation.ImplementRegrIntercept; +import io.trino.plugin.jdbc.aggregation.ImplementRegrSlope; +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 io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.AggregateFunction; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitSource; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.JoinCondition; +import io.trino.spi.connector.JoinStatistics; +import io.trino.spi.connector.JoinType; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.TableNotFoundException; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.statistics.TableStatistics; +import io.trino.spi.type.CharType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Types; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; +import java.util.StringJoiner; +import java.util.function.BiFunction; +import java.util.stream.Stream; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; +import static io.trino.plugin.adb.AdbPluginConfig.ArrayMapping.AS_ARRAY; +import static io.trino.plugin.adb.AdbPluginConfig.IDENTIFIER_QUOTE; +import static io.trino.plugin.adb.connector.AdbSessionProperties.getArrayMapping; +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static io.trino.plugin.jdbc.JdbcJoinPushdownUtil.implementJoinCostAware; +import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling; +import static io.trino.plugin.jdbc.UnsupportedTypeHandling.IGNORE; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static java.lang.Math.max; +import static java.lang.String.format; +import static java.sql.DatabaseMetaData.columnNoNulls; +import static java.util.stream.Collectors.joining; + +public class AdbSqlClient + extends BaseJdbcClient +{ + private static final Logger log = Logger.get(AdbSqlClient.class); + private static final String DUPLICATE_TABLE_SQLSTATE = "42P07"; + private final AdbMetadataDao metadata; + private final List tableTypes; + private final Integer fetchSize; + private final ConnectorExpressionRewriter connectorExpressionRewriter; + private final AggregateFunctionRewriter aggregateFunctionRewriter; + private final DataTypeMapper dataTypeMapper; + private final StatisticsManager statisticsManager; + private final SplitSourceManager splitSourceManager; + + @Inject + public AdbSqlClient(ConnectionFactory connectionFactory, + QueryBuilder queryBuilder, + BaseJdbcConfig jdbcConfig, + IdentifierMapping identifierMapping, + RemoteQueryModifier remoteQueryModifier, + AdbPluginConfig config, + AdbMetadataDao metadata, + DataTypeMapper dataTypeMapper, + StatisticsManager statisticsManager, + SplitSourceManager splitSourceManager) + { + super(IDENTIFIER_QUOTE, + connectionFactory, + queryBuilder, + jdbcConfig.getJdbcTypesMappedToVarchar(), + identifierMapping, + remoteQueryModifier, + true); + this.dataTypeMapper = dataTypeMapper; + this.statisticsManager = statisticsManager; + this.splitSourceManager = splitSourceManager; + ImmutableList.Builder tableTypes = ImmutableList.builder(); + tableTypes.add("TABLE", "PARTITIONED TABLE", "VIEW", "MATERIALIZED VIEW", "FOREIGN TABLE"); + if (config.isIncludeSystemTables()) { + tableTypes.add("SYSTEM TABLE", "SYSTEM VIEW"); + } + this.tableTypes = tableTypes.build(); + this.metadata = metadata; + this.fetchSize = config.getFetchSize(); + connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder() + .addStandardRules(this::quoted) + .withTypeClass("integer_type", ImmutableSet.of("tinyint", "smallint", "integer", "bigint")) + .withTypeClass("decimal_type", ImmutableSet.of("decimal")) + .withTypeClass("double_type", ImmutableSet.of("real", "double")) + .withTypeClass("numeric_type", + ImmutableSet.of("tinyint", "smallint", "integer", "bigint", "decimal", "real", "double")) + .withTypeClass("string_type", ImmutableSet.of("char", "varchar")) + .withTypeClass("datetime_type", + ImmutableSet.of("date", "time", "timestamp", "timestamp with time zone")) + .map("$not($is_null(value))") + .to("value IS NOT NULL") + .map("$not(value: boolean)") + .to("NOT value") + .add(new RewriteIn()) + .map("$is_null(value)") + .to("value IS NULL") + .map("$nullif(first, second)") + .to("NULLIF(first, second)") + .map("$equal(left, right)") + .to("left = right") + .map("$not_equal(left, right)") + .to("left <> right") + .map("$is_distinct_from(left, right)") + .to("left IS DISTINCT FROM right") + .map("$less_than(left: numeric_type, right: numeric_type)") + .to("left < right") + .map("$less_than_or_equal(left: numeric_type, right: numeric_type)") + .to("left <= right") + .map("$greater_than(left: numeric_type, right: numeric_type)") + .to("left > right") + .map("$greater_than_or_equal(left: numeric_type, right: numeric_type)") + .to("left >= right") + .build(); + JdbcTypeHandle bigintTypeHandle = new JdbcTypeHandle(Types.BIGINT, + Optional.of("bigint"), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty()); + this.aggregateFunctionRewriter = new AggregateFunctionRewriter<>( + this.connectorExpressionRewriter, + ImmutableSet.>builder() + .add(new ImplementCountAll(bigintTypeHandle)) + .add(new ImplementCount(bigintTypeHandle)) + .add(new ImplementCountDistinct(bigintTypeHandle, false)) + .add(new ImplementSum(AdbSqlClient::toDecimalTypeToTypeHandle)) + .add(new ImplementAvgFloatingPoint()) + .add(new ImplementAvgDecimal()) + .add(new ImplementCovarianceSamp()) + .add(new ImplementCovariancePop()) + .add(new ImplementCorr()) + .add(new ImplementRegrIntercept()) + .add(new ImplementRegrSlope()) + .build()); + } + + @Override + public void execute(ConnectorSession session, String query) + { + super.execute(session, query); + } + + @Override + public void execute(ConnectorSession session, Connection connection, String query) + throws SQLException + { + super.execute(session, connection, query); + } + + @Override + public PreparedQuery prepareQuery( + ConnectorSession session, + Connection connection, + JdbcTableHandle table, + Optional>> groupingSets, + List columns, + Map columnExpressions, + Optional split) + { + return super.prepareQuery(session, connection, table, groupingSets, columns, columnExpressions, split); + } + + public void executeAsPreparedStatement(ConnectorSession session, Connection connection, PreparedQuery preparedQuery) + throws SQLException + { + try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(this, + session, + connection, + preparedQuery, + Optional.empty())) { + preparedStatement.executeUpdate(); + } + } + + @Override + public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle tableHandle) + { + return splitSourceManager.getSplits(session, tableHandle); + } + + private static Optional toDecimalTypeToTypeHandle(DecimalType decimalType) + { + return Optional.of( + new JdbcTypeHandle(Types.NUMERIC, + Optional.of("decimal"), + Optional.of(decimalType.getPrecision()), + Optional.of(decimalType.getScale()), + Optional.empty(), + Optional.empty())); + } + + @Override + public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) + { + try { + createTable(session, tableMetadata, tableMetadata.getTable().getTableName()); + } + catch (SQLException e) { + boolean exists = DUPLICATE_TABLE_SQLSTATE.equals(e.getSQLState()); + throw new TrinoException((exists ? StandardErrorCode.ALREADY_EXISTS : JdbcErrorCode.JDBC_ERROR), e); + } + } + + @Override + protected JdbcOutputTableHandle createTable( + ConnectorSession session, + Connection connection, + ConnectorTableMetadata tableMetadata, + RemoteIdentifiers remoteIdentifiers, + String catalog, + String remoteSchema, + String remoteTable, + String remoteTargetTableName, + Optional pageSinkIdColumn) + throws SQLException + { + List columns = tableMetadata.getColumns(); + ImmutableList.Builder columnNames = ImmutableList.builderWithExpectedSize(columns.size()); + ImmutableList.Builder columnTypes = ImmutableList.builderWithExpectedSize(columns.size()); + ImmutableList.Builder columnList = ImmutableList.builderWithExpectedSize(columns.size() + + (pageSinkIdColumn.isPresent() ? 1 : 0)); + IdentifierMapping identifierMapping = this.getIdentifierMapping(); + + for (ColumnMetadata column : columns) { + String columnName = identifierMapping.toRemoteColumnName(remoteIdentifiers, column.getName()); + this.verifyColumnName(connection.getMetaData(), columnName); + columnNames.add(columnName); + columnTypes.add(column.getType()); + columnList.add(this.getColumnDefinitionSql(session, column, columnName)); + } + + Optional pageSinkIdColumnName = Optional.empty(); + if (pageSinkIdColumn.isPresent()) { + String columnName = + identifierMapping.toRemoteColumnName(remoteIdentifiers, pageSinkIdColumn.get().getName()); + pageSinkIdColumnName = Optional.of(columnName); + this.verifyColumnName(connection.getMetaData(), columnName); + columnList.add(this.getColumnDefinitionSql(session, pageSinkIdColumn.get(), columnName)); + } + + List columnNamesList = columnNames.build(); + RemoteTableName remoteTableName = new RemoteTableName(Optional.ofNullable(catalog), + Optional.ofNullable(remoteSchema), + remoteTargetTableName); + + for (String sql : this.createTableSqls(remoteTableName, + columnNamesList, + columnList.build(), + tableMetadata, + remoteIdentifiers)) { + log.debug("Execute create table query:" + sql); + this.execute(session, connection, sql); + } + + return new JdbcOutputTableHandle(catalog, + remoteSchema, + remoteTable, + columnNamesList, + columnTypes.build(), + Optional.empty(), + Optional.of(remoteTargetTableName), + pageSinkIdColumnName); + } + + @Override + protected List createTableSqls(RemoteTableName remoteTableName, + List columns, + ConnectorTableMetadata tableMetadata) + { + //this method id override super method, because method of the super class can throw incorrect exception about nonsupporting creating table with comment + throw new UnsupportedOperationException("Unsupported"); + } + + private List createTableSqls( + RemoteTableName remoteTableName, + List columnNames, + List columns, + ConnectorTableMetadata tableMetadata, + RemoteIdentifiers remoteIdentifiers) + { + ImmutableList.Builder createTableSqlsBuilder = ImmutableList.builder(); + String createTableSql = + format("CREATE TABLE %s (%s)", this.quoted(remoteTableName), String.join(", ", columns)); + Map tableProperties = tableMetadata.getProperties(); + Map storageProperties = new HashMap<>(); + Optional appendOptimized = AdbTableProperties.getAppendOptimized(tableProperties); + appendOptimized.ifPresent( + aBoolean -> storageProperties.put("appendoptimized", aBoolean.toString().toUpperCase(Locale.ENGLISH))); + + Optional blockSize = AdbTableProperties.getBlockSize(tableProperties); + blockSize.ifPresent(integer -> storageProperties.put("blocksize", integer.toString())); + + Optional orientation = AdbTableProperties.getOrientation(tableProperties); + orientation.ifPresent( + adbTableStorageOrientation -> storageProperties.put("orientation", adbTableStorageOrientation.name())); + + Optional checksum = AdbTableProperties.getChecksum(tableProperties); + checksum.ifPresent( + aBoolean -> storageProperties.put("checksum", aBoolean.toString().toUpperCase(Locale.ENGLISH))); + + Optional compressType = AdbTableProperties.getCompressType(tableProperties); + compressType.ifPresent(adbTableStorageCompressType -> storageProperties.put("compresstype", + adbTableStorageCompressType.name())); + + Optional compressLevel = AdbTableProperties.getCompressLevel(tableProperties); + compressLevel.ifPresent(integer -> storageProperties.put("compresslevel", integer.toString())); + + Optional fillFactor = AdbTableProperties.getFillFactor(tableProperties); + fillFactor.ifPresent(integer -> storageProperties.put("fillfactor", integer.toString())); + + if (!storageProperties.isEmpty()) { + StringJoiner withJoiner = new StringJoiner(", ", " WITH (", ")"); + + for (Map.Entry entry : storageProperties.entrySet()) { + withJoiner.add(entry.getKey() + " = " + entry.getValue()); + } + createTableSql = createTableSql + withJoiner; + } + + Optional distributed = AdbTableProperties.getDistributed(tableProperties); + Optional> distributedBy = AdbTableProperties.getDistributedBy(tableProperties); + if (distributed.isPresent()) { + if (distributedBy.isPresent()) { + throw new TrinoException( + StandardErrorCode.INVALID_TABLE_PROPERTY, + format("%s property cannot be used when %s is set", "distributed", "distributed_by")); + } + + createTableSql = createTableSql + " DISTRIBUTED " + distributed.get().name(); + } + + if (distributedBy.isPresent()) { + if (distributedBy.get().isEmpty()) { + throw new TrinoException(StandardErrorCode.INVALID_TABLE_PROPERTY, + format("%s property must contain at least one column", "distributed_by")); + } + + StringJoiner distributionColumnJoiner = new StringJoiner(", ", " DISTRIBUTED BY (", ")"); + + for (String distributionColumn : distributedBy.get()) { + String remoteDistributionColumn = + this.getIdentifierMapping().toRemoteColumnName(remoteIdentifiers, distributionColumn); + if (!columnNames.contains(remoteDistributionColumn)) { + throw new TrinoException( + StandardErrorCode.INVALID_TABLE_PROPERTY, + format("%s property references non-existent column %s", "distributed_by", + distributionColumn)); + } + + distributionColumnJoiner.add(remoteDistributionColumn); + } + + createTableSql = createTableSql + distributionColumnJoiner; + } + + createTableSqlsBuilder.add(createTableSql); + Optional tableComment = tableMetadata.getComment(); + if (tableComment.isPresent()) { + createTableSqlsBuilder.add(this.buildTableCommentSql(remoteTableName, tableComment)); + } + + return createTableSqlsBuilder.build(); + } + + @Override + public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) + { + if (tableHandle.isNamedRelation()) { + String objectName = this.quoted(tableHandle.getRequiredNamedRelation().getRemoteTableName()); + return this.metadata.getTableProperties(session, objectName, this.getIdentifierMapping()); + } + else { + return Map.of(); + } + } + + @Override + public void setTableComment(ConnectorSession session, JdbcTableHandle handle, Optional comment) + { + this.execute(session, this.buildTableCommentSql(handle.asPlainTable().getRemoteTableName(), comment)); + } + + private String buildTableCommentSql(RemoteTableName remoteTableName, Optional comment) + { + return format("COMMENT ON TABLE %s IS %s", this.quoted(remoteTableName), + comment.map(BaseJdbcClient::varcharLiteral).orElse("NULL")); + } + + @Override + protected void renameTable( + ConnectorSession session, + Connection connection, + String catalogName, + String remoteSchemaName, + String remoteTableName, + String newRemoteSchemaName, + String newRemoteTableName) + throws SQLException + { + if (!remoteSchemaName.equals(newRemoteSchemaName)) { + throw new TrinoException(NOT_SUPPORTED, "Adb connector does not support renaming tables across schemas"); + } + else { + this.execute( + session, + connection, + format("ALTER TABLE %s RENAME TO %s", this.quoted(catalogName, remoteSchemaName, remoteTableName), + this.quoted(newRemoteTableName))); + } + } + + @Override + public PreparedStatement getPreparedStatement(Connection connection, String sql, Optional columnCount) + throws SQLException + { + // fetch-size is ignored when connection is in auto-commit + connection.setAutoCommit(false); + PreparedStatement statement = connection.prepareStatement(sql); + // This is a heuristic, not exact science. A better formula can perhaps be found with measurements. + // Column count is not known for non-SELECT queries. Not setting fetch size for these. + Optional fetchSize = Optional.ofNullable(Optional.ofNullable(this.fetchSize).orElseGet(() -> + columnCount.map(count -> max(100_000 / count, 1_000)) + .orElse(null))); + if (fetchSize.isPresent()) { + statement.setFetchSize(fetchSize.get()); + } + return statement; + } + + @Override + protected Optional> getTableTypes() + { + return Optional.of(this.tableTypes); + } + + @Override + public List getColumns(ConnectorSession session, JdbcTableHandle tableHandle) + { + if (tableHandle.getColumns().isPresent()) { + return tableHandle.getColumns().get(); + } + checkArgument(tableHandle.isNamedRelation(), "Cannot get columns for %s", tableHandle); + SchemaTableName schemaTableName = tableHandle.getRequiredNamedRelation().getSchemaTableName(); + + try (Connection connection = connectionFactory.openConnection(session)) { + Map arrayColumnDimensions = ImmutableMap.of(); + if (getArrayMapping(session) == AS_ARRAY) { + arrayColumnDimensions = getArrayColumnDimensions(connection, tableHandle); + } + try (ResultSet resultSet = getColumns(tableHandle, connection.getMetaData())) { + int allColumns = 0; + List columns = new ArrayList<>(); + while (resultSet.next()) { + allColumns++; + String columnName = resultSet.getString("COLUMN_NAME"); + JdbcTypeHandle typeHandle = new JdbcTypeHandle( + getInteger(resultSet, "DATA_TYPE").orElseThrow( + () -> new IllegalStateException("DATA_TYPE is null")), + Optional.of(resultSet.getString("TYPE_NAME")), + getInteger(resultSet, "COLUMN_SIZE"), + getInteger(resultSet, "DECIMAL_DIGITS"), + Optional.ofNullable(arrayColumnDimensions.get(columnName)), + Optional.empty()); + Optional columnMapping = toColumnMapping(session, connection, typeHandle); + log.debug("Mapping data type of '%s' column '%s': %s mapped to %s", schemaTableName, columnName, + typeHandle, columnMapping); + // skip unsupported column types + if (columnMapping.isPresent()) { + boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls); + Optional comment = Optional.ofNullable(resultSet.getString("REMARKS")); + columns.add(JdbcColumnHandle.builder() + .setColumnName(columnName) + .setJdbcTypeHandle(typeHandle) + .setColumnType(columnMapping.get().getType()) + .setNullable(nullable) + .setComment(comment) + .build()); + } + if (columnMapping.isEmpty()) { + UnsupportedTypeHandling unsupportedTypeHandling = getUnsupportedTypeHandling(session); + verify( + unsupportedTypeHandling == IGNORE, + "Unsupported type handling is set to %s, but toColumnMapping() returned empty for %s", + unsupportedTypeHandling, + typeHandle); + } + } + if (columns.isEmpty()) { + // A table may have no supported columns. In rare cases a table might have no columns at all. + throw new TableNotFoundException( + schemaTableName, + format("Table '%s' has no supported columns (all %s columns are not supported)", + schemaTableName, allColumns)); + } + return ImmutableList.copyOf(columns); + } + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + private static Map getArrayColumnDimensions(Connection connection, JdbcTableHandle tableHandle) + throws SQLException + { + String sql = "SELECT att.attname, greatest(att.attndims, 1) AS attndims " + + "FROM pg_attribute att " + + " JOIN pg_type attyp ON att.atttypid = attyp.oid" + + " JOIN pg_class tbl ON tbl.oid = att.attrelid " + + " JOIN pg_namespace ns ON tbl.relnamespace = ns.oid " + + "WHERE ns.nspname = ? " + + "AND tbl.relname = ? " + + "AND attyp.typcategory = 'A' "; + try (PreparedStatement statement = connection.prepareStatement(sql)) { + RemoteTableName remoteTableName = tableHandle.getRequiredNamedRelation().getRemoteTableName(); + statement.setString(1, remoteTableName.getSchemaName().orElse(null)); + statement.setString(2, remoteTableName.getTableName()); + + Map arrayColumnDimensions = new HashMap<>(); + try (ResultSet resultSet = statement.executeQuery()) { + while (resultSet.next()) { + arrayColumnDimensions.put(resultSet.getString("attname"), resultSet.getInt("attndims")); + } + } + return arrayColumnDimensions; + } + } + + @Override + public Optional toColumnMapping(ConnectorSession session, Connection connection, + JdbcTypeHandle typeHandle) + { + return dataTypeMapper.toColumnMapping(session, connection, typeHandle); + } + + public ColumnDataType getColumnDataType(ConnectorSession session, JdbcTypeHandle typeHandle) + { + return dataTypeMapper.getColumnDataType(session, typeHandle); + } + + @Override + public WriteMapping toWriteMapping(ConnectorSession session, Type type) + { + return dataTypeMapper.toWriteMapping(session, type); + } + + @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 boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, List sortOrder) + { + for (JdbcSortItem sortItem : sortOrder) { + Type sortItemType = sortItem.column().getColumnType(); + if (sortItemType instanceof CharType || sortItemType instanceof VarcharType) { + if (!isCollatable(sortItem.column())) { + return false; + } + } + } + return true; + } + + @Override + protected Optional topNFunction() + { + return Optional.of((query, sortItems, limit) -> { + String orderBy = sortItems.stream() + .map(sortItem -> { + String ordering = sortItem.sortOrder().isAscending() ? "ASC" : "DESC"; + String nullsHandling = sortItem.sortOrder().isNullsFirst() ? "NULLS FIRST" : "NULLS LAST"; + String collation = ""; + if (isCollatable(sortItem.column())) { + collation = "COLLATE \"C\""; + } + return format("%s %s %s %s", quoted(sortItem.column().getColumnName()), collation, ordering, + nullsHandling); + }) + .collect(joining(", ")); + return format("%s ORDER BY %s LIMIT %d", query, orderBy, limit); + }); + } + + public static boolean isCollatable(JdbcColumnHandle column) + { + if (column.getColumnType() instanceof CharType || column.getColumnType() instanceof VarcharType) { + String jdbcTypeName = column.getJdbcTypeHandle().jdbcTypeName() + .orElseThrow(() -> new TrinoException(JDBC_ERROR, + "Type name is missing: " + column.getJdbcTypeHandle())); + return isCollatable(jdbcTypeName); + } + // non-textual types don't have the concept of collation + return false; + } + + private static boolean isCollatable(String jdbcTypeName) + { + // Only char (internally named bpchar)/varchar/text are the built-in collatable types + return "bpchar".equals(jdbcTypeName) || "varchar".equals(jdbcTypeName) || "text".equals(jdbcTypeName); + } + + @Override + public boolean isTopNGuaranteed(ConnectorSession session) + { + return this.isLimitGuaranteed(session); + } + + @Override + protected Optional> limitFunction() + { + return Optional.of((sql, limit) -> sql + " LIMIT " + limit); + } + + @Override + public boolean isLimitGuaranteed(ConnectorSession session) + { + return AdbSessionProperties.getMaxScanParallelism(session) == 1; + } + + @Override + public OptionalLong delete(ConnectorSession session, JdbcTableHandle handle) + { + checkArgument(handle.isNamedRelation(), "Unable to delete from synthetic table: %s", handle); + checkArgument(handle.getLimit().isEmpty(), "Unable to delete when limit is set: %s", handle); + checkArgument(handle.getSortOrder().isEmpty(), "Unable to delete when sort order is set: %s", handle); + checkArgument(handle.getUpdateAssignments().isEmpty(), "Unable to delete when update assignments are set: %s", + handle); + try (Connection connection = connectionFactory.openConnection(session)) { + verify(connection.getAutoCommit()); + PreparedQuery preparedQuery = queryBuilder.prepareDeleteQuery( + this, + session, + connection, + handle.getRequiredNamedRelation(), + handle.getConstraint(), + getAdditionalPredicate(handle.getConstraintExpressions(), Optional.empty())); + try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(this, session, connection, + preparedQuery, Optional.empty())) { + int affectedRowsCount = preparedStatement.executeUpdate(); + // In getPreparedStatement we set autocommit to false so here we need an explicit commit + connection.commit(); + return OptionalLong.of(affectedRowsCount); + } + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + @Override + public OptionalLong update(ConnectorSession session, JdbcTableHandle handle) + { + checkArgument(handle.isNamedRelation(), "Unable to update from synthetic table: %s", handle); + checkArgument(handle.getLimit().isEmpty(), "Unable to update when limit is set: %s", handle); + checkArgument(handle.getSortOrder().isEmpty(), "Unable to update when sort order is set: %s", handle); + checkArgument(!handle.getUpdateAssignments().isEmpty(), + "Unable to update when update assignments are not set: %s", handle); + try (Connection connection = connectionFactory.openConnection(session)) { + verify(connection.getAutoCommit()); + PreparedQuery preparedQuery = queryBuilder.prepareUpdateQuery( + this, + session, + connection, + handle.getRequiredNamedRelation(), + handle.getConstraint(), + getAdditionalPredicate(handle.getConstraintExpressions(), Optional.empty()), + handle.getUpdateAssignments()); + try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(this, session, connection, + preparedQuery, Optional.empty())) { + int affectedRows = preparedStatement.executeUpdate(); + connection.commit(); + return OptionalLong.of(affectedRows); + } + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + @Override + public OptionalInt getMaxColumnNameLength(ConnectorSession session) + { + return getMaxColumnNameLengthFromDatabaseMetaData(session); + } + + @Override + public TableStatistics getTableStatistics(ConnectorSession session, JdbcTableHandle handle) + { + return statisticsManager.getTableStatistics(session, handle, getColumns(session, handle)); + } + + @Override + public Optional implementJoin( + ConnectorSession session, + JoinType joinType, + PreparedQuery leftSource, + Map leftProjections, + PreparedQuery rightSource, + Map rightProjections, + List joinConditions, + JoinStatistics statistics) + { + if (joinType == JoinType.FULL_OUTER) { + // FULL JOIN is only supported with merge-joinable or hash-joinable join conditions + return Optional.empty(); + } + return implementJoinCostAware( + session, + joinType, + leftSource, + rightSource, + statistics, + () -> super.implementJoin(session, joinType, leftSource, leftProjections, rightSource, rightProjections, + joinConditions, statistics)); + } + + @Override + public Optional legacyImplementJoin( + ConnectorSession session, + JoinType joinType, + PreparedQuery leftSource, + PreparedQuery rightSource, + List joinConditions, + Map rightAssignments, + Map leftAssignments, + JoinStatistics statistics) + { + if (joinType == JoinType.FULL_OUTER) { + // FULL JOIN is only supported with merge-joinable or hash-joinable join conditions + return Optional.empty(); + } + return implementJoinCostAware( + session, + joinType, + leftSource, + rightSource, + statistics, + () -> super.legacyImplementJoin(session, joinType, leftSource, rightSource, joinConditions, + rightAssignments, leftAssignments, statistics)); + } + + @Override + protected boolean isSupportedJoinCondition(ConnectorSession session, JdbcJoinCondition joinCondition) + { + boolean isVarchar = Stream.of(joinCondition.getLeftColumn(), joinCondition.getRightColumn()) + .map(JdbcColumnHandle::getColumnType) + .anyMatch(type -> type instanceof CharType || type instanceof VarcharType); + if (isVarchar) { + JoinCondition.Operator operator = joinCondition.getOperator(); + return switch (operator) { + case LESS_THAN, LESS_THAN_OR_EQUAL, GREATER_THAN, GREATER_THAN_OR_EQUAL -> + AdbSessionProperties.isEnableStringPushdownWithCollate(session); + case EQUAL, NOT_EQUAL, IDENTICAL -> true; + }; + } + + return true; + } + + @Override + protected void verifySchemaName(DatabaseMetaData databaseMetadata, String schemaName) + throws SQLException + { + if (schemaName.length() > databaseMetadata.getMaxSchemaNameLength()) { + throw new TrinoException(NOT_SUPPORTED, + format("Schema name must be shorter than or equal to '%s' characters but got '%s'", + databaseMetadata.getMaxSchemaNameLength(), schemaName.length())); + } + } + + @Override + protected void verifyTableName(DatabaseMetaData databaseMetadata, String tableName) + throws SQLException + { + if (tableName.length() > databaseMetadata.getMaxTableNameLength()) { + throw new TrinoException(NOT_SUPPORTED, + format("Table name must be shorter than or equal to '%s' characters but got '%s'", + databaseMetadata.getMaxTableNameLength(), tableName.length())); + } + } + + @Override + protected void verifyColumnName(DatabaseMetaData databaseMetadata, String columnName) + throws SQLException + { + // PostgreSQL truncates table name to 63 chars silently + // PostgreSQL driver caches the max column name length in a DatabaseMetaData object. The cost to call this method per column is low. + if (columnName.length() > databaseMetadata.getMaxColumnNameLength()) { + throw new TrinoException(NOT_SUPPORTED, + format("Column name must be shorter than or equal to '%s' characters but got '%s': '%s'", + databaseMetadata.getMaxColumnNameLength(), columnName.length(), columnName)); + } + } + + @Override + public void setColumnComment(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column, + Optional comment) + { + // adb doesn't support prepared statement for COMMENT statement + String sql = format( + "COMMENT ON COLUMN %s.%s IS %s", + quoted(handle.asPlainTable().getRemoteTableName()), + quoted(column.getColumnName()), + comment.map(BaseJdbcClient::varcharLiteral).orElse("NULL")); + execute(session, sql); + } + + public List getColumnDataTypes(ConnectorSession session, JdbcOutputTableHandle outputTableHandle) + { + return dataTypeMapper.getColumnDataTypes(session, outputTableHandle); + } + + public String getTargetTableName(JdbcOutputTableHandle tableHandle) + { + return this.quoted(tableHandle.getCatalogName(), + tableHandle.getSchemaName(), + tableHandle.getTemporaryTableName().orElseGet(tableHandle::getTableName)); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/CollationAwareQueryBuilder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/CollationAwareQueryBuilder.java new file mode 100644 index 000000000000..586a8a04f5d8 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/CollationAwareQueryBuilder.java @@ -0,0 +1,74 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector; + +import com.google.inject.Inject; +import io.trino.plugin.jdbc.DefaultQueryBuilder; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcJoinCondition; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.QueryParameter; +import io.trino.plugin.jdbc.WriteFunction; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.type.Type; + +import java.util.Optional; +import java.util.function.Consumer; +import java.util.stream.Stream; + +public class CollationAwareQueryBuilder + extends DefaultQueryBuilder +{ + @Inject + public CollationAwareQueryBuilder(RemoteQueryModifier queryModifier) + { + super(queryModifier); + } + + protected String formatJoinCondition(JdbcClient client, String leftRelationAlias, String rightRelationAlias, JdbcJoinCondition condition) + { + boolean isCollatable = Stream.of(condition.getLeftColumn(), condition.getRightColumn()).anyMatch(AdbSqlClient::isCollatable); + return isCollatable + ? String.format( + "%s.%s COLLATE \"C\" %s %s.%s COLLATE \"C\"", + leftRelationAlias, + this.buildJoinColumn(client, condition.getLeftColumn()), + condition.getOperator().getValue(), + rightRelationAlias, + this.buildJoinColumn(client, condition.getRightColumn())) + : super.formatJoinCondition(client, leftRelationAlias, rightRelationAlias, condition); + } + + protected String toPredicate( + JdbcClient client, + ConnectorSession session, + JdbcColumnHandle column, + JdbcTypeHandle jdbcType, + Type type, + WriteFunction writeFunction, + String operator, + Object value, + Consumer accumulator) + { + if (AdbSqlClient.isCollatable(column) && AdbSessionProperties.isEnableStringPushdownWithCollate(session)) { + accumulator.accept(new QueryParameter(jdbcType, type, Optional.of(value))); + return String.format("%s %s %s COLLATE \"C\"", client.quoted(column.getColumnName()), operator, writeFunction.getBindExpression()); + } + else { + return super.toPredicate(client, session, column, jdbcType, type, writeFunction, operator, value, accumulator); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BigintDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BigintDataType.java new file mode 100644 index 000000000000..966b8e840b10 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BigintDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class BigintDataType + implements ColumnDataType +{ + private final String name; + + public BigintDataType() + { + this.name = "bigint"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.BIGINT; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BitDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BitDataType.java new file mode 100644 index 000000000000..567525425e4a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BitDataType.java @@ -0,0 +1,39 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import java.util.Optional; + +public class BitDataType + implements ColumnDataType +{ + private final String name; + + public BitDataType(Optional length) + { + this.name = length.map(len -> String.format("bit(%d)", len)).orElse("bit"); + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.BIT; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BooleanDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BooleanDataType.java new file mode 100644 index 000000000000..e72c64a7c2dd --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/BooleanDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class BooleanDataType + implements ColumnDataType +{ + private final String name; + + public BooleanDataType() + { + this.name = "boolean"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.BOOLEAN; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ByteaDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ByteaDataType.java new file mode 100644 index 000000000000..1c19b45d2956 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ByteaDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class ByteaDataType + implements ColumnDataType +{ + private final String name; + + public ByteaDataType() + { + this.name = "bytea"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.BYTEA; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/CharDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/CharDataType.java new file mode 100644 index 000000000000..cf4c799d70b8 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/CharDataType.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.CharType; + +public class CharDataType + implements ColumnDataType +{ + private final String name; + private final CharType charType; + + public CharDataType(CharType charType) + { + this.name = String.format("char(%d)", charType.getLength()); + this.charType = charType; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.CHAR; + } + + public CharType getCharType() + { + return charType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ColumnDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ColumnDataType.java new file mode 100644 index 000000000000..21e9371197e9 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ColumnDataType.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public interface ColumnDataType +{ + String getName(); + + ConnectorDataType getType(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ConnectorDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ConnectorDataType.java new file mode 100644 index 000000000000..2c7ccb00cac5 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/ConnectorDataType.java @@ -0,0 +1,41 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public enum ConnectorDataType +{ + BOOLEAN, + MONEY, + UUID, + JSONB, + BIT, + BIGINT, + BYTEA, + CHAR, + VARCHAR, + DECIMAL_SHORT, + DECIMAL_LONG, + INTEGER, + SMALLINT, + REAL, + DOUBLE_PRECISION, + DATE, + TIME, + TIMESTAMP_SHORT_WITH_TIME_ZONE, + TIMESTAMP_LONG_WITH_TIME_ZONE, + TIMESTAMP_WITHOUT_TIME_ZONE, + ENUM, + ARRAY, + UNKNOWN +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DateDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DateDataType.java new file mode 100644 index 000000000000..21c4d2ef5460 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DateDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class DateDataType + implements ColumnDataType +{ + private final String name; + + public DateDataType() + { + this.name = "date"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.DATE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalDataType.java new file mode 100644 index 000000000000..9591b71d0172 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalDataType.java @@ -0,0 +1,40 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.DecimalType; + +public abstract class DecimalDataType + implements ColumnDataType +{ + private final String name; + private final DecimalType decimalType; + + public DecimalDataType(DecimalType decimalType) + { + this.name = String.format("decimal(%d,%d)", decimalType.getPrecision(), decimalType.getScale()); + this.decimalType = decimalType; + } + + @Override + public String getName() + { + return name; + } + + public DecimalType getDecimalType() + { + return decimalType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalLongDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalLongDataType.java new file mode 100644 index 000000000000..f6de444fcb81 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalLongDataType.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.DecimalType; + +public class DecimalLongDataType + extends DecimalDataType +{ + public DecimalLongDataType(DecimalType decimalType) + { + super(decimalType); + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.DECIMAL_LONG; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalShortDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalShortDataType.java new file mode 100644 index 000000000000..dc26e98a4732 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DecimalShortDataType.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.DecimalType; + +public class DecimalShortDataType + extends DecimalDataType +{ + public DecimalShortDataType(DecimalType decimalType) + { + super(decimalType); + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.DECIMAL_SHORT; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DoubleDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DoubleDataType.java new file mode 100644 index 000000000000..7cc6bd93359b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/DoubleDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class DoubleDataType + implements ColumnDataType +{ + private final String name; + + public DoubleDataType() + { + this.name = "double precision"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.DOUBLE_PRECISION; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/IntegerDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/IntegerDataType.java new file mode 100644 index 000000000000..90d6d0dcfa53 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/IntegerDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class IntegerDataType + implements ColumnDataType +{ + private final String name; + + public IntegerDataType() + { + this.name = "integer"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.INTEGER; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/JsonbDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/JsonbDataType.java new file mode 100644 index 000000000000..80f6f9b92378 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/JsonbDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class JsonbDataType + implements ColumnDataType +{ + private final String name; + + public JsonbDataType() + { + this.name = "jsonb"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.JSONB; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/MoneyDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/MoneyDataType.java new file mode 100644 index 000000000000..a69dc5acc83a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/MoneyDataType.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.VarcharType; + +public class MoneyDataType + implements ColumnDataType +{ + private final String name; + private final VarcharType varcharType; + + public MoneyDataType() + { + this.name = "money"; + this.varcharType = VarcharType.createUnboundedVarcharType(); + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.MONEY; + } + + public VarcharType getVarcharType() + { + return varcharType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/RealDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/RealDataType.java new file mode 100644 index 000000000000..790bcb5ee26b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/RealDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class RealDataType + implements ColumnDataType +{ + private final String name; + + public RealDataType() + { + this.name = "real"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.REAL; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/SmallintDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/SmallintDataType.java new file mode 100644 index 000000000000..ba77a83fd586 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/SmallintDataType.java @@ -0,0 +1,53 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.SmallintType; +import io.trino.spi.type.TinyintType; +import io.trino.spi.type.Type; + +import static com.google.common.base.Verify.verify; + +public class SmallintDataType + implements ColumnDataType +{ + private final String name; + private final Type trinoType; + + public SmallintDataType(Type trinoType) + { + verify(trinoType == TinyintType.TINYINT + || trinoType == SmallintType.SMALLINT, + "Unexpected trino type " + trinoType + " for " + getType() + " data type"); + this.name = "smallint"; + this.trinoType = trinoType; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.SMALLINT; + } + + public Type getTrinoType() + { + return trinoType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimeDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimeDataType.java new file mode 100644 index 000000000000..d6eb88121d4f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimeDataType.java @@ -0,0 +1,44 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class TimeDataType + implements ColumnDataType +{ + private final String name; + private final int precision; + + public TimeDataType(int precision) + { + this.name = String.format("time(%d)", precision); + this.precision = precision; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.TIME; + } + + public int getPrecision() + { + return precision; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampLongWithTimeZoneDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampLongWithTimeZoneDataType.java new file mode 100644 index 000000000000..0d36134ab164 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampLongWithTimeZoneDataType.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class TimestampLongWithTimeZoneDataType + extends TimestampWithTimeZoneDataType +{ + public TimestampLongWithTimeZoneDataType(int precision) + { + super(precision); + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.TIMESTAMP_LONG_WITH_TIME_ZONE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampShortWithTimeZoneDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampShortWithTimeZoneDataType.java new file mode 100644 index 000000000000..8dbabbc62e4b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampShortWithTimeZoneDataType.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class TimestampShortWithTimeZoneDataType + extends TimestampWithTimeZoneDataType +{ + public TimestampShortWithTimeZoneDataType(int precision) + { + super(precision); + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.TIMESTAMP_SHORT_WITH_TIME_ZONE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithTimeZoneDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithTimeZoneDataType.java new file mode 100644 index 000000000000..ac4ab87c7bd3 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithTimeZoneDataType.java @@ -0,0 +1,48 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.TimestampWithTimeZoneType; +import io.trino.spi.type.Type; + +public abstract class TimestampWithTimeZoneDataType + implements ColumnDataType +{ + private final String name; + private final Type trinoType; + private final int precision; + + public TimestampWithTimeZoneDataType(int precision) + { + this.name = String.format("timestamptz(%d)", precision); + this.trinoType = TimestampWithTimeZoneType.createTimestampWithTimeZoneType(precision); + this.precision = precision; + } + + @Override + public String getName() + { + return name; + } + + public Type getTrinoType() + { + return trinoType; + } + + public int getPrecision() + { + return precision; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithoutTimeZoneDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithoutTimeZoneDataType.java new file mode 100644 index 000000000000..22a92e060220 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/TimestampWithoutTimeZoneDataType.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.TimestampType; + +public class TimestampWithoutTimeZoneDataType + implements ColumnDataType +{ + private final String name; + private final TimestampType timestampType; + + public TimestampWithoutTimeZoneDataType(TimestampType timestampType) + { + this.name = String.format("timestamp(%d)", timestampType.getPrecision()); + this.timestampType = timestampType; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.TIMESTAMP_WITHOUT_TIME_ZONE; + } + + public TimestampType getTimestampType() + { + return timestampType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UnknownDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UnknownDataType.java new file mode 100644 index 000000000000..102c448c1c28 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UnknownDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class UnknownDataType + implements ColumnDataType +{ + private final String name; + + public UnknownDataType() + { + this.name = "unknown"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.UNKNOWN; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UuidDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UuidDataType.java new file mode 100644 index 000000000000..6eaf9b0d84ed --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/UuidDataType.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +public class UuidDataType + implements ColumnDataType +{ + private final String name; + + public UuidDataType() + { + this.name = "uuid"; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.UUID; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/VarcharDataType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/VarcharDataType.java new file mode 100644 index 000000000000..b0a14640da15 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/VarcharDataType.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype; + +import io.trino.spi.type.VarcharType; + +public class VarcharDataType + implements ColumnDataType +{ + private final String name; + private final VarcharType varcharType; + + public VarcharDataType(VarcharType varcharType) + { + this.name = varcharType.isUnbounded() ? "varchar" : String.format("varchar(%d)", varcharType.getBoundedLength()); + this.varcharType = varcharType; + } + + @Override + public String getName() + { + return name; + } + + @Override + public ConnectorDataType getType() + { + return ConnectorDataType.VARCHAR; + } + + public VarcharType getVarcharType() + { + return varcharType; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapper.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapper.java new file mode 100644 index 000000000000..1e4d435a6151 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapper.java @@ -0,0 +1,39 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype.mapper; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.jdbc.ColumnMapping; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.WriteMapping; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.type.Type; + +import java.sql.Connection; +import java.util.List; +import java.util.Optional; + +public interface DataTypeMapper +{ + Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle); + + WriteMapping toWriteMapping(ConnectorSession session, Type type); + + ColumnDataType getColumnDataType(ConnectorSession session, JdbcTypeHandle typeHandle); + + List getColumnDataTypes(ConnectorSession session, JdbcOutputTableHandle outputTableHandle); + + Optional fromTrinoType(Type type); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapperImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapperImpl.java new file mode 100644 index 000000000000..611ec987c415 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/datatype/mapper/DataTypeMapperImpl.java @@ -0,0 +1,920 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.datatype.mapper; + +import com.google.common.collect.ImmutableSortedSet; +import com.google.common.math.LongMath; +import com.google.inject.Inject; +import io.airlift.slice.Slice; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.AdbColumnMapping; +import io.trino.plugin.adb.connector.AdbSessionProperties; +import io.trino.plugin.adb.connector.datatype.BigintDataType; +import io.trino.plugin.adb.connector.datatype.BitDataType; +import io.trino.plugin.adb.connector.datatype.BooleanDataType; +import io.trino.plugin.adb.connector.datatype.ByteaDataType; +import io.trino.plugin.adb.connector.datatype.CharDataType; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.datatype.ConnectorDataType; +import io.trino.plugin.adb.connector.datatype.DateDataType; +import io.trino.plugin.adb.connector.datatype.DecimalLongDataType; +import io.trino.plugin.adb.connector.datatype.DecimalShortDataType; +import io.trino.plugin.adb.connector.datatype.DoubleDataType; +import io.trino.plugin.adb.connector.datatype.IntegerDataType; +import io.trino.plugin.adb.connector.datatype.JsonbDataType; +import io.trino.plugin.adb.connector.datatype.MoneyDataType; +import io.trino.plugin.adb.connector.datatype.RealDataType; +import io.trino.plugin.adb.connector.datatype.SmallintDataType; +import io.trino.plugin.adb.connector.datatype.TimeDataType; +import io.trino.plugin.adb.connector.datatype.TimestampLongWithTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.TimestampShortWithTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.TimestampWithoutTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.UnknownDataType; +import io.trino.plugin.adb.connector.datatype.UuidDataType; +import io.trino.plugin.adb.connector.datatype.VarcharDataType; +import io.trino.plugin.jdbc.BaseJdbcConfig; +import io.trino.plugin.jdbc.BooleanReadFunction; +import io.trino.plugin.jdbc.ColumnMapping; +import io.trino.plugin.jdbc.DoubleReadFunction; +import io.trino.plugin.jdbc.JdbcMetadataSessionProperties; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.JdbcTypeHandle; +import io.trino.plugin.jdbc.LongReadFunction; +import io.trino.plugin.jdbc.LongWriteFunction; +import io.trino.plugin.jdbc.ObjectReadFunction; +import io.trino.plugin.jdbc.ObjectWriteFunction; +import io.trino.plugin.jdbc.PredicatePushdownController; +import io.trino.plugin.jdbc.ReadFunction; +import io.trino.plugin.jdbc.SliceReadFunction; +import io.trino.plugin.jdbc.SliceWriteFunction; +import io.trino.plugin.jdbc.WriteMapping; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.predicate.Domain; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.CharType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Decimals; +import io.trino.spi.type.LongTimestamp; +import io.trino.spi.type.LongTimestampWithTimeZone; +import io.trino.spi.type.TimeType; +import io.trino.spi.type.TimestampType; +import io.trino.spi.type.TimestampWithTimeZoneType; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeSignature; +import io.trino.spi.type.UuidType; +import io.trino.spi.type.VarcharType; +import org.postgresql.core.TypeInfo; +import org.postgresql.jdbc.PgConnection; + +import java.io.IOException; +import java.math.RoundingMode; +import java.sql.Array; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.sql.Types; +import java.time.Instant; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; +import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.adb.AdbPluginConfig.ArrayMapping.AS_ARRAY; +import static io.trino.plugin.adb.AdbPluginConfig.ArrayMapping.AS_JSON; +import static io.trino.plugin.adb.AdbPluginConfig.ArrayMapping.DISABLED; +import static io.trino.plugin.adb.TypeUtil.TIME_TYPE_FORMATTER; +import static io.trino.plugin.adb.TypeUtil.arrayDepth; +import static io.trino.plugin.adb.TypeUtil.getArrayElementPgTypeName; +import static io.trino.plugin.adb.TypeUtil.getJdbcObjectArray; +import static io.trino.plugin.adb.TypeUtil.toPgTimestamp; +import static io.trino.plugin.adb.connector.AdbSessionProperties.getArrayMapping; +import static io.trino.plugin.base.util.JsonTypeUtil.jsonParse; +import static io.trino.plugin.base.util.JsonTypeUtil.toJsonValue; +import static io.trino.plugin.jdbc.DecimalConfig.DecimalMapping.ALLOW_OVERFLOW; +import static io.trino.plugin.jdbc.DecimalSessionSessionProperties.getDecimalDefaultScale; +import static io.trino.plugin.jdbc.DecimalSessionSessionProperties.getDecimalRounding; +import static io.trino.plugin.jdbc.DecimalSessionSessionProperties.getDecimalRoundingMode; +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static io.trino.plugin.jdbc.PredicatePushdownController.DISABLE_PUSHDOWN; +import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping; +import static io.trino.plugin.jdbc.StandardColumnMappings.bigintWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.booleanColumnMapping; +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.dateColumnMappingUsingLocalDate; +import static io.trino.plugin.jdbc.StandardColumnMappings.dateWriteFunctionUsingLocalDate; +import static io.trino.plugin.jdbc.StandardColumnMappings.decimalColumnMapping; +import static io.trino.plugin.jdbc.StandardColumnMappings.doubleColumnMapping; +import static io.trino.plugin.jdbc.StandardColumnMappings.doubleWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.fromTrinoTimestamp; +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.timestampReadFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.tinyintWriteFunction; +import static io.trino.plugin.jdbc.StandardColumnMappings.varbinaryColumnMapping; +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 static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; +import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc; +import static io.trino.spi.type.DateType.DATE; +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.TimeType.createTimeType; +import static io.trino.spi.type.TimeZoneKey.UTC_KEY; +import static io.trino.spi.type.TimestampType.createTimestampType; +import static io.trino.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType; +import static io.trino.spi.type.Timestamps.MILLISECONDS_PER_SECOND; +import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_DAY; +import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MILLISECOND; +import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_DAY; +import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_MICROSECOND; +import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND; +import static io.trino.spi.type.Timestamps.round; +import static io.trino.spi.type.TinyintType.TINYINT; +import static io.trino.spi.type.UuidType.javaUuidToTrinoUuid; +import static io.trino.spi.type.UuidType.trinoUuidToJavaUuid; +import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; +import static java.lang.Math.floorDiv; +import static java.lang.Math.floorMod; +import static java.lang.Math.max; +import static java.lang.Math.min; +import static java.lang.String.CASE_INSENSITIVE_ORDER; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class DataTypeMapperImpl + implements DataTypeMapper +{ + private static final String COLUMN_TYPE_NOT_SUPPORTED_ERROR_MSG_TEMPLATE = "Column type %s is not supported"; + private static final int POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION = 6; + private static final int PRECISION_OF_UNSPECIFIED_DECIMAL = 0; + private static final int ARRAY_RESULT_SET_VALUE_COLUMN = 2; + private static final PredicatePushdownController ADB_STRING_COLLATION_AWARE_PUSHDOWN = (session, domain) -> { + if (domain.isOnlyNull()) { + return PredicatePushdownController.FULL_PUSHDOWN.apply(session, domain); + } + else if (AdbSessionProperties.isEnableStringPushdownWithCollate(session)) { + return PredicatePushdownController.FULL_PUSHDOWN.apply(session, domain); + } + else { + Domain simplifiedDomain = domain.simplify(JdbcMetadataSessionProperties.getDomainCompactionThreshold(session)); + return !simplifiedDomain.getValues().isDiscreteSet() + ? DISABLE_PUSHDOWN.apply(session, domain) + : PredicatePushdownController.FULL_PUSHDOWN.apply(session, simplifiedDomain); + } + }; + private final Type jsonType; + private final Type uuidType; + private final Set jdbcTypesMappedToVarchar; + + @Inject + public DataTypeMapperImpl(TypeManager typeManager, BaseJdbcConfig jdbcConfig) + { + this.jsonType = typeManager.getType(new TypeSignature("json")); + this.uuidType = typeManager.getType(new TypeSignature("uuid")); + this.jdbcTypesMappedToVarchar = ImmutableSortedSet.orderedBy(CASE_INSENSITIVE_ORDER) + .addAll(requireNonNull(jdbcConfig.getJdbcTypesMappedToVarchar(), "jdbcTypesMappedToVarchar is null")) + .build(); + } + + @Override + public Optional toColumnMapping(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) + { + return Optional.ofNullable(toColumnMappingInternal(session, Optional.of(connection), typeHandle)) + .map(AdbColumnMapping::columnMapping); + } + + @Override + public WriteMapping toWriteMapping(ConnectorSession session, Type type) + { + if (type == BOOLEAN) { + return WriteMapping.booleanMapping("boolean", booleanWriteFunction()); + } + if (type == TINYINT) { + return WriteMapping.longMapping("smallint", tinyintWriteFunction()); + } + if (type == SMALLINT) { + return WriteMapping.longMapping("smallint", smallintWriteFunction()); + } + if (type == INTEGER) { + return WriteMapping.longMapping("integer", integerWriteFunction()); + } + if (type == BIGINT) { + return WriteMapping.longMapping("bigint", bigintWriteFunction()); + } + if (type == REAL) { + return WriteMapping.longMapping("real", realWriteFunction()); + } + if (type == DOUBLE) { + return WriteMapping.doubleMapping("double precision", doubleWriteFunction()); + } + 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)); + } + if (type instanceof CharType) { + return WriteMapping.sliceMapping("char(" + ((CharType) type).getLength() + ")", charWriteFunction()); + } + if (type instanceof VarcharType varcharType) { + String dataType; + if (varcharType.isUnbounded()) { + dataType = "varchar"; + } + else { + dataType = "varchar(" + varcharType.getBoundedLength() + ")"; + } + return WriteMapping.sliceMapping(dataType, varcharWriteFunction()); + } + if (VARBINARY.equals(type)) { + return WriteMapping.sliceMapping("bytea", varbinaryWriteFunction()); + } + if (type == DATE) { + return WriteMapping.longMapping("date", dateWriteFunctionUsingLocalDate()); + } + if (type instanceof TimeType timeType) { + verify(timeType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION); + return WriteMapping.longMapping(format("time(%s)", timeType.getPrecision()), timeWriteFunction(timeType.getPrecision())); + } + if (type instanceof TimestampType timestampType) { + if (timestampType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) { + return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), DataTypeMapperImpl::shortTimestampWriteFunction); + } + else { + return WriteMapping.objectMapping(format("timestamp(%s)", POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION), longTimestampWriteFunction()); + } + } + if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType) { + verify(timestampWithTimeZoneType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION); + String dataType = format("timestamptz(%d)", timestampWithTimeZoneType.getPrecision()); + if (timestampWithTimeZoneType.isShort()) { + return WriteMapping.longMapping(dataType, shortTimestampWithTimeZoneWriteFunction()); + } + return WriteMapping.objectMapping(dataType, longTimestampWithTimeZoneWriteFunction()); + } + if (type.equals(jsonType)) { + return WriteMapping.sliceMapping("jsonb", typedVarcharWriteFunction("json")); + } + if (type.equals(uuidType)) { + return WriteMapping.sliceMapping("uuid", uuidWriteFunction()); + } + if (type instanceof ArrayType arrayType && getArrayMapping(session) == AS_ARRAY) { + Type elementType = arrayType.getElementType(); + String elementDataType = toWriteMapping(session, elementType).getDataType(); + return WriteMapping.objectMapping(elementDataType + "[]", arrayWriteFunction(session, elementType, getArrayElementPgTypeName(session, this, elementType))); + } + throw new TrinoException(NOT_SUPPORTED, "Unsupported column type: " + type.getDisplayName()); + } + + @Override + public ColumnDataType getColumnDataType(ConnectorSession session, JdbcTypeHandle typeHandle) + { + return Optional.ofNullable(toColumnMappingInternal(session, Optional.empty(), typeHandle)) + .map(AdbColumnMapping::columnDataType) + .orElseThrow(() -> new IllegalArgumentException("Failed to get column data type for type: " + typeHandle)); + } + + private AdbColumnMapping toColumnMappingInternal(ConnectorSession session, Optional connection, JdbcTypeHandle typeHandle) + { + String jdbcTypeName = typeHandle.jdbcTypeName() + .orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + typeHandle)); + Optional mapping = getForcedMappingToVarchar(typeHandle); + if (mapping.isPresent()) { + return new AdbColumnMapping(mapping.get(), new UnknownDataType()); + } + switch (jdbcTypeName) { + case "bool": + return new AdbColumnMapping(booleanColumnMapping(), new BooleanDataType()); + case "money": + return new AdbColumnMapping(moneyColumnMapping(), new MoneyDataType()); + case "uuid": + return new AdbColumnMapping(uuidColumnMapping(), new UuidDataType()); + case "jsonb": + case "json": + return new AdbColumnMapping(jsonColumnMapping(), new JsonbDataType()); + case "timestamptz": + int decimalDigits = typeHandle.requiredDecimalDigits(); + return timestampWithTimeZoneColumnMapping(decimalDigits); + } + switch (typeHandle.jdbcType()) { + case Types.BIT: + return new AdbColumnMapping(booleanColumnMapping(), new BitDataType(typeHandle.columnSize())); + case Types.SMALLINT: + return new AdbColumnMapping(smallintColumnMapping(), new SmallintDataType(SMALLINT)); + case Types.INTEGER: + return new AdbColumnMapping(integerColumnMapping(), new IntegerDataType()); + case Types.BIGINT: + return new AdbColumnMapping(bigintColumnMapping(), new BigintDataType()); + case Types.REAL: + return new AdbColumnMapping(realColumnMapping(), new RealDataType()); + case Types.DOUBLE: + return new AdbColumnMapping(doubleColumnMapping(), new DoubleDataType()); + case Types.NUMERIC: { + int columnSize = typeHandle.requiredColumnSize(); + int precision; + int decimalDigits = typeHandle.decimalDigits().orElse(0); + if (getDecimalRounding(session) == ALLOW_OVERFLOW) { + if (columnSize == PRECISION_OF_UNSPECIFIED_DECIMAL) { + // decimal type with unspecified scale - up to 131072 digits before the decimal point; up to 16383 digits after the decimal point) + DecimalType decimalType = createDecimalType(Decimals.MAX_PRECISION, getDecimalDefaultScale(session)); + return new AdbColumnMapping(decimalColumnMapping(decimalType, getDecimalRoundingMode(session)), new UnknownDataType()); + } + precision = columnSize; + if (precision > Decimals.MAX_PRECISION) { + int scale = min(decimalDigits, getDecimalDefaultScale(session)); + DecimalType decimalType = createDecimalType(Decimals.MAX_PRECISION, scale); + return new AdbColumnMapping(decimalColumnMapping(decimalType, getDecimalRoundingMode(session)), new UnknownDataType()); + } + } + precision = columnSize + max(-decimalDigits, 0); + if (columnSize != PRECISION_OF_UNSPECIFIED_DECIMAL && precision <= Decimals.MAX_PRECISION) { + DecimalType decimalType = createDecimalType(precision, max(decimalDigits, 0)); + ColumnDataType columnDataType = decimalType.isShort() + ? new DecimalShortDataType(decimalType) + : new DecimalLongDataType(decimalType); + return new AdbColumnMapping(decimalColumnMapping(decimalType, RoundingMode.UNNECESSARY), + columnDataType); + } + break; + } + case Types.CHAR: + ColumnMapping charColumnMapping = charColumnMapping(typeHandle.requiredColumnSize()); + return new AdbColumnMapping(charColumnMapping, new CharDataType((CharType) charColumnMapping.getType())); + case Types.VARCHAR: + if (!jdbcTypeName.equals("varchar")) { + return new AdbColumnMapping(enumColumnMapping(session, jdbcTypeName), new UnknownDataType()); + } + ColumnMapping varcharColumnMapping = varcharColumnMapping(typeHandle.requiredColumnSize()); + return new AdbColumnMapping(varcharColumnMapping, new VarcharDataType((VarcharType) varcharColumnMapping.getType())); + case Types.BINARY: + ColumnMapping columnMapping = varbinaryColumnMapping(); + if (jdbcTypeName.equals("bytea")) { + return new AdbColumnMapping(columnMapping, new ByteaDataType()); + } + return new AdbColumnMapping(varbinaryColumnMapping(), new UnknownDataType()); + case Types.DATE: + return new AdbColumnMapping(dateColumnMappingUsingLocalDate(), new DateDataType()); + case Types.TIME: + int requiredDecimalDigits = typeHandle.requiredDecimalDigits(); + return new AdbColumnMapping(timeColumnMapping(requiredDecimalDigits), new TimeDataType(requiredDecimalDigits)); + case Types.TIMESTAMP: + TimestampType timestampType = createTimestampType(typeHandle.requiredDecimalDigits()); + return new AdbColumnMapping( + ColumnMapping.longMapping( + timestampType, timestampReadFunction(timestampType), DataTypeMapperImpl::shortTimestampWriteFunction), + new TimestampWithoutTimeZoneDataType(timestampType)); + case Types.ARRAY: + if (connection.isPresent()) { + Optional arrayColumnMapping = arrayToTrinoType(session, connection.get(), typeHandle); + if (arrayColumnMapping.isPresent()) { + return new AdbColumnMapping(arrayColumnMapping.get(), new UnknownDataType()); + } + } + break; + } + if (getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR) { + Optional columnMapping = mapToUnboundedVarchar(typeHandle); + if (columnMapping.isPresent()) { + return new AdbColumnMapping(columnMapping.get(), new UnknownDataType()); + } + } + return null; + } + + @Override + public Optional fromTrinoType(Type type) + { + if (type == BOOLEAN) { + return Optional.of(new BigintDataType()); + } + else if (type == TINYINT) { + return Optional.of(new SmallintDataType(TINYINT)); + } + else if (type == SMALLINT) { + return Optional.of(new SmallintDataType(SMALLINT)); + } + else if (type == INTEGER) { + return Optional.of(new IntegerDataType()); + } + else if (type == BIGINT) { + return Optional.of(new BigintDataType()); + } + else if (type == REAL) { + return Optional.of(new RealDataType()); + } + else if (type == DOUBLE) { + return Optional.of(new DoubleDataType()); + } + else if (type instanceof DecimalType decimalType) { + return decimalType.isShort() + ? Optional.of(new DecimalShortDataType(decimalType)) + : Optional.of(new DecimalLongDataType(decimalType)); + } + else if (type == DATE) { + return Optional.of(new DateDataType()); + } + else if (type instanceof TimeType timeType) { + return Optional.of(new TimeDataType(timeType.getPrecision())); + } + else if (type instanceof TimestampType timestampType) { + return Optional.of(new TimestampWithoutTimeZoneDataType(timestampType)); + } + else if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType) { + return timestampWithTimeZoneType.isShort() + ? Optional.of(new TimestampShortWithTimeZoneDataType(timestampWithTimeZoneType.getPrecision())) + : Optional.of(new TimestampLongWithTimeZoneDataType(timestampWithTimeZoneType.getPrecision())); + } + else if (type instanceof CharType charType) { + return Optional.of(new CharDataType(charType)); + } + else if (type instanceof VarcharType varcharType) { + return Optional.of(new VarcharDataType(varcharType)); + } + else if (type == VARBINARY) { + return Optional.of(new ByteaDataType()); + } + else { + return type == UuidType.UUID ? Optional.of(new UuidDataType()) : Optional.empty(); + } + } + + @Override + public List getColumnDataTypes(ConnectorSession session, JdbcOutputTableHandle outputTableHandle) + { + List columnDataTypes = new ArrayList<>(); + for (int i = 0; i < outputTableHandle.getColumnNames().size(); i++) { + ColumnDataType columnDataType; + if (outputTableHandle.getJdbcColumnTypes().isEmpty()) { + Type columnType = outputTableHandle.getColumnTypes().get(i); + columnDataType = fromTrinoType(columnType) + .orElseThrow(() -> new TrinoException(NOT_SUPPORTED, + format(COLUMN_TYPE_NOT_SUPPORTED_ERROR_MSG_TEMPLATE, columnType))); + } + else { + JdbcTypeHandle columnType = (outputTableHandle.getJdbcColumnTypes().get()).get(i); + columnDataType = Optional.ofNullable(toColumnMappingInternal(session, + Optional.empty(), + columnType)) + .orElseThrow(() -> new TrinoException(NOT_SUPPORTED, + format(COLUMN_TYPE_NOT_SUPPORTED_ERROR_MSG_TEMPLATE, columnType))) + .columnDataType(); + if (columnDataType.getType() == ConnectorDataType.UNKNOWN) { + throw new TrinoException(NOT_SUPPORTED, + format(COLUMN_TYPE_NOT_SUPPORTED_ERROR_MSG_TEMPLATE, columnType)); + } + } + columnDataTypes.add(columnDataType); + } + return columnDataTypes; + } + + private Optional getForcedMappingToVarchar(JdbcTypeHandle typeHandle) + { + if (typeHandle.jdbcTypeName().isPresent() && jdbcTypesMappedToVarchar.contains(typeHandle.jdbcTypeName().get())) { + return mapToUnboundedVarchar(typeHandle); + } + return Optional.empty(); + } + + protected static Optional mapToUnboundedVarchar(JdbcTypeHandle typeHandle) + { + VarcharType unboundedVarcharType = createUnboundedVarcharType(); + return Optional.of(ColumnMapping.sliceMapping( + unboundedVarcharType, + varcharReadFunction(unboundedVarcharType), + (statement, index, value) -> { + throw new TrinoException( + NOT_SUPPORTED, + "Underlying type that is mapped to VARCHAR is not supported for INSERT: " + typeHandle.jdbcTypeName().get()); + }, + DISABLE_PUSHDOWN)); + } + + private static ColumnMapping moneyColumnMapping() + { + return ColumnMapping.sliceMapping( + VARCHAR, + new SliceReadFunction() + { + @Override + public boolean isNull(ResultSet resultSet, int columnIndex) + throws SQLException + { + resultSet.getString(columnIndex); + return resultSet.wasNull(); + } + + @Override + public Slice readSlice(ResultSet resultSet, int columnIndex) + throws SQLException + { + return utf8Slice(resultSet.getString(columnIndex)); + } + }, + (statement, index, value) -> { + throw new TrinoException(NOT_SUPPORTED, "Money type is not supported for INSERT"); + }, + DISABLE_PUSHDOWN); + } + + private static LongWriteFunction timeWriteFunction(int precision) + { + checkArgument(precision <= 6, "Unsupported precision: %s", precision); // PostgreSQL limit but also assumption within this method + String bindExpression = format("CAST(? AS time(%s))", precision); + return new LongWriteFunction() + { + @Override + public String getBindExpression() + { + return bindExpression; + } + + @Override + public void set(PreparedStatement statement, int index, long picosOfDay) + throws SQLException + { + picosOfDay = round(picosOfDay, 12 - precision); + if (picosOfDay == PICOSECONDS_PER_DAY) { + picosOfDay = 0; + } + LocalTime localTime = LocalTime.ofNanoOfDay(picosOfDay / PICOSECONDS_PER_NANOSECOND); + // statement.setObject(.., localTime) would yield incorrect end result for 23:59:59.999000 + statement.setString(index, TIME_TYPE_FORMATTER.format(localTime)); + } + }; + } + + private static SliceWriteFunction typedVarcharWriteFunction(String jdbcTypeName) + { + requireNonNull(jdbcTypeName, "jdbcTypeName is null"); + String quotedJdbcTypeName = jdbcTypeName.startsWith("\"") && jdbcTypeName.endsWith("\"") ? jdbcTypeName : "\"%s\"".formatted(jdbcTypeName.replace("\"", "\"\"")); + String bindExpression = format("CAST(? AS %s)", quotedJdbcTypeName); + + return new SliceWriteFunction() + { + @Override + public String getBindExpression() + { + return bindExpression; + } + + @Override + public void set(PreparedStatement statement, int index, Slice value) + throws SQLException + { + statement.setString(index, value.toStringUtf8()); + } + }; + } + + private ColumnMapping uuidColumnMapping() + { + return ColumnMapping.sliceMapping( + uuidType, + (resultSet, columnIndex) -> javaUuidToTrinoUuid((UUID) resultSet.getObject(columnIndex)), + uuidWriteFunction()); + } + + private ColumnMapping jsonColumnMapping() + { + return ColumnMapping.sliceMapping( + jsonType, + (resultSet, columnIndex) -> jsonParse(utf8Slice(resultSet.getString(columnIndex))), + typedVarcharWriteFunction("json"), + DISABLE_PUSHDOWN); + } + + private static AdbColumnMapping timestampWithTimeZoneColumnMapping(int precision) + { + // Adb supports timestamptz precision up to microseconds + checkArgument(precision <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION, "unsupported precision value %s", precision); + TimestampWithTimeZoneType trinoType = createTimestampWithTimeZoneType(precision); + if (precision <= TimestampWithTimeZoneType.MAX_SHORT_PRECISION) { + return new AdbColumnMapping(ColumnMapping.longMapping( + trinoType, + shortTimestampWithTimeZoneReadFunction(), + shortTimestampWithTimeZoneWriteFunction()), + new TimestampShortWithTimeZoneDataType(precision)); + } + return new AdbColumnMapping(ColumnMapping.objectMapping( + trinoType, + longTimestampWithTimeZoneReadFunction(), + longTimestampWithTimeZoneWriteFunction()), + new TimestampLongWithTimeZoneDataType(precision)); + } + + private static LongReadFunction shortTimestampWithTimeZoneReadFunction() + { + return (resultSet, columnIndex) -> { + // Adb does not store zone information in "timestamp with time zone" data type + long millisUtc = resultSet.getTimestamp(columnIndex).getTime(); + return packDateTimeWithZone(millisUtc, UTC_KEY); + }; + } + + private static LongWriteFunction shortTimestampWithTimeZoneWriteFunction() + { + return (statement, index, value) -> { + // Adb does not store zone information in "timestamp with time zone" data type + long millisUtc = unpackMillisUtc(value); + statement.setTimestamp(index, new Timestamp(millisUtc)); + }; + } + + private static ObjectReadFunction longTimestampWithTimeZoneReadFunction() + { + return ObjectReadFunction.of( + LongTimestampWithTimeZone.class, + (resultSet, columnIndex) -> { + // Adb does not store zone information in "timestamp with time zone" data type + OffsetDateTime offsetDateTime = resultSet.getObject(columnIndex, OffsetDateTime.class); + return LongTimestampWithTimeZone.fromEpochSecondsAndFraction( + offsetDateTime.toEpochSecond(), + (long) offsetDateTime.getNano() * PICOSECONDS_PER_NANOSECOND, + UTC_KEY); + }); + } + + private static ObjectWriteFunction longTimestampWithTimeZoneWriteFunction() + { + return ObjectWriteFunction.of( + LongTimestampWithTimeZone.class, + (statement, index, value) -> { + // Adb does not store zone information in "timestamp with time zone" data type + long epochSeconds = floorDiv(value.getEpochMillis(), MILLISECONDS_PER_SECOND); + long nanosOfSecond = (long) floorMod(value.getEpochMillis(), MILLISECONDS_PER_SECOND) * NANOSECONDS_PER_MILLISECOND + value.getPicosOfMilli() / PICOSECONDS_PER_NANOSECOND; + statement.setObject(index, OffsetDateTime.ofInstant(Instant.ofEpochSecond(epochSeconds, nanosOfSecond), UTC_KEY.getZoneId())); + }); + } + + private static SliceWriteFunction uuidWriteFunction() + { + return (statement, index, value) -> statement.setObject(index, trinoUuidToJavaUuid(value), Types.OTHER); + } + + private static ColumnMapping charColumnMapping(int charLength) + { + if (charLength > 65536) { + return varcharColumnMapping(charLength); + } + else { + CharType charType = CharType.createCharType(charLength); + return ColumnMapping.sliceMapping( + charType, charReadFunction(charType), charWriteFunction(), ADB_STRING_COLLATION_AWARE_PUSHDOWN); + } + } + + private static ColumnMapping varcharColumnMapping(int varcharLength) + { + VarcharType varcharType = varcharLength <= 2147483646 ? VarcharType.createVarcharType(varcharLength) : VarcharType.createUnboundedVarcharType(); + return ColumnMapping.sliceMapping( + varcharType, + varcharReadFunction(varcharType), + varcharWriteFunction(), + ADB_STRING_COLLATION_AWARE_PUSHDOWN); + } + + private static ColumnMapping enumColumnMapping(ConnectorSession session, String jdbcTypeName) + { + //todo implement AdbAdvancedPushdownSessionProperties.isPushdownEnums(session); + boolean pushdownEnums = false; + PredicatePushdownController pushdownController = pushdownEnums ? ADB_STRING_COLLATION_AWARE_PUSHDOWN : DISABLE_PUSHDOWN; + return ColumnMapping.sliceMapping( + VARCHAR, + (resultSet, columnIndex) -> utf8Slice(resultSet.getString(columnIndex)), + typedVarcharWriteFunction(jdbcTypeName), + pushdownController); + } + + private static ColumnMapping timeColumnMapping(int precision) + { + // adb limit but also assumption within this method + verify(precision <= 6, "Unsupported precision: %s", precision); + return ColumnMapping.longMapping( + createTimeType(precision), + (resultSet, columnIndex) -> { + LocalTime time = resultSet.getObject(columnIndex, LocalTime.class); + long nanosOfDay = time.toNanoOfDay(); + if (nanosOfDay == NANOSECONDS_PER_DAY - 1) { + // adb 24:00:00 is returned as 23:59:59.999999999, regardless of column precision + nanosOfDay = NANOSECONDS_PER_DAY - LongMath.pow(10, 9 - precision); + } + + long picosOfDay = nanosOfDay * PICOSECONDS_PER_NANOSECOND; + return round(picosOfDay, 12 - precision); + }, + timeWriteFunction(precision), + // Pushdown disabled because adb distinguishes TIME '24:00:00' and TIME '00:00:00' whereas Trino does not. + DISABLE_PUSHDOWN); + } + + private static void shortTimestampWriteFunction(PreparedStatement statement, int index, long epochMicros) + throws SQLException + { + LocalDateTime localDateTime = fromTrinoTimestamp(epochMicros); + statement.setObject(index, toPgTimestamp(localDateTime)); + } + + private static ObjectWriteFunction longTimestampWriteFunction() + { + return ObjectWriteFunction.of(LongTimestamp.class, (statement, index, timestamp) -> { + // adb supports up to 6 digits of precision + //noinspection ConstantConditions + verify(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION == 6); + long epochMicros = timestamp.getEpochMicros(); + if (timestamp.getPicosOfMicro() >= PICOSECONDS_PER_MICROSECOND / 2) { + epochMicros++; + } + shortTimestampWriteFunction(statement, index, epochMicros); + }); + } + + private Optional arrayToTrinoType(ConnectorSession session, Connection connection, JdbcTypeHandle typeHandle) + { + checkArgument(typeHandle.jdbcType() == Types.ARRAY, "Not array type"); + AdbPluginConfig.ArrayMapping arrayMapping = getArrayMapping(session); + if (arrayMapping == DISABLED) { + return Optional.empty(); + } + JdbcTypeHandle baseElementTypeHandle = getArrayElementTypeHandle(connection, typeHandle); + String baseElementTypeName = baseElementTypeHandle.jdbcTypeName() + .orElseThrow(() -> new TrinoException(JDBC_ERROR, "Element type name is missing: " + baseElementTypeHandle)); + if (baseElementTypeHandle.jdbcType() == Types.BINARY) { + // adb jdbc driver doesn't currently support array of varbinary (bytea[]) + return Optional.empty(); + } + Optional baseElementMapping = toColumnMapping(session, connection, baseElementTypeHandle); + + if (arrayMapping == AS_ARRAY) { + if (typeHandle.arrayDimensions().isEmpty()) { + return Optional.empty(); + } + return baseElementMapping + .map(elementMapping -> { + ArrayType trinoArrayType = new ArrayType(elementMapping.getType()); + ColumnMapping arrayColumnMapping = arrayColumnMapping(session, trinoArrayType, elementMapping, baseElementTypeName); + + int arrayDimensions = typeHandle.arrayDimensions().get(); + for (int i = 1; i < arrayDimensions; i++) { + trinoArrayType = new ArrayType(trinoArrayType); + arrayColumnMapping = arrayColumnMapping(session, trinoArrayType, arrayColumnMapping, baseElementTypeName); + } + return arrayColumnMapping; + }); + } + if (arrayMapping == AS_JSON) { + return baseElementMapping + .map(elementMapping -> arrayAsJsonColumnMapping(session, elementMapping)); + } + throw new IllegalStateException("Unsupported array mapping type: " + arrayMapping); + } + + private static ColumnMapping arrayColumnMapping(ConnectorSession session, ArrayType arrayType, ColumnMapping arrayElementMapping, String baseElementJdbcTypeName) + { + return ColumnMapping.objectMapping( + arrayType, + arrayReadFunction(arrayType.getElementType(), arrayElementMapping.getReadFunction()), + arrayWriteFunction(session, arrayType.getElementType(), baseElementJdbcTypeName)); + } + + private static ObjectWriteFunction arrayWriteFunction(ConnectorSession session, Type elementType, String baseElementJdbcTypeName) + { + return ObjectWriteFunction.of(Block.class, (statement, index, block) -> { + Array jdbcArray = statement.getConnection().createArrayOf(baseElementJdbcTypeName, getJdbcObjectArray(session, elementType, block)); + statement.setArray(index, jdbcArray); + }); + } + + private ColumnMapping arrayAsJsonColumnMapping(ConnectorSession session, ColumnMapping baseElementMapping) + { + return ColumnMapping.sliceMapping( + jsonType, + arrayAsJsonReadFunction(session, baseElementMapping), + (statement, index, block) -> { + throw new TrinoException(NOT_SUPPORTED, "Writing to array type is unsupported"); + }, + DISABLE_PUSHDOWN); + } + + private static JdbcTypeHandle getArrayElementTypeHandle(Connection connection, JdbcTypeHandle arrayTypeHandle) + { + String jdbcTypeName = arrayTypeHandle.jdbcTypeName() + .orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + arrayTypeHandle)); + try { + TypeInfo typeInfo = connection.unwrap(PgConnection.class).getTypeInfo(); + int pgElementOid = typeInfo.getPGArrayElement(typeInfo.getPGType(jdbcTypeName)); + verify(arrayTypeHandle.caseSensitivity().isEmpty(), "Case sensitivity not supported"); + return new JdbcTypeHandle( + typeInfo.getSQLType(pgElementOid), + Optional.of(typeInfo.getPGType(pgElementOid)), + arrayTypeHandle.columnSize(), + arrayTypeHandle.decimalDigits(), + Optional.empty(), + Optional.empty()); + } + catch (SQLException e) { + throw new TrinoException(JDBC_ERROR, e); + } + } + + private static SliceReadFunction arrayAsJsonReadFunction(ConnectorSession session, ColumnMapping baseElementMapping) + { + return (resultSet, columnIndex) -> { + Object jdbcArray = resultSet.getArray(columnIndex).getArray(); + int arrayDimensions = arrayDepth(jdbcArray); + + ReadFunction readFunction = baseElementMapping.getReadFunction(); + Type type = baseElementMapping.getType(); + for (int i = 0; i < arrayDimensions; i++) { + readFunction = arrayReadFunction(type, readFunction); + type = new ArrayType(type); + } + Block block = (Block) ((ObjectReadFunction) readFunction).readObject(resultSet, columnIndex); + BlockBuilder builder = type.createBlockBuilder(null, 1); + type.writeObject(builder, block); + Object value = type.getObjectValue(session, builder.build(), 0); + + try { + return toJsonValue(value); + } + catch (IOException e) { + throw new TrinoException(JDBC_ERROR, "Conversion to JSON failed for " + type.getDisplayName(), e); + } + }; + } + + private static ObjectReadFunction arrayReadFunction(Type elementType, ReadFunction elementReadFunction) + { + return ObjectReadFunction.of(Block.class, (resultSet, columnIndex) -> { + Array array = resultSet.getArray(columnIndex); + BlockBuilder builder = elementType.createBlockBuilder(null, 10); + try (ResultSet arrayAsResultSet = array.getResultSet()) { + while (arrayAsResultSet.next()) { + if (elementReadFunction.isNull(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)) { + builder.appendNull(); + } + else if (elementType.getJavaType() == boolean.class) { + elementType.writeBoolean(builder, ((BooleanReadFunction) elementReadFunction).readBoolean(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)); + } + else if (elementType.getJavaType() == long.class) { + elementType.writeLong(builder, ((LongReadFunction) elementReadFunction).readLong(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)); + } + else if (elementType.getJavaType() == double.class) { + elementType.writeDouble(builder, ((DoubleReadFunction) elementReadFunction).readDouble(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)); + } + else if (elementType.getJavaType() == Slice.class) { + elementType.writeSlice(builder, ((SliceReadFunction) elementReadFunction).readSlice(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)); + } + else { + elementType.writeObject(builder, ((ObjectReadFunction) elementReadFunction).readObject(arrayAsResultSet, ARRAY_RESULT_SET_VALUE_COLUMN)); + } + } + } + return builder.build(); + }); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/ColumnValue.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/ColumnValue.java new file mode 100644 index 000000000000..9fd5c439a85a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/ColumnValue.java @@ -0,0 +1,18 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.decode; + +public record ColumnValue(Object value, long estimatedSize) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoder.java new file mode 100644 index 000000000000..3cb1a52c87e0 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoder.java @@ -0,0 +1,24 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.decode; + +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; + +public interface RowDecoder +{ + ConnectorRow decode(T[] data); + + DataFormat getFormat(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoderFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoderFactory.java new file mode 100644 index 000000000000..a83e9f8253a6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/RowDecoderFactory.java @@ -0,0 +1,24 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.decode; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; + +public interface RowDecoderFactory +{ + RowDecoder create(ConnectorSession session, List columnDataTypes); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoder.java new file mode 100644 index 000000000000..cfdfaf6eafce --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoder.java @@ -0,0 +1,403 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.decode.csv; + +import com.google.common.math.LongMath; +import io.airlift.slice.SizeOf; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.plugin.adb.connector.datatype.CharDataType; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.datatype.ConnectorDataType; +import io.trino.plugin.adb.connector.datatype.DecimalLongDataType; +import io.trino.plugin.adb.connector.datatype.DecimalShortDataType; +import io.trino.plugin.adb.connector.datatype.TimeDataType; +import io.trino.plugin.adb.connector.datatype.TimestampWithoutTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.VarcharDataType; +import io.trino.plugin.adb.connector.decode.ColumnValue; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.GpfdistConnectorRow; +import io.trino.plugin.base.util.JsonTypeUtil; +import io.trino.plugin.jdbc.StandardColumnMappings; +import io.trino.spi.type.CharType; +import io.trino.spi.type.DateTimeEncoding; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Decimals; +import io.trino.spi.type.Int128; +import io.trino.spi.type.LongTimestampWithTimeZone; +import io.trino.spi.type.TimeZoneKey; +import io.trino.spi.type.TimestampType; +import io.trino.spi.type.Timestamps; +import io.trino.spi.type.Type; +import io.trino.spi.type.UuidType; +import io.trino.spi.type.VarcharType; +import org.postgresql.util.PGbytea; + +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.nio.charset.StandardCharsets; +import java.sql.SQLException; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.function.BiFunction; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.io.BaseEncoding.base16; +import static io.airlift.slice.SliceUtf8.countCodePoints; +import static io.trino.plugin.adb.TypeUtil.DATE_TYPE_FORMATTER; +import static io.trino.plugin.adb.TypeUtil.TIMESTAMP_TYPE_FORMATTER; +import static java.lang.String.format; + +public class CsvRowDecoder + implements RowDecoder +{ + private static final String DECODE_VALUE_ERROR_MSG_TEMPLATE = "Unexpected value %s for type %s"; + private static final Map> DECODE_FUNCTIONS_MAP = + Map.ofEntries( + Map.entry(ConnectorDataType.BOOLEAN, decodeBoolean()), + Map.entry(ConnectorDataType.MONEY, decodeMoney()), + Map.entry(ConnectorDataType.UUID, decodeUuid()), + Map.entry(ConnectorDataType.JSONB, decodeJsonb()), + Map.entry(ConnectorDataType.BIT, decodeBit()), + Map.entry(ConnectorDataType.BIGINT, decodeBigint()), + Map.entry(ConnectorDataType.BYTEA, decodeBytes()), + Map.entry(ConnectorDataType.CHAR, decodeChar()), + Map.entry(ConnectorDataType.VARCHAR, decodeVarchar()), + Map.entry(ConnectorDataType.DECIMAL_SHORT, decodeDecimalShort()), + Map.entry(ConnectorDataType.DECIMAL_LONG, decodeDecimalLong()), + Map.entry(ConnectorDataType.INTEGER, decodeInteger()), + Map.entry(ConnectorDataType.SMALLINT, decodeSmallint()), + Map.entry(ConnectorDataType.REAL, decodeReal()), + Map.entry(ConnectorDataType.DOUBLE_PRECISION, decodeDouble()), + Map.entry(ConnectorDataType.DATE, decodeDate()), + Map.entry(ConnectorDataType.TIME, decodeTime()), + Map.entry(ConnectorDataType.TIMESTAMP_SHORT_WITH_TIME_ZONE, decodeTimestampShortWithTimeZone()), + Map.entry(ConnectorDataType.TIMESTAMP_LONG_WITH_TIME_ZONE, decodeTimestampLongWithTimeZone()), + Map.entry(ConnectorDataType.TIMESTAMP_WITHOUT_TIME_ZONE, decodeTimestampWithoutTimeZone())); + + private final List columnDataTypes; + private final List> decodeFunctions; + + public CsvRowDecoder(List columnDataTypes) + { + this.columnDataTypes = columnDataTypes; + decodeFunctions = columnDataTypes.stream() + .map(columnDataType -> { + BiFunction decodeFunc = + DECODE_FUNCTIONS_MAP.get(columnDataType.getType()); + if (decodeFunc == null) { + throw new IllegalArgumentException(format("Unsupported column %s with type %s", + columnDataType.getName(), + columnDataType)); + } + else { + return decodeFunc; + } + }) + .toList(); + checkArgument(columnDataTypes.size() == decodeFunctions.size(), + "Column list size does not match decode function list size"); + } + + @Override + public ConnectorRow decode(String[] data) + { + try { + checkArgument(data.length == decodeFunctions.size(), + format("Parsed string values array length does not match decode function list size. Expected: %d, actual: %s", + decodeFunctions.size(), data.length)); + return decodeColumnsValues(data); + } + catch (Exception e) { + throw new IllegalArgumentException("Failed to decode data: " + e.getMessage(), e); + } + } + + private ConnectorRow decodeColumnsValues(String[] tokens) + { + ColumnValue[] values = new ColumnValue[tokens.length]; + long estimatedSize = 0L; + for (int i = 0; i < values.length; i++) { + ColumnValue columnValue = decodeValue(tokens[i], i); + values[i] = columnValue; + estimatedSize += columnValue.estimatedSize(); + } + return new GpfdistConnectorRow(estimatedSize, values); + } + + private ColumnValue decodeValue(String data, int position) + { + if (data != null) { + return decodeFunctions.get(position).apply(columnDataTypes.get(position), data); + } + else { + return new ColumnValue(null, 0); + } + } + + @Override + public DataFormat getFormat() + { + return DataFormat.CSV; + } + + private static BiFunction decodeBoolean() + { + return (_, data) -> new ColumnValue(getBooleanFromString(data), SizeOf.BOOLEAN_INSTANCE_SIZE); + } + + private static BiFunction decodeMoney() + { + return (_, data) -> { + Slice value = Slices.utf8Slice(data); + return new ColumnValue(value, value.getRetainedSize()); + }; + } + + private static BiFunction decodeBigint() + { + return (_, data) -> new ColumnValue(Long.parseLong(data), SizeOf.LONG_INSTANCE_SIZE); + } + + private static BiFunction decodeUuid() + { + return (_, data) -> { + Slice value = UuidType.javaUuidToTrinoUuid(UUID.fromString(data)); + return new ColumnValue(value, value.getRetainedSize()); + }; + } + + private static BiFunction decodeJsonb() + { + return (_, data) -> { + Slice value = JsonTypeUtil.jsonParse(Slices.utf8Slice(data)); + return new ColumnValue(value, value.getRetainedSize()); + }; + } + + private static BiFunction decodeBit() + { + return (_, data) -> new ColumnValue(getBooleanFromBitString(data), SizeOf.BOOLEAN_INSTANCE_SIZE); + } + + private static BiFunction decodeBytes() + { + return (_, data) -> { + try { + Slice value = Slices.wrappedBuffer(PGbytea.toBytes(data.getBytes(StandardCharsets.UTF_8))); + return new ColumnValue(value, value.getRetainedSize()); + } + catch (SQLException e) { + throw new IllegalArgumentException(format(DECODE_VALUE_ERROR_MSG_TEMPLATE, data, ConnectorDataType.BYTEA)); + } + }; + } + + private static BiFunction decodeChar() + { + return (dataType, data) -> { + CharType charType = ((CharDataType) dataType).getCharType(); + int trimCount = 0; + for (int i = data.length() - 1; i >= 0 && data.charAt(i) == ' '; i--) { + trimCount++; + } + if (trimCount > 0) { + data = data.substring(0, data.length() - trimCount); + } + Slice value = Slices.utf8Slice(data); + checkLengthInCodePoints(value, charType, charType.getLength()); + return new ColumnValue(value, value.getRetainedSize()); + }; + } + + private static BiFunction decodeVarchar() + { + return (dataType, data) -> { + VarcharType varcharType = ((VarcharDataType) dataType).getVarcharType(); + Slice value = Slices.utf8Slice(data); + if (!varcharType.isUnbounded()) { + checkLengthInCodePoints(value, varcharType, varcharType.getBoundedLength()); + } + return new ColumnValue(value, value.getRetainedSize()); + }; + } + + private static BiFunction decodeDecimalShort() + { + return (dataType, data) -> { + DecimalType decimalType = ((DecimalShortDataType) dataType).getDecimalType(); + return new ColumnValue(Decimals.encodeShortScaledValue(new BigDecimal(data), + decimalType.getScale(), + RoundingMode.UNNECESSARY), + SizeOf.LONG_INSTANCE_SIZE); + }; + } + + private static BiFunction decodeDecimalLong() + { + return (dataType, data) -> { + DecimalType decimalType = ((DecimalLongDataType) dataType).getDecimalType(); + return new ColumnValue(Decimals.valueOf(new BigDecimal(data).setScale(decimalType.getScale(), + RoundingMode.UNNECESSARY)), + Int128.INSTANCE_SIZE); + }; + } + + private static BiFunction decodeInteger() + { + return (_, data) -> { + long value = Integer.parseInt(data); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeSmallint() + { + return (_, data) -> { + long value = Short.parseShort(data); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeReal() + { + return (_, data) -> { + long value = Float.floatToRawIntBits(Float.parseFloat(data)); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeDouble() + { + return (_, data) -> { + double value = Double.parseDouble(data); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeDate() + { + return (_, data) -> new ColumnValue(LocalDate.parse(data, DATE_TYPE_FORMATTER).toEpochDay(), SizeOf.LONG_INSTANCE_SIZE); + } + + private static BiFunction decodeTime() + { + return (dataType, data) -> { + int precision = ((TimeDataType) dataType).getPrecision(); + LocalTime time; + if ("24:00:00".equals(data)) { + time = LocalTime.MAX; + } + else { + time = LocalTime.parse(data); + } + long nanosOfDay = time.toNanoOfDay(); + if (nanosOfDay == 86399999999999L) { + nanosOfDay = 86400000000000L - LongMath.pow(10L, 9 - precision); + } + long picosOfDay = nanosOfDay * 1000L; + long value = Timestamps.round(picosOfDay, 12 - precision); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeTimestampShortWithTimeZone() + { + return (_, data) -> { + LongTimestampWithTimeZone timestampValue = decodeTimestampWithTimeZoneLong(data); + long millisUtc = timestampValue.getEpochMillis(); + long value = DateTimeEncoding.packDateTimeWithZone(millisUtc, TimeZoneKey.getTimeZoneKey(timestampValue.getTimeZoneKey())); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static BiFunction decodeTimestampLongWithTimeZone() + { + return (_, data) -> new ColumnValue(decodeTimestampWithTimeZoneLong(data), LongTimestampWithTimeZone.INSTANCE_SIZE); + } + + private static BiFunction decodeTimestampWithoutTimeZone() + { + return (dataType, data) -> { + TimestampType timestampType = ((TimestampWithoutTimeZoneDataType) dataType).getTimestampType(); + LocalDateTime res = LocalDateTime.parse(data, TIMESTAMP_TYPE_FORMATTER); + long value = StandardColumnMappings.toTrinoTimestamp(timestampType, res); + return new ColumnValue(value, SizeOf.sizeOf(value)); + }; + } + + private static boolean getBooleanFromString(String data) + { + if ("t".equals(data)) { + return true; + } + else if ("f".equals(data)) { + return false; + } + else { + throw new IllegalArgumentException(format(DECODE_VALUE_ERROR_MSG_TEMPLATE, data, ConnectorDataType.BOOLEAN)); + } + } + + private static boolean getBooleanFromBitString(String data) + { + if ("1".equals(data) + || "true".equalsIgnoreCase(data) + || "t".equalsIgnoreCase(data) + || "yes".equalsIgnoreCase(data) + || "y".equalsIgnoreCase(data) + || "on".equalsIgnoreCase(data)) { + return true; + } + else if (!"0".equals(data) + && !"false".equalsIgnoreCase(data) + && !"f".equalsIgnoreCase(data) + && !"no".equalsIgnoreCase(data) + && !"n".equalsIgnoreCase(data) + && !"off".equalsIgnoreCase(data)) { + throw new IllegalArgumentException(format(DECODE_VALUE_ERROR_MSG_TEMPLATE, data, ConnectorDataType.BIT)); + } + else { + return false; + } + } + + private static void checkLengthInCodePoints(Slice value, Type characterDataType, int lengthLimit) + { + if (value.length() > lengthLimit) { + if (countCodePoints(value) > lengthLimit) { + throw new IllegalStateException( + String.format("Illegal value for trino type %s: '%s' [%s]", + characterDataType, + value.toStringUtf8(), + base16().encode(value.getBytes()))); + } + } + } + + private static LongTimestampWithTimeZone decodeTimestampWithTimeZoneLong(String data) + { + OffsetDateTime offsetDateTime = OffsetDateTime.parse(data, TIMESTAMP_TYPE_FORMATTER); + return LongTimestampWithTimeZone.fromEpochSecondsAndFraction(offsetDateTime.toEpochSecond(), + (long) offsetDateTime.getNano() * 1000L, + TimeZoneKey.UTC_KEY); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoderFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoderFactory.java new file mode 100644 index 000000000000..2c544da54268 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/decode/csv/CsvRowDecoderFactory.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.decode.csv; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.decode.RowDecoderFactory; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; + +public class CsvRowDecoderFactory + implements RowDecoderFactory +{ + @Override + public RowDecoder create(ConnectorSession session, List columnDataTypes) + { + return new CsvRowDecoder(columnDataTypes); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/AbstractRowEncoder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/AbstractRowEncoder.java new file mode 100644 index 000000000000..25d3d9cabb66 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/AbstractRowEncoder.java @@ -0,0 +1,301 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +import com.google.common.collect.ImmutableList; +import com.google.common.primitives.Shorts; +import io.trino.plugin.adb.connector.datatype.CharDataType; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.datatype.ConnectorDataType; +import io.trino.plugin.adb.connector.datatype.DecimalShortDataType; +import io.trino.plugin.adb.connector.datatype.MoneyDataType; +import io.trino.plugin.adb.connector.datatype.SmallintDataType; +import io.trino.plugin.adb.connector.datatype.TimeDataType; +import io.trino.plugin.adb.connector.datatype.TimestampLongWithTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.TimestampShortWithTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.TimestampWithoutTimeZoneDataType; +import io.trino.plugin.adb.connector.datatype.VarcharDataType; +import io.trino.spi.block.Block; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.type.BooleanType; +import io.trino.spi.type.DateTimeEncoding; +import io.trino.spi.type.DateType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.Int128; +import io.trino.spi.type.LongTimestampWithTimeZone; +import io.trino.spi.type.RealType; +import io.trino.spi.type.TimeType; +import io.trino.spi.type.TimeZoneKey; +import io.trino.spi.type.Timestamps; +import io.trino.spi.type.Type; +import io.trino.spi.type.UuidType; +import io.trino.spi.type.VarbinaryType; +import io.trino.spi.type.VarcharType; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.math.MathContext; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.jdbc.StandardColumnMappings.fromTrinoTimestamp; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.IntegerType.INTEGER; +import static java.lang.Float.intBitsToFloat; +import static java.lang.Math.toIntExact; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public abstract class AbstractRowEncoder + implements RowEncoder +{ + private static final String UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE = "Unsupported type '%s' for column '%s'"; + private static final long PICO_SECONDS_PER_DAY = 86400000000000000L; + protected final ConnectorSession session; + protected final List columnDataTypes; + private final Map> map; + protected int currentColumnIndex; + + protected AbstractRowEncoder(ConnectorSession session, List columnDataTypes) + { + requireNonNull(columnDataTypes, "columnDataTypes is null"); + this.session = session; + this.columnDataTypes = ImmutableList.copyOf(columnDataTypes); + this.currentColumnIndex = 0; + map = Map.ofEntries( + Map.entry(ConnectorDataType.BOOLEAN, (encoder, pageBlock) -> encoder.appendBoolean(BOOLEAN.getBoolean(pageBlock.block(), pageBlock.position()))), + Map.entry(ConnectorDataType.MONEY, (encoder, pageBlock) -> encoder.appendString(((MoneyDataType) pageBlock.columnDataType()).getVarcharType().getSlice(pageBlock.block(), pageBlock.position()).toStringUtf8())), + Map.entry(ConnectorDataType.UUID, (encoder, pageBlock) -> encoder.appendString(UuidType.trinoUuidToJavaUuid(UuidType.UUID.getSlice(pageBlock.block(), pageBlock.position())).toString())), + Map.entry(ConnectorDataType.JSONB, (encoder, pageBlock) -> encoder.appendString(VarcharType.VARCHAR.getSlice(pageBlock.block(), pageBlock.position()).toStringUtf8())), + Map.entry(ConnectorDataType.BIT, (encoder, pageBlock) -> encoder.appendString(BooleanType.BOOLEAN.getBoolean(pageBlock.block(), pageBlock.position()) ? "1" : "0")), + Map.entry(ConnectorDataType.BIGINT, (encoder, pageBlock) -> encoder.appendLong(BIGINT.getLong(pageBlock.block(), pageBlock.position()))), + Map.entry(ConnectorDataType.BYTEA, (encoder, pageBlock) -> encoder.appendBytes(VarbinaryType.VARBINARY.getSlice(pageBlock.block(), pageBlock.position()).getBytes())), + Map.entry(ConnectorDataType.CHAR, (encoder, pageBlock) -> encoder.appendString(((CharDataType) pageBlock.columnDataType()).getCharType().getSlice(pageBlock.block(), pageBlock.position()).toStringUtf8())), + Map.entry(ConnectorDataType.VARCHAR, (encoder, pageBlock) -> encoder.appendString(((VarcharDataType) pageBlock.columnDataType()).getVarcharType().getSlice(pageBlock.block(), pageBlock.position()).toStringUtf8())), + Map.entry(ConnectorDataType.DECIMAL_SHORT, AbstractRowEncoder::encodeDecimalShort), + Map.entry(ConnectorDataType.DECIMAL_LONG, AbstractRowEncoder::encodeDecimalLong), + Map.entry(ConnectorDataType.INTEGER, (encoder, pageBlock) -> encoder.appendInt(INTEGER.getInt(pageBlock.block(), pageBlock.position()))), + Map.entry(ConnectorDataType.SMALLINT, AbstractRowEncoder::encodeSmallint), + Map.entry(ConnectorDataType.REAL, (encoder, pageBlock) -> encoder.appendFloat(intBitsToFloat(toIntExact(RealType.REAL.getLong(pageBlock.block(), pageBlock.position()))))), + Map.entry(ConnectorDataType.DOUBLE_PRECISION, (encoder, pageBlock) -> encoder.appendDouble(DoubleType.DOUBLE.getDouble(pageBlock.block(), pageBlock.position()))), + Map.entry(ConnectorDataType.DATE, (encoder, pageBlock) -> encoder.appendDate(LocalDate.ofEpochDay(DateType.DATE.getLong(pageBlock.block(), pageBlock.position())))), + Map.entry(ConnectorDataType.TIME, AbstractRowEncoder::encodeTime), + Map.entry(ConnectorDataType.TIMESTAMP_SHORT_WITH_TIME_ZONE, AbstractRowEncoder::encodeTimestampShortWithTimeZone), + Map.entry(ConnectorDataType.TIMESTAMP_LONG_WITH_TIME_ZONE, AbstractRowEncoder::encodeTimestampLongWithTimeZone), + Map.entry(ConnectorDataType.TIMESTAMP_WITHOUT_TIME_ZONE, AbstractRowEncoder::encodeTimestampWithoutTimeZone)); + } + + @Override + public void appendColumnValue(Block block, int position) + { + checkArgument(currentColumnIndex < columnDataTypes.size(), + format("currentColumnIndex '%d' is greater than number of columns '%d'", + currentColumnIndex, + columnDataTypes.size())); + ColumnDataType columnDataType = columnDataTypes.get(currentColumnIndex); + ConnectorDataType type = columnDataType.getType(); + if (block.isNull(position)) { + appendNullValue(); + } + else { + BiConsumer func = map.get(type); + checkArgument(func != null, + format("Encoder for column type '%s' is not supported", type)); + func.accept(this, new EncoderMetadata(columnDataType, block, position)); + } + currentColumnIndex++; + } + + private static void encodeTimestampShortWithTimeZone(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + TimestampShortWithTimeZoneDataType dataType = (TimestampShortWithTimeZoneDataType) pageBlock.columnDataType(); + long dateTimeWithTimeZone = dataType.getTrinoType().getLong(pageBlock.block(), pageBlock.position()); + long millis = DateTimeEncoding.unpackMillisUtc(dateTimeWithTimeZone); + Instant instant = Instant.ofEpochMilli(millis); + encoder.appendDateTimeWithTimeZone(OffsetDateTime.ofInstant(instant, TimeZoneKey.UTC_KEY.getZoneId())); + } + + private static void encodeTimestampLongWithTimeZone(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + TimestampLongWithTimeZoneDataType dataType = (TimestampLongWithTimeZoneDataType) pageBlock.columnDataType(); + LongTimestampWithTimeZone value = (LongTimestampWithTimeZone) dataType.getTrinoType().getObject(pageBlock.block(), pageBlock.position()); + long epochSeconds = Math.floorDiv(value.getEpochMillis(), 1000); + long nanosOfSecond = (long) Math.floorMod(value.getEpochMillis(), 1000) * 1000000L + (long) (value.getPicosOfMilli() / 1000); + Instant instant = Instant.ofEpochSecond(epochSeconds, nanosOfSecond); + encoder.appendDateTimeWithTimeZone(OffsetDateTime.ofInstant(instant, TimeZoneKey.UTC_KEY.getZoneId())); + } + + private static void encodeTimestampWithoutTimeZone(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + TimestampWithoutTimeZoneDataType dataType = (TimestampWithoutTimeZoneDataType) pageBlock.columnDataType(); + long epochMicros = dataType.getTimestampType().getLong(pageBlock.block(), pageBlock.position()); + LocalDateTime localDateTime = fromTrinoTimestamp(epochMicros); + encoder.appendDateTimeWithoutTimeZone(localDateTime); + } + + private static void encodeDecimalShort(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + DecimalType decimalType = ((DecimalShortDataType) pageBlock.columnDataType()).getDecimalType(); + BigInteger unscaledValue = BigInteger.valueOf(decimalType.getLong(pageBlock.block(), pageBlock.position())); + BigDecimal value = new BigDecimal(unscaledValue, decimalType.getScale(), new MathContext(decimalType.getPrecision())); + encoder.appendBigDecimal(value); + } + + private static void encodeDecimalLong(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + DecimalType decimalType = ((DecimalShortDataType) pageBlock.columnDataType()).getDecimalType(); + BigInteger unscaledValue = ((Int128) decimalType.getObject(pageBlock.block(), pageBlock.position())).toBigInteger(); + BigDecimal value = new BigDecimal(unscaledValue, decimalType.getScale(), new MathContext(decimalType.getPrecision())); + encoder.appendBigDecimal(value); + } + + private static void encodeSmallint(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + Type trinoType = ((SmallintDataType) pageBlock.columnDataType()).getTrinoType(); + encoder.appendShort(Shorts.checkedCast(trinoType.getLong(pageBlock.block(), pageBlock.position()))); + } + + private static void encodeTime(AbstractRowEncoder encoder, EncoderMetadata pageBlock) + { + int precision = ((TimeDataType) pageBlock.columnDataType()).getPrecision(); + long picosOfDay = TimeType.createTimeType(precision).getLong(pageBlock.block(), pageBlock.position()); + picosOfDay = Timestamps.round(picosOfDay, 12 - precision); + if (picosOfDay == PICO_SECONDS_PER_DAY) { + picosOfDay = 0L; + } + LocalTime localTime = LocalTime.ofNanoOfDay(picosOfDay / 1000L); + encoder.appendTime(localTime); + } + + protected void appendNullValue() + { + throw new UnsupportedOperationException(format("Column '%s' does not support 'null' value", + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendBoolean(boolean value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + boolean.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendLong(long value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + long.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendInt(int value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + int.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendShort(short value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + short.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendFloat(float value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + float.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendDouble(double value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + double.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendString(String value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendDate(LocalDate value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendTime(LocalTime value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendDateTimeWithTimeZone(OffsetDateTime value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendDateTimeWithoutTimeZone(LocalDateTime value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendBytes(byte[] value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + byte.class.getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void appendBigDecimal(BigDecimal value) + { + throw new UnsupportedOperationException(format(UNSUPPORTED_TYPE_ERROR_MSG_TEMPLATE, + value.getClass().getName(), + columnDataTypes.get(currentColumnIndex).getName())); + } + + protected void resetColumnIndex() + { + currentColumnIndex = 0; + } + + @Override + public void close() + { + // do nothing + } + + private record EncoderMetadata(ColumnDataType columnDataType, Block block, int position) + { + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormat.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormat.java new file mode 100644 index 000000000000..e1ca94d1c151 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormat.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +public enum DataFormat +{ + CSV +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatConfig.java new file mode 100644 index 000000000000..52a29ab27fe6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatConfig.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +public interface DataFormatConfig +{ + DataFormat getDataFormat(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatModule.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatModule.java new file mode 100644 index 000000000000..482bb4a2f973 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/DataFormatModule.java @@ -0,0 +1,36 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +import com.google.inject.Binder; +import com.google.inject.Scopes; +import com.google.inject.TypeLiteral; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.adb.connector.decode.RowDecoderFactory; +import io.trino.plugin.adb.connector.decode.csv.CsvRowDecoderFactory; +import io.trino.plugin.adb.connector.encode.csv.CsvFormatConfig; +import io.trino.plugin.adb.connector.encode.csv.CsvRowEncoderFactory; + +public class DataFormatModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + binder.bind(DataFormatConfig.class).to(CsvFormatConfig.class).in(Scopes.SINGLETON); + binder.bind(RowEncoderFactory.class).to(CsvRowEncoderFactory.class).in(Scopes.SINGLETON); + binder.bind(new TypeLiteral>() {}).to(CsvRowDecoderFactory.class) + .in(Scopes.SINGLETON); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoder.java new file mode 100644 index 000000000000..c6208a748770 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoder.java @@ -0,0 +1,28 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +import io.trino.spi.block.Block; + +import java.io.Closeable; + +public interface RowEncoder + extends Closeable +{ + void appendColumnValue(Block block, int position); + + byte[] toByteArray(); + + DataFormat getFormat(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoderFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoderFactory.java new file mode 100644 index 000000000000..cfb769cebf50 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/RowEncoderFactory.java @@ -0,0 +1,24 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; + +public interface RowEncoderFactory +{ + RowEncoder create(ConnectorSession session, List columnDataTypes); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvFormatConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvFormatConfig.java new file mode 100644 index 000000000000..59a72ee2e6a2 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvFormatConfig.java @@ -0,0 +1,75 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode.csv; + +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.encode.DataFormatConfig; + +public class CsvFormatConfig + implements DataFormatConfig +{ + private char delimiter = '|'; + private String nullValue; + private String encoding = "UTF-8"; + + public CsvFormatConfig() + { + //default value, otherwise checkstyle plugin will raise error + nullValue = null; + } + + public static CsvFormatConfig create() + { + return new CsvFormatConfig(); + } + + public CsvFormatConfig delimiter(char delimiter) + { + this.delimiter = delimiter; + return this; + } + + public CsvFormatConfig nullValue(String nullValue) + { + this.nullValue = nullValue; + return this; + } + + public char getDelimiter() + { + return delimiter; + } + + public String getNullValue() + { + return nullValue; + } + + @Override + public DataFormat getDataFormat() + { + return DataFormat.CSV; + } + + public CsvFormatConfig encoding(String encoding) + { + this.encoding = encoding; + return this; + } + + public String getEncoding() + { + return encoding; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoder.java new file mode 100644 index 000000000000..046d4215b362 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoder.java @@ -0,0 +1,172 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode.csv; + +import com.opencsv.CSVWriterBuilder; +import com.opencsv.ICSVWriter; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.datatype.ConnectorDataType; +import io.trino.plugin.adb.connector.encode.AbstractRowEncoder; +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.spi.connector.ConnectorSession; +import org.postgresql.util.PGbytea; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStreamWriter; +import java.io.UncheckedIOException; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.util.List; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.adb.TypeUtil.DATE_TYPE_FORMATTER; +import static io.trino.plugin.adb.TypeUtil.TIMESTAMP_TYPE_FORMATTER; +import static io.trino.plugin.adb.TypeUtil.TIME_TYPE_FORMATTER; +import static java.lang.String.format; + +public class CsvRowEncoder + extends AbstractRowEncoder +{ + private final CsvFormatConfig encoderConfig; + private final String[] row; + + public CsvRowEncoder(ConnectorSession session, List columnDataTypes, CsvFormatConfig encoderConfig) + { + super(session, columnDataTypes); + for (ColumnDataType columnDataType : this.columnDataTypes) { + checkArgument(columnDataType.getType() != ConnectorDataType.UNKNOWN, + "Unexpected data type '%s' defined for column '%s'", + columnDataType, + columnDataType.getName()); + } + this.row = new String[this.columnDataTypes.size()]; + this.encoderConfig = encoderConfig; + } + + @Override + protected void appendNullValue() + { + row[currentColumnIndex] = encoderConfig.getNullValue(); + } + + @Override + protected void appendBoolean(boolean value) + { + row[currentColumnIndex] = Boolean.toString(value); + } + + @Override + protected void appendLong(long value) + { + row[currentColumnIndex] = Long.toString(value); + } + + @Override + protected void appendInt(int value) + { + row[currentColumnIndex] = Integer.toString(value); + } + + @Override + protected void appendShort(short value) + { + row[currentColumnIndex] = Short.toString(value); + } + + @Override + protected void appendFloat(float value) + { + row[currentColumnIndex] = Float.toString(value); + } + + @Override + protected void appendDouble(double value) + { + row[currentColumnIndex] = Double.toString(value); + } + + @Override + protected void appendDate(LocalDate value) + { + row[currentColumnIndex] = DATE_TYPE_FORMATTER.format(value); + } + + @Override + protected void appendTime(LocalTime value) + { + row[currentColumnIndex] = TIME_TYPE_FORMATTER.format(value); + } + + @Override + protected void appendDateTimeWithTimeZone(OffsetDateTime value) + { + row[currentColumnIndex] = TIMESTAMP_TYPE_FORMATTER.format(value); + } + + @Override + protected void appendDateTimeWithoutTimeZone(LocalDateTime value) + { + row[currentColumnIndex] = TIMESTAMP_TYPE_FORMATTER.format(value); + } + + @Override + protected void appendString(String value) + { + row[currentColumnIndex] = value; + } + + @Override + protected void appendBytes(byte[] value) + { + row[currentColumnIndex] = PGbytea.toPGString(value); + } + + @Override + protected void appendBigDecimal(BigDecimal value) + { + row[currentColumnIndex] = value.toString(); + } + + @Override + public byte[] toByteArray() + { + checkArgument(currentColumnIndex == columnDataTypes.size(), + format("Missing %d columns", columnDataTypes.size() - currentColumnIndex + 1)); + + try (ByteArrayOutputStream byteArrayOuts = new ByteArrayOutputStream(); + OutputStreamWriter outsWriter = new OutputStreamWriter(byteArrayOuts, StandardCharsets.UTF_8); + ICSVWriter writer = new CSVWriterBuilder(outsWriter) + .withSeparator(encoderConfig.getDelimiter()) + .build()) { + writer.writeNext(row); + writer.flush(); + resetColumnIndex(); + return byteArrayOuts.toByteArray(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + @Override + public DataFormat getFormat() + { + return DataFormat.CSV; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoderFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoderFactory.java new file mode 100644 index 000000000000..f469a2d57e6d --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/encode/csv/CsvRowEncoderFactory.java @@ -0,0 +1,48 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.encode.csv; + +import com.google.inject.Inject; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.encode.DataFormatConfig; +import io.trino.plugin.adb.connector.encode.RowEncoder; +import io.trino.plugin.adb.connector.encode.RowEncoderFactory; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; + +import static com.google.common.base.Preconditions.checkArgument; + +public class CsvRowEncoderFactory + implements RowEncoderFactory +{ + private final DataFormatConfig dataFormatConfig; + + @Inject + public CsvRowEncoderFactory(DataFormatConfig dataFormatConfig) + { + checkArgument(dataFormatConfig.getDataFormat() == DataFormat.CSV, + "Unexpected encoder config format '%s' defined for current encoder factory '%s'", + dataFormatConfig.getDataFormat(), + getClass().getSimpleName()); + this.dataFormatConfig = dataFormatConfig; + } + + @Override + public RowEncoder create(ConnectorSession session, List columnDataTypes) + { + return new CsvRowEncoder(session, columnDataTypes, (CsvFormatConfig) dataFormatConfig); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/AdbMetadataDao.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/AdbMetadataDao.java new file mode 100644 index 000000000000..4fa4d511bfc1 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/AdbMetadataDao.java @@ -0,0 +1,30 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.metadata; + +import io.trino.plugin.base.mapping.IdentifierMapping; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; + +import java.util.Map; + +public interface AdbMetadataDao +{ + int getSegmentCount(ConnectorSession session) + throws TrinoException; + + boolean isSegmentedTable(ConnectorSession session, String objectName); + + Map getTableProperties(ConnectorSession session, String objectName, IdentifierMapping identifierMapping); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/impl/AdbMetadataDaoImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/impl/AdbMetadataDaoImpl.java new file mode 100644 index 000000000000..f6278abfabaf --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/metadata/impl/AdbMetadataDaoImpl.java @@ -0,0 +1,199 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.metadata.impl; + +import com.google.inject.Inject; +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.metadata.AdbMetadataDao; +import io.trino.plugin.adb.connector.table.AdbTableDistributed; +import io.trino.plugin.adb.connector.table.AdbTableStorageCompressType; +import io.trino.plugin.adb.connector.table.AdbTableStorageOrientation; +import io.trino.plugin.base.mapping.IdentifierMapping; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcErrorCode; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public class AdbMetadataDaoImpl + implements AdbMetadataDao +{ + private static final Logger log = Logger.get(AdbMetadataDaoImpl.class); + public static final String DISTRIBUTED = "DISTRIBUTED"; + public static final String DISTRIBUTED_BY = "DISTRIBUTED BY"; + public static final String DISTRIBUTED_RANDOMLY = "DISTRIBUTED RANDOMLY"; + public static final String DISTRIBUTED_REPLICATED = "DISTRIBUTED REPLICATED"; + private final ConnectionFactory connectionFactory; + + @Inject + public AdbMetadataDaoImpl(ConnectionFactory connectionFactory) + { + this.connectionFactory = requireNonNull(connectionFactory, "connectionFactory is null"); + } + + @Override + public int getSegmentCount(ConnectorSession session) + { + try { + int segmentCount = 0; + try (Connection connection = this.connectionFactory.openConnection(session); + ResultSet rs = connection.createStatement() + .executeQuery("SELECT MAX(content) + 1 FROM pg_catalog.gp_segment_configuration")) { + if (!rs.next()) { + return segmentCount; + } + segmentCount = rs.getInt(1); + if (segmentCount < 0) { + return 0; + } + } + return segmentCount; + } + catch (SQLException e) { + throw new TrinoException(JdbcErrorCode.JDBC_ERROR, "Failed to get segment count", e); + } + } + + @Override + public boolean isSegmentedTable(ConnectorSession session, String objectName) + { + String sql = String.format("WITH oid AS (SELECT '%s'::regclass::oid oid)\n" + + "SELECT pg_catalog.pg_get_table_distributedby(oid) FROM oid", objectName); + try { + boolean isDistributed; + try (Connection connection = this.connectionFactory.openConnection(session); + ResultSet rs = connection.createStatement().executeQuery(sql)) { + if (!rs.next()) { + return false; + } + String distribution = rs.getString(1).trim(); + isDistributed = + distribution.startsWith("DISTRIBUTED BY") || distribution.equals("DISTRIBUTED RANDOMLY"); + } + return isDistributed; + } + catch (SQLException e) { + throw new TrinoException(JdbcErrorCode.JDBC_ERROR, + "Failed to determine whether the table contains the segment ID column.", e); + } + } + + @Override + public Map getTableProperties(ConnectorSession session, String objectName, + IdentifierMapping identifierMapping) + { + String sql = "WITH oid AS (SELECT '" + objectName + "'::regclass::oid oid)\n" + + "SELECT pg_catalog.pg_get_table_distributedby(oid) FROM oid\n" + + "UNION ALL\n" + + "SELECT UNNEST(reloptions) FROM pg_catalog.pg_class\n" + + "WHERE oid IN (SELECT oid FROM oid)"; + List rows = new ArrayList<>(); + try (Connection connection = this.connectionFactory.openConnection(session); + ResultSet rs = connection.createStatement().executeQuery(sql)) { + while (rs.next()) { + String row = rs.getString(1).trim(); + if (!row.isEmpty()) { + rows.add(rs.getString(1).trim()); + } + } + } + catch (SQLException e) { + throw new TrinoException(JdbcErrorCode.JDBC_ERROR, "Failed to get table properties: " + objectName, e); + } + + if (rows.isEmpty()) { + return Map.of(); + } + Map res = new HashMap<>(rows.size()); + + for (String property : rows) { + try { + this.parseProperty(property, res, identifierMapping); + } + catch (Exception e) { + log.warn(e, "Failed to parse table property: " + property); + } + } + + return res; + } + + private void parseProperty(String property, Map res, IdentifierMapping identifierMapping) + { + if (property.startsWith(DISTRIBUTED)) { + if (property.startsWith(DISTRIBUTED_BY)) { + property = property.substring(DISTRIBUTED_BY.length()).trim(); + property = property.substring(1, property.length() - 1); + String[] remoteColumns = property.split(", "); + List columns = new ArrayList<>(remoteColumns.length); + + for (String remoteColumn : remoteColumns) { + columns.add(identifierMapping.fromRemoteColumnName(remoteColumn)); + } + + res.put("distributed_by", columns); + } + else if (property.equals(DISTRIBUTED_RANDOMLY)) { + res.put("distributed", AdbTableDistributed.RANDOMLY); + } + else if (property.equals(DISTRIBUTED_REPLICATED)) { + res.put("distributed", AdbTableDistributed.REPLICATED); + } + } + else { + String[] keyValue = property.split("="); + if (keyValue.length != 2) { + throw new IllegalArgumentException("Unsupported table property: " + property); + } + + String key = keyValue[0].trim(); + String value = keyValue[1].trim(); + switch (key) { + case "appendonly": + res.put("appendoptimized", Boolean.parseBoolean(value)); + break; + case "blocksize": + res.put("blocksize", Integer.parseInt(value)); + break; + case "orientation": + res.put("orientation", AdbTableStorageOrientation.valueOf(value.toUpperCase(Locale.ENGLISH))); + break; + case "checksum": + res.put("checksum", Boolean.parseBoolean(value)); + break; + case "compresstype": + res.put("compresstype", AdbTableStorageCompressType.valueOf(value.toUpperCase(Locale.ENGLISH))); + break; + case "compresslevel": + res.put("compresslevel", Integer.parseInt(value)); + break; + case "fillfactor": + res.put("fillfactor", Integer.parseInt(value)); + break; + default: + throw new IllegalArgumentException("Unsupported table property: " + property); + } + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/TransferDataProtocol.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/TransferDataProtocol.java new file mode 100644 index 000000000000..df63298b09a6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/TransferDataProtocol.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol; + +public enum TransferDataProtocol +{ + GPFDIST +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractContextManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractContextManager.java new file mode 100644 index 000000000000..3ba9b97b53e5 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractContextManager.java @@ -0,0 +1,48 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; + +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; + +public abstract class AbstractContextManager + implements ContextManager +{ + private static final Logger log = Logger.get(AbstractContextManager.class); + private final Map contextMap = new ConcurrentHashMap<>(); + + @Override + public Optional get(ContextId contextId) + { + return Optional.ofNullable(contextMap.get(contextId)); + } + + @Override + public void add(T context) + { + contextMap.put(context.getId(), context); + log.debug("Added context %s", context.getId()); + } + + @Override + public void remove(ContextId contextId) + { + Optional.ofNullable(contextMap.remove(contextId)) + .ifPresent(ctx -> log.info("Removed context %s", ctx.getId())); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractDataTransferQueryExecutor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractDataTransferQueryExecutor.java new file mode 100644 index 000000000000..45726673ff72 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractDataTransferQueryExecutor.java @@ -0,0 +1,82 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; + +public abstract class AbstractDataTransferQueryExecutor + implements DataTransferQueryExecutor +{ + private static final Logger log = Logger.get(AbstractDataTransferQueryExecutor.class); + protected final AdbSqlClient client; + protected final ConnectorSession session; + protected final ExecutorService executor; + protected final CreateExternalTableQueryFactory externalTableQueryFactory; + protected final InsertDataQueryFactory insertDataQueryFactory; + protected Connection connection; + + public AbstractDataTransferQueryExecutor(AdbSqlClient client, + ConnectorSession session, + ExecutorService executor, + CreateExternalTableQueryFactory externalTableQueryFactory, + InsertDataQueryFactory insertDataQueryFactory) + { + this.client = client; + this.session = session; + this.executor = executor; + this.externalTableQueryFactory = externalTableQueryFactory; + this.insertDataQueryFactory = insertDataQueryFactory; + } + + @Override + public CompletableFuture execute() + { + return CompletableFuture.supplyAsync(() -> { + try { + connection = client.getConnection(session); + connection.setAutoCommit(false); + connection.setReadOnly(false); + executeQueries(); + connection.commit(); + connection.close(); + return null; + } + catch (Exception e) { + String errMsg = "Failed to execute data transfer query: " + e.getMessage(); + log.error(errMsg, e); + throw new RuntimeException(errMsg, e); + } + finally { + if (connection != null) { + try { + connection.close(); + } + catch (SQLException e) { + log.warn("Failed to close connection", e); + } + } + } + }, executor); + } + + protected abstract void executeQueries() + throws SQLException; +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractExternalTableQueryFactory.java new file mode 100644 index 000000000000..8ac40af9df07 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/AbstractExternalTableQueryFactory.java @@ -0,0 +1,47 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; + +import java.util.stream.IntStream; + +import static java.lang.String.format; +import static java.util.stream.Collectors.joining; + +public abstract class AbstractExternalTableQueryFactory + implements CreateExternalTableQueryFactory +{ + protected String createCommonQuery(GpfdistMetadata metadata) + { + String columnsDefinition = IntStream.range(0, metadata.getColumnNames().size()) + .boxed() + .map(i -> { + String columnName = metadata.getColumnNames().get(i); + String typeName = metadata.getDataTypes().get(i).getName(); + return columnName + " " + typeName; + }) + .collect(joining(",")); + return format( + "CREATE %s EXTERNAL TEMPORARY TABLE %s (%s) LOCATION ('%s') FORMAT '%s' (DELIMITER '%s' NULL AS '%s') ENCODING '%s'", + getExternalTableType().name(), + metadata.getSourceTable(), + columnsDefinition, + metadata.getGpfdistLocation(), + metadata.getExternalTableFormatConfig().dataFormat().name(), + metadata.getExternalTableFormatConfig().delimiter(), + metadata.getExternalTableFormatConfig().nullValue() == null ? "" : metadata.getExternalTableFormatConfig().nullValue(), + metadata.getExternalTableFormatConfig().encoding()); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ConnectorRow.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ConnectorRow.java new file mode 100644 index 000000000000..5915b921ae6a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ConnectorRow.java @@ -0,0 +1,25 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.decode.ColumnValue; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.WithEstimatedSize; + +public interface ConnectorRow + extends WithEstimatedSize +{ + ColumnValue[] getColumnValues(); + + long getEstimatedSize(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/Context.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/Context.java new file mode 100644 index 000000000000..5a39a57d04cd --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/Context.java @@ -0,0 +1,23 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; + +public interface Context +{ + ContextId getId(); + + void close(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ContextManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ContextManager.java new file mode 100644 index 000000000000..ca5e8d572852 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ContextManager.java @@ -0,0 +1,27 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; + +import java.util.Optional; + +public interface ContextManager +{ + Optional get(ContextId contextId); + + void add(T context); + + void remove(ContextId contextId); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/CreateExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/CreateExternalTableQueryFactory.java new file mode 100644 index 000000000000..027128d748f8 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/CreateExternalTableQueryFactory.java @@ -0,0 +1,24 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; + +public interface CreateExternalTableQueryFactory +{ + String createQuery(GpfdistMetadata metadata); + + ExternalTableType getExternalTableType(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/DataTransferQueryExecutor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/DataTransferQueryExecutor.java new file mode 100644 index 000000000000..3d2508c86f9b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/DataTransferQueryExecutor.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import java.util.concurrent.CompletableFuture; + +public interface DataTransferQueryExecutor +{ + CompletableFuture execute(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ExecutorServiceProvider.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ExecutorServiceProvider.java new file mode 100644 index 000000000000..43de53b280ed --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/ExecutorServiceProvider.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.airlift.concurrent.Threads; + +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +public interface ExecutorServiceProvider +{ + ExecutorService WRITE_DATA_EXECUTOR_SERVICE = + Executors.newCachedThreadPool(Threads.daemonThreadsNamed("adb-write-data-%s")); + ExecutorService LOAD_DATA_QUERY_EXECUTOR_SERVICE = + Executors.newCachedThreadPool(Threads.daemonThreadsNamed("adb-load-query-%s")); + ExecutorService GPFDIST_HTTP_REQUEST_EXECUTOR_SERVICE = + Executors.newCachedThreadPool(Threads.daemonThreadsNamed("adb-gpfdist-http-response-%s")); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistModule.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistModule.java new file mode 100644 index 000000000000..7922b222fa3f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistModule.java @@ -0,0 +1,141 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import com.google.inject.Binder; +import com.google.inject.Provides; +import com.google.inject.Scopes; +import com.google.inject.Singleton; +import com.google.inject.TypeLiteral; +import com.google.inject.multibindings.MapBinder; +import com.google.inject.multibindings.Multibinder; +import com.google.inject.multibindings.OptionalBinder; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.http.server.HttpServer; +import io.airlift.http.server.HttpServerInfo; +import io.airlift.http.server.HttpServerProvider; +import io.airlift.http.server.RequestStats; +import io.airlift.node.NodeConfig; +import io.airlift.node.NodeInfo; +import io.opentelemetry.api.OpenTelemetry; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.process.GpfdistPageSinkProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.query.CreateReadableExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.query.InsertDataFromExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableFormatConfigFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableFormatConfigFactoryImpl; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadataFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadataFactoryImpl; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLocationFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLocationFactoryImpl; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadataFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadataFactoryImpl; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.GpfdistServerConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.GpfdistServerModule; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessorFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistInputDataProcessorFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistRecordSetProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.query.CreateWritableExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.query.InsertDataToExternalTableQueryFactory; +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.credential.CredentialProvider; +import io.trino.spi.NodeManager; +import io.trino.spi.connector.ConnectorPageSinkProvider; +import io.trino.spi.connector.ConnectorRecordSetProvider; +import org.postgresql.Driver; +import org.weakref.jmx.guice.ExportBinder; + +import java.util.Properties; + +import static com.google.inject.multibindings.MapBinder.newMapBinder; + +public class GpfdistModule + extends AbstractConfigurationAwareModule +{ + @Override + public void setup(Binder binder) + { + install(new GpfdistServerModule()); + + Multibinder createExtTableQueryFactories = Multibinder.newSetBinder(binder, CreateExternalTableQueryFactory.class); + Multibinder insertDataQueryFactories = Multibinder.newSetBinder(binder, InsertDataQueryFactory.class); + + createExtTableQueryFactories.addBinding().to(CreateReadableExternalTableQueryFactory.class).in(Scopes.SINGLETON); + createExtTableQueryFactories.addBinding().to(CreateWritableExternalTableQueryFactory.class).in(Scopes.SINGLETON); + insertDataQueryFactories.addBinding().to(InsertDataFromExternalTableQueryFactory.class).in(Scopes.SINGLETON); + insertDataQueryFactories.addBinding().to(InsertDataToExternalTableQueryFactory.class).in(Scopes.SINGLETON); + + binder.bind(ExternalTableFormatConfigFactory.class).to(ExternalTableFormatConfigFactoryImpl.class).in(Scopes.SINGLETON); + + OptionalBinder.newOptionalBinder(binder, ConnectorPageSinkProvider.class).setBinding().to(GpfdistPageSinkProvider.class).in(Scopes.SINGLETON); + OptionalBinder.newOptionalBinder(binder, ConnectorRecordSetProvider.class).setBinding().to(GpfdistRecordSetProvider.class).in(Scopes.SINGLETON); + + binder.bind(GpfdistLoadMetadataFactory.class).to(GpfdistLoadMetadataFactoryImpl.class).in(Scopes.SINGLETON); + binder.bind(GpfdistUnloadMetadataFactory.class).to(GpfdistUnloadMetadataFactoryImpl.class).in(Scopes.SINGLETON); + binder.bind(GpfdistLocationFactory.class).to(GpfdistLocationFactoryImpl.class).in(Scopes.SINGLETON); + binder.bind(GpfdistUnloadMetadataFactory.class).to(GpfdistUnloadMetadataFactoryImpl.class).in(Scopes.SINGLETON); + binder.bind(new TypeLiteral>() {}).to(ReadContextManager.class).in(Scopes.SINGLETON); + binder.bind(new TypeLiteral>() {}).to(WriteContextManager.class).in(Scopes.SINGLETON); + binder.bind(new TypeLiteral>() {}).to(GpfdistInputDataProcessorFactory.class).in(Scopes.SINGLETON); + binder.bind(NodeInfo.class).in(Scopes.SINGLETON); + binder.bind(HttpServerInfo.class).in(Scopes.SINGLETON); + binder.bind(RequestStats.class).in(Scopes.SINGLETON); + + ExportBinder.newExporter(binder) + .export(RequestStats.class) + .as(generator -> generator.generatedNameOf(RequestStats.class, "adb-gpfdist-server-request-stats")); + ExportBinder.newExporter(binder).export(HttpServer.class).as(generator -> generator.generatedNameOf(HttpServer.class, "adb-gpfdist-server")); + binder.bind(HttpServer.class).toProvider(HttpServerProvider.class).in(Scopes.SINGLETON); + } + + public static MapBinder externalTableQueriesFactoryMap(Binder binder) + { + return newMapBinder(binder, ExternalTableType.class, CreateExternalTableQueryFactory.class); + } + + @Provides + @Singleton + @ForBaseJdbc + public static ConnectionFactory getConnectionFactory(BaseJdbcConfig config, + OpenTelemetry openTelemetry, + CredentialProvider credentialProvider) + { + //todo implement cached connection pool with using hikari data source + Properties connectionProperties = new Properties(); + return DriverConnectionFactory.builder(new Driver(), config.getConnectionUrl(), credentialProvider) + .setConnectionProperties(connectionProperties) + .setOpenTelemetry(openTelemetry) + .build(); + } + + @Provides + @Singleton + public static NodeConfig getNodeConfig(GpfdistServerConfig config, NodeManager nodeManager) + { + String internalHost = config.getServerHost() != null ? config.getServerHost() : nodeManager.getCurrentNode().getHost(); + String externalHost = config.getServerExternalHost() != null ? config.getServerExternalHost() : internalHost; + return new NodeConfig() + .setEnvironment("adb_gpfdist") + .setNodeId(nodeManager.getCurrentNode().getNodeIdentifier()) + .setNodeInternalAddress(internalHost) + .setNodeExternalAddress(externalHost); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistUtil.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistUtil.java new file mode 100644 index 000000000000..6f5644df0b43 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/GpfdistUtil.java @@ -0,0 +1,26 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +public final class GpfdistUtil +{ + private GpfdistUtil() + { + } + + public static String createGpfdistFileName(String externalTableName) + { + return "/adb/" + externalTableName; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/InsertDataQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/InsertDataQueryFactory.java new file mode 100644 index 000000000000..0ed5d45bb8ca --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/InsertDataQueryFactory.java @@ -0,0 +1,28 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist; + +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.Connection; + +public interface InsertDataQueryFactory +{ + PreparedQuery create(ConnectorSession session, Connection connection, GpfdistMetadata gpfdistMetadata); + + ExternalTableType getExternalTableType(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/GpfdistPacketBuilder.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/GpfdistPacketBuilder.java new file mode 100644 index 000000000000..c00a253d2104 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/GpfdistPacketBuilder.java @@ -0,0 +1,99 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load; + +import com.google.common.primitives.Bytes; +import io.airlift.log.Logger; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; + +/** + * Gpfdist v1 packet structure: + * byte 0: type (can be 'F'ilename, 'O'ffset, 'D'ata, 'E'rror, 'L'inenumber) + * byte 1-4: length. # bytes of following data block. in network-order. + * byte 5-X: the block itself + * Data packets: + * F + fileName length + fileName + O + offset length + offset + L + lineNumber length + lineNumber + D + data length + data + * End packet: + * D + zero length + * Error packet: + * E + error data length + error data + */ +public class GpfdistPacketBuilder +{ + private static final Logger log = Logger.get(GpfdistPacketBuilder.class); + private static final byte GPFDIST_PACKAGE_FILE_NAME_MSG_TYPE = (byte) 70; + private static final byte GPFDIST_PACKAGE_OFFSET_MSG_TYPE = (byte) 79; + private static final byte GPFDIST_PACKAGE_LINE_NUMBER_MSG_TYPE = (byte) 76; + private static final byte GPFDIST_PACKAGE_END_MSG_TYPE = (byte) 68; + private static final byte GPFDIST_PACKAGE_ERROR_MSG_TYPE = (byte) 69; + private static final int HEADER_METADATA_TYPE_BYTES_LENGTH = 5; + private static final int OFFSET_BYTES_LENGTH = 8; + private static final int LINE_NUMBER_BYTES_LENGTH = 8; + private final String gpfdistFileName; + private final ByteBuffer headerBuffer; + private final ByteBuffer endPacketBuffer; + + public GpfdistPacketBuilder(String gpfdistFileName) + { + this.gpfdistFileName = gpfdistFileName; + headerBuffer = ByteBuffer.allocate(HEADER_METADATA_TYPE_BYTES_LENGTH * 4 + + OFFSET_BYTES_LENGTH + + LINE_NUMBER_BYTES_LENGTH + + gpfdistFileName.getBytes(StandardCharsets.UTF_8).length); + endPacketBuffer = ByteBuffer.allocate(HEADER_METADATA_TYPE_BYTES_LENGTH); + } + + public byte[] createDataPacket(PageSerializationResult serializationResult) + { + headerBuffer.clear(); + byte[] fileNameBytes = gpfdistFileName.getBytes(StandardCharsets.UTF_8); + headerBuffer.put(GPFDIST_PACKAGE_FILE_NAME_MSG_TYPE); + headerBuffer.putInt(fileNameBytes.length); + headerBuffer.put(fileNameBytes); + headerBuffer.put(GPFDIST_PACKAGE_OFFSET_MSG_TYPE); + headerBuffer.putInt(OFFSET_BYTES_LENGTH); + headerBuffer.putLong(0); + headerBuffer.put(GPFDIST_PACKAGE_LINE_NUMBER_MSG_TYPE); + headerBuffer.putInt(LINE_NUMBER_BYTES_LENGTH); + headerBuffer.putLong(serializationResult.rowCount()); + headerBuffer.put(GPFDIST_PACKAGE_END_MSG_TYPE); + headerBuffer.putInt(serializationResult.data().length); + return Bytes.concat(headerBuffer.array(), serializationResult.data()); + } + + public byte[] createSingleEmptyDataPacket() + { + return createDataPacket(new PageSerializationResult(new byte[] {}, 0)); + } + + public byte[] createEndPacket() + { + endPacketBuffer.clear(); + endPacketBuffer.put(GPFDIST_PACKAGE_END_MSG_TYPE); + endPacketBuffer.putInt(0); + return endPacketBuffer.array(); + } + + public byte[] createErrorPacket(Throwable error) + { + byte[] msgBytes = error.getMessage().getBytes(StandardCharsets.UTF_8); + ByteBuffer buffer = ByteBuffer.allocate(HEADER_METADATA_TYPE_BYTES_LENGTH + msgBytes.length); + buffer.put(GPFDIST_PACKAGE_ERROR_MSG_TYPE); + buffer.putInt(msgBytes.length); + buffer.put(msgBytes); + return buffer.array(); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessor.java new file mode 100644 index 000000000000..08f0dc44b54c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessor.java @@ -0,0 +1,23 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load; + +import io.trino.spi.Page; + +public interface PageProcessor +{ + void process(Page page); + + void stop(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessorProvider.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessorProvider.java new file mode 100644 index 000000000000..b0b7f93559a0 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageProcessorProvider.java @@ -0,0 +1,27 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load; + +import java.util.Queue; + +public interface PageProcessorProvider +{ + void add(PageProcessor processor); + + PageProcessor take(); + + Queue getAll(); + + void clear(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializationResult.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializationResult.java new file mode 100644 index 000000000000..c0b307d8f9f7 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializationResult.java @@ -0,0 +1,18 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load; + +public record PageSerializationResult(byte[] data, long rowCount) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializer.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializer.java new file mode 100644 index 000000000000..050e0e71bde4 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/PageSerializer.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load; + +import io.trino.spi.Page; + +public interface PageSerializer +{ + PageSerializationResult serialize(Page page); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContext.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContext.java new file mode 100644 index 000000000000..cb8bb8ed4d86 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContext.java @@ -0,0 +1,129 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.context; + +import io.airlift.log.Logger; +import io.airlift.units.DataSize; +import io.trino.plugin.adb.connector.encode.RowEncoder; +import io.trino.plugin.adb.connector.protocol.gpfdist.Context; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.process.GpfdistPageProcessorProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadata; + +import java.util.Queue; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static java.lang.String.format; + +public class WriteContext + implements Context +{ + private static final Logger log = Logger.get(WriteContext.class); + private final AtomicReference adbQueryException = new AtomicReference<>(); + private final AtomicLong completedBytes = new AtomicLong(); + private final AtomicLong memoryUsage = new AtomicLong(); + private final AtomicReference error = new AtomicReference<>(); + private final ContextId id; + private final GpfdistLoadMetadata metadata; + private final RowEncoder rowEncoder; + private final DataSize writeBufferSize; + private final GpfdistPageProcessorProvider pageProcessorProvider; + + public WriteContext(GpfdistLoadMetadata metadata, RowEncoder rowEncoder, DataSize writeBufferSize, + GpfdistPageProcessorProvider pageProcessorProvider) + { + this(new ContextId(metadata.getSourceTable()), metadata, rowEncoder, writeBufferSize, pageProcessorProvider); + } + + public WriteContext(ContextId id, + GpfdistLoadMetadata metadata, + RowEncoder rowEncoder, + DataSize writeBufferSize, + GpfdistPageProcessorProvider pageProcessorProvider) + { + this.id = id; + this.metadata = metadata; + this.rowEncoder = rowEncoder; + this.writeBufferSize = writeBufferSize; + this.pageProcessorProvider = pageProcessorProvider; + } + + @Override + public ContextId getId() + { + return id; + } + + public GpfdistLoadMetadata getMetadata() + { + return metadata; + } + + public RowEncoder getRowEncoder() + { + return rowEncoder; + } + + public AtomicReference getAdbQueryException() + { + return adbQueryException; + } + + public AtomicLong getCompletedBytes() + { + return completedBytes; + } + + public AtomicLong getMemoryUsage() + { + return memoryUsage; + } + + public DataSize getWriteBufferSize() + { + return writeBufferSize; + } + + public AtomicReference getError() + { + return error; + } + + public GpfdistPageProcessorProvider getPageProcessorProvider() + { + return pageProcessorProvider; + } + + @Override + public void close() + { + Queue pageProcessors = pageProcessorProvider.getAll(); + StringBuilder sb = new StringBuilder(); + pageProcessors.forEach(processor -> { + try { + processor.stop(); + } + catch (Exception e) { + sb.append(format("Failed to stop page processor %s. Error: %s;", processor, e.getMessage())); + } + }); + pageProcessorProvider.clear(); + if (!sb.isEmpty()) { + throw new RuntimeException(sb.toString()); + } + log.debug("Closed write context %s", id); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContextManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContextManager.java new file mode 100644 index 000000000000..3156fa2f6924 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/context/WriteContextManager.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.context; + +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractContextManager; + +public class WriteContextManager + extends AbstractContextManager +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessor.java new file mode 100644 index 000000000000..4525c316a306 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessor.java @@ -0,0 +1,130 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.process; + +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.encode.RowEncoder; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.GpfdistPacketBuilder; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageSerializationResult; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageSerializer; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistReadableRequest; +import io.trino.spi.Page; +import io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; + +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; + +import static io.trino.plugin.adb.connector.protocol.gpfdist.GpfdistUtil.createGpfdistFileName; + +public class GpfdistPageProcessor + implements PageProcessor +{ + private static final Logger log = Logger.get(GpfdistPageProcessor.class); + private final OutputStream outputStream; + private final GpfdistReadableRequest request; + private final WriteContext writeContext; + private final RowEncoder rowEncoder; + private long bytesWritten; + private final PageSerializer pageSerializer; + private final ByteBuffer pageLengthBuffer; + private final GpfdistPacketBuilder packetBuilder; + + public GpfdistPageProcessor(OutputStream outputStream, + GpfdistReadableRequest request, + WriteContext writeContext, + PageSerializer pageSerializer) + { + this.outputStream = outputStream; + this.writeContext = writeContext; + this.rowEncoder = writeContext.getRowEncoder(); + this.request = request; + this.pageSerializer = pageSerializer; + pageLengthBuffer = ByteBuffer.allocate(4); + packetBuilder = new GpfdistPacketBuilder(createGpfdistFileName(writeContext.getMetadata().getSourceTable())); + } + + @Override + public void process(Page page) + { + try { + writeContext.getMemoryUsage().addAndGet(page.getRetainedSizeInBytes()); + PageSerializationResult serializationResult = pageSerializer.serialize(page); + outputStream.write(packetBuilder.createDataPacket(serializationResult)); + bytesWritten += page.getRetainedSizeInBytes(); + pageLengthBuffer.clear(); + if (bytesWritten >= writeContext.getWriteBufferSize().toBytes()) { + outputStream.flush(); + log.debug("Sent data packet with size %d for request %s", bytesWritten, request); + writeContext.getCompletedBytes().addAndGet(bytesWritten); + writeContext.getMemoryUsage().addAndGet(-bytesWritten); + bytesWritten = 0; + } + } + catch (Exception e) { + writePacket(packetBuilder.createErrorPacket(e)); + close(); + throw new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, "Failed to process page " + page, e); + } + } + + @Override + public void stop() + { + Throwable error = writeContext.getError().get(); + if (error != null) { + writePacket(packetBuilder.createErrorPacket(error)); + log.warn("Stopped writing process with error: %s. Sent error packet data", error); + } + else { + writePacket(packetBuilder.createEndPacket()); + log.info("Stopped writing process for requestId %s. Sent finished packet data", request); + } + close(); + } + + private void writePacket(byte[] dataPacket) + { + try { + outputStream.write(dataPacket); + outputStream.flush(); + } + catch (IOException e) { + throw new RuntimeException("Failed to write end packet data: " + e.getMessage(), e); + } + } + + private void close() + { + try { + outputStream.close(); + rowEncoder.close(); + } + catch (IOException e) { + throw new RuntimeException("Failed to close processing data stream" + e.getMessage(), e); + } + } + + @Override + public String toString() + { + return "GpfdistPageProcessor{" + + "request=" + request + + ", writeContext.id=" + writeContext.getId() + + '}'; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessorProvider.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessorProvider.java new file mode 100644 index 000000000000..1cecdcdd6443 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageProcessorProvider.java @@ -0,0 +1,107 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.process; + +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageProcessorProvider; + +import java.util.LinkedList; +import java.util.Queue; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; + +import static java.lang.String.format; + +public class GpfdistPageProcessorProvider + implements PageProcessorProvider +{ + private static final long ADB_SEGMENT_WAIT_TIMEOUT = 60000L; + private final Queue pageProcessors = new LinkedList<>(); + private final AtomicBoolean isReadyForProcessing = new AtomicBoolean(false); + private final ReentrantLock lock = new ReentrantLock(); + private final Condition isReadyForProcessingCondition = lock.newCondition(); + + public GpfdistPageProcessorProvider() + { + } + + @Override + public void add(PageProcessor processor) + { + lock.lock(); + try { + pageProcessors.add(processor); + isReadyForProcessingCondition.signalAll(); + } + finally { + lock.unlock(); + } + } + + @Override + public PageProcessor take() + { + lock.lock(); + try { + if (!isReadyForProcessing.get()) { + long startTime = System.currentTimeMillis(); + while (pageProcessors.isEmpty()) { + try { + if (currentTimeMsProvider().get() - startTime > ADB_SEGMENT_WAIT_TIMEOUT) { + throw new RuntimeException( + format("Timeout :%d ms waiting for segments responses is exceeded", + ADB_SEGMENT_WAIT_TIMEOUT)); + } + isReadyForProcessingCondition.await(); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + isReadyForProcessing.set(true); + } + PageProcessor pageProcessor = pageProcessors.poll(); + pageProcessors.offer(pageProcessor); + return pageProcessor; + } + finally { + lock.unlock(); + } + } + + public static Supplier currentTimeMsProvider() + { + return System::currentTimeMillis; + } + + @Override + public Queue getAll() + { + return pageProcessors; + } + + @Override + public void clear() + { + lock.lock(); + try { + pageProcessors.clear(); + } + finally { + lock.unlock(); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSerializer.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSerializer.java new file mode 100644 index 000000000000..6f46094da2ae --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSerializer.java @@ -0,0 +1,59 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.process; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.adb.connector.encode.RowEncoder; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageSerializationResult; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageSerializer; +import io.trino.spi.Page; +import io.trino.spi.block.Block; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.List; + +public class GpfdistPageSerializer + implements PageSerializer +{ + private final List dataTypes; + private final RowEncoder rowEncoder; + + public GpfdistPageSerializer(List dataTypes, RowEncoder rowEncoder) + { + this.dataTypes = dataTypes; + this.rowEncoder = rowEncoder; + } + + @Override + public PageSerializationResult serialize(Page page) + { + long rowCount = 0L; + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + for (int position = 0; position < page.getPositionCount(); position++) { + for (int channel = 0; channel < dataTypes.size(); channel++) { + Block block = page.getBlock(channel); + rowEncoder.appendColumnValue(block, position); + } + out.write(rowEncoder.toByteArray()); + rowCount++; + } + out.flush(); + return new PageSerializationResult(out.toByteArray(), rowCount); + } + catch (IOException e) { + throw new RuntimeException("Failed to serialize data page: " + e.getMessage(), e); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSink.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSink.java new file mode 100644 index 000000000000..9719b6ca4744 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSink.java @@ -0,0 +1,130 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.process; + +import com.google.common.collect.ImmutableList; +import io.airlift.log.Logger; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.ExecutorServiceProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.PageProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContext; +import io.trino.spi.Page; +import io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorPageSink; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutorService; + +public class GpfdistPageSink + implements ConnectorPageSink +{ + private static final Logger log = Logger.get(GpfdistPageSink.class); + private final ContextManager writeContextManager; + private final WriteContext writeContext; + private final CompletableFuture queryLoadFuture; + private final ExecutorService executorService; + private CompletableFuture pageProcessingFuture; + + public GpfdistPageSink(ContextManager writeContextManager, + WriteContext writeContext, + CompletableFuture queryLoadFuture) + { + this.writeContext = writeContext; + this.queryLoadFuture = queryLoadFuture; + this.writeContextManager = writeContextManager; + executorService = ExecutorServiceProvider.WRITE_DATA_EXECUTOR_SERVICE; + pageProcessingFuture = CompletableFuture.completedFuture(null); + } + + @Override + public CompletableFuture appendPage(Page page) + { + pageProcessingFuture = CompletableFuture.runAsync(() -> { + if (writeContext.getAdbQueryException().get() == null) { + PageProcessor pageProcessor = writeContext.getPageProcessorProvider().take(); + pageProcessor.process(page); + } + else { + throw writeContext.getAdbQueryException().get(); + } + }, executorService); + return pageProcessingFuture; + } + + @Override + public CompletableFuture> finish() + { + return pageProcessingFuture + .thenCompose(_ -> closeCtx()) + .thenCompose(_ -> queryLoadFuture) + .thenApply(_ -> createFinishResult()) + .whenComplete((_, e) -> { + if (e != null) { + log.error("Failed to load data into adb: %s", e); + failContext(e); + throw new CompletionException(e); + } + writeContextManager.remove(writeContext.getId()); + }); + } + + private CompletableFuture closeCtx() + { + return CompletableFuture.runAsync(writeContext::close); + } + + private Collection createFinishResult() + { + Slice value = Slices.allocate(8); + value.setLong(0, writeContext.getMetadata().getPageSinkId().getId()); + return ImmutableList.of(value); + } + + @Override + public void abort() + { + TrinoException error = new TrinoException(StandardErrorCode.USER_CANCELED, "Query has been aborted"); + queryLoadFuture.completeExceptionally(error); + pageProcessingFuture.completeExceptionally(error); + failContext(error); + } + + private void failContext(Throwable e) + { + try { + writeContext.getError().set(e); + writeContext.close(); + } + finally { + writeContextManager.remove(writeContext.getId()); + } + } + + @Override + public long getCompletedBytes() + { + return writeContext.getCompletedBytes().get(); + } + + @Override + public long getMemoryUsage() + { + return writeContext.getMemoryUsage().get(); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSinkProvider.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSinkProvider.java new file mode 100644 index 000000000000..799d2dea24a8 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/process/GpfdistPageSinkProvider.java @@ -0,0 +1,140 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.process; + +import com.google.inject.Inject; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.encode.RowEncoderFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.CreateExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.DataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.ExecutorServiceProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.query.GpfdistLoadDataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableFormatConfigFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadataFactory; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.spi.connector.ConnectorInsertTableHandle; +import io.trino.spi.connector.ConnectorOutputTableHandle; +import io.trino.spi.connector.ConnectorPageSink; +import io.trino.spi.connector.ConnectorPageSinkId; +import io.trino.spi.connector.ConnectorPageSinkProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTransactionHandle; + +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static com.google.common.base.Preconditions.checkArgument; + +public class GpfdistPageSinkProvider + implements ConnectorPageSinkProvider +{ + private static final ExternalTableType EXTERNAL_TABLE_TYPE = ExternalTableType.READABLE; + private final AdbSqlClient client; + private final AdbPluginConfig pluginConfig; + private final GpfdistLoadMetadataFactory loadMetadataFactory; + private final ExecutorService loadQueryThreadExecutor; + private final ContextManager contextManager; + private final RowEncoderFactory rowEncoderFactory; + private final ExternalTableFormatConfigFactory externalTableFormatConfigFactory; + private final CreateExternalTableQueryFactory externalTableCreateQueryFactory; + private final InsertDataQueryFactory insertDataQueryFactory; + + @Inject + public GpfdistPageSinkProvider(@ForBaseJdbc JdbcClient client, + AdbPluginConfig pluginConfig, + GpfdistLoadMetadataFactory loadMetadataFactory, + ContextManager contextManager, + RowEncoderFactory rowEncoderFactory, + ExternalTableFormatConfigFactory externalTableFormatConfigFactory, + Set createExternalTableQueryFactories, + Set insertDataQueryFactories) + { + this.client = (AdbSqlClient) client; + this.pluginConfig = pluginConfig; + this.loadMetadataFactory = loadMetadataFactory; + this.contextManager = contextManager; + this.rowEncoderFactory = rowEncoderFactory; + this.externalTableFormatConfigFactory = externalTableFormatConfigFactory; + this.loadQueryThreadExecutor = ExecutorServiceProvider.LOAD_DATA_QUERY_EXECUTOR_SERVICE; + Map externalTableQueryFactoryMap = + createExternalTableQueryFactories.stream() + .collect(Collectors.toMap(CreateExternalTableQueryFactory::getExternalTableType, + Function.identity())); + externalTableCreateQueryFactory = externalTableQueryFactoryMap.get(EXTERNAL_TABLE_TYPE); + checkArgument(externalTableCreateQueryFactory != null, + "failed to get writable table query factory by externalTableType %s", + EXTERNAL_TABLE_TYPE); + Map insertDataFactoryMap = insertDataQueryFactories.stream() + .collect(Collectors.toMap(InsertDataQueryFactory::getExternalTableType, Function.identity())); + insertDataQueryFactory = insertDataFactoryMap.get(EXTERNAL_TABLE_TYPE); + checkArgument(insertDataQueryFactory != null, + "failed to get insert data query factory by externalTableType %s", + EXTERNAL_TABLE_TYPE); + } + + @Override + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, + ConnectorSession session, + ConnectorOutputTableHandle outputTableHandle, + ConnectorPageSinkId pageSinkId) + { + return createPageSinkInternal(session, outputTableHandle, pageSinkId); + } + + @Override + public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHandle, + ConnectorSession session, + ConnectorInsertTableHandle insertTableHandle, + ConnectorPageSinkId pageSinkId) + { + return createPageSinkInternal(session, (ConnectorOutputTableHandle) insertTableHandle, + pageSinkId); + } + + private ConnectorPageSink createPageSinkInternal(ConnectorSession session, + ConnectorOutputTableHandle outputTableHandle, + ConnectorPageSinkId pageSinkId) + { + GpfdistLoadMetadata loadMetadata = loadMetadataFactory.create(session, + (JdbcOutputTableHandle) outputTableHandle, + pageSinkId, + externalTableFormatConfigFactory.create()); + WriteContext writeContext = new WriteContext( + loadMetadata, + rowEncoderFactory.create(session, loadMetadata.getDataTypes()), + pluginConfig.getWriteBufferSize(), + new GpfdistPageProcessorProvider()); + DataTransferQueryExecutor loadDataExecutor = new GpfdistLoadDataTransferQueryExecutor(client, + session, + loadQueryThreadExecutor, + loadMetadata, + externalTableCreateQueryFactory, + insertDataQueryFactory); + contextManager.add(writeContext); + CompletableFuture queryLoadFuture = loadDataExecutor.execute(); + return new GpfdistPageSink(contextManager, writeContext, queryLoadFuture); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/CreateReadableExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/CreateReadableExternalTableQueryFactory.java new file mode 100644 index 000000000000..0e636574ad43 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/CreateReadableExternalTableQueryFactory.java @@ -0,0 +1,34 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.query; + +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; + +public class CreateReadableExternalTableQueryFactory + extends AbstractExternalTableQueryFactory +{ + @Override + public String createQuery(GpfdistMetadata metadata) + { + return createCommonQuery(metadata); + } + + @Override + public ExternalTableType getExternalTableType() + { + return ExternalTableType.READABLE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/GpfdistLoadDataTransferQueryExecutor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/GpfdistLoadDataTransferQueryExecutor.java new file mode 100644 index 000000000000..ff01b1432980 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/GpfdistLoadDataTransferQueryExecutor.java @@ -0,0 +1,67 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.query; + +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractDataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.CreateExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadata; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.SQLException; +import java.util.concurrent.ExecutorService; + +public class GpfdistLoadDataTransferQueryExecutor + extends AbstractDataTransferQueryExecutor +{ + private static final Logger log = Logger.get(GpfdistLoadDataTransferQueryExecutor.class); + private final GpfdistLoadMetadata loadMetadata; + + public GpfdistLoadDataTransferQueryExecutor(AdbSqlClient client, + ConnectorSession session, + ExecutorService executor, GpfdistLoadMetadata loadMetadata, + CreateExternalTableQueryFactory externalTableQueryFactory, + InsertDataQueryFactory insertDataQueryFactory) + { + super(client, session, executor, externalTableQueryFactory, insertDataQueryFactory); + this.loadMetadata = loadMetadata; + } + + @Override + protected void executeQueries() + throws SQLException + { + createReadableExternalTable(); + insertIntoExternalTable(); + } + + private void createReadableExternalTable() + throws SQLException + { + String sql = externalTableQueryFactory.createQuery(loadMetadata); + client.execute(session, connection, sql); + log.info("Executed create readable external table query: %s", sql); + } + + private void insertIntoExternalTable() + throws SQLException + { + PreparedQuery query = insertDataQueryFactory.create(session, connection, loadMetadata); + client.executeAsPreparedStatement(session, connection, query); + log.info("Executed insert into target table from external table query: %s", query.query()); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/InsertDataFromExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/InsertDataFromExternalTableQueryFactory.java new file mode 100644 index 000000000000..a25fbe9ddd0f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/load/query/InsertDataFromExternalTableQueryFactory.java @@ -0,0 +1,59 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.load.query; + +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistLoadMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.Connection; +import java.util.List; + +import static java.lang.String.format; + +public class InsertDataFromExternalTableQueryFactory + implements InsertDataQueryFactory +{ + private static final String COLUMN_DELIMITER = ", "; + + @Override + public PreparedQuery create(ConnectorSession session, Connection connection, GpfdistMetadata metadata) + { + GpfdistLoadMetadata loadMetadata = (GpfdistLoadMetadata) metadata; + String delimiter = COLUMN_DELIMITER; + String externalTableColumnNames = String.join(delimiter, loadMetadata.getColumnNames()); + String targetTableColumnNames; + if (loadMetadata.getPageSinkIdColumn().isPresent()) { + targetTableColumnNames = externalTableColumnNames + delimiter + loadMetadata.getPageSinkIdColumn().get(); + externalTableColumnNames = externalTableColumnNames + delimiter + loadMetadata.getPageSinkId().getId(); + } + else { + targetTableColumnNames = externalTableColumnNames; + } + return new PreparedQuery(format("INSERT INTO %s (%s) SELECT %s FROM %s", + loadMetadata.getTargetTable(), + targetTableColumnNames, + externalTableColumnNames, + loadMetadata.getSourceTable()), List.of()); + } + + @Override + public ExternalTableType getExternalTableType() + { + return ExternalTableType.READABLE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ContextId.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ContextId.java new file mode 100644 index 000000000000..483ca84a6f31 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ContextId.java @@ -0,0 +1,18 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +public record ContextId(String externalTable) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfig.java new file mode 100644 index 000000000000..9caa4c80078c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfig.java @@ -0,0 +1,20 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.adb.connector.encode.DataFormat; + +public record ExternalTableFormatConfig(char delimiter, String encoding, String nullValue, DataFormat dataFormat) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactory.java new file mode 100644 index 000000000000..c1297bec906f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactory.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +public interface ExternalTableFormatConfigFactory +{ + ExternalTableFormatConfig create(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactoryImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactoryImpl.java new file mode 100644 index 000000000000..8ddbe7a8a2e2 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableFormatConfigFactoryImpl.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import com.google.inject.Inject; +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.encode.DataFormatConfig; +import io.trino.plugin.adb.connector.encode.csv.CsvFormatConfig; + +public class ExternalTableFormatConfigFactoryImpl + implements ExternalTableFormatConfigFactory +{ + private final DataFormatConfig dataFormatConfig; + + @Inject + public ExternalTableFormatConfigFactoryImpl(DataFormatConfig dataFormatConfig) + { + this.dataFormatConfig = dataFormatConfig; + } + + @Override + public ExternalTableFormatConfig create() + { + if (dataFormatConfig.getDataFormat() == DataFormat.CSV) { + CsvFormatConfig csvFormatConfig = (CsvFormatConfig) dataFormatConfig; + return new ExternalTableFormatConfig(csvFormatConfig.getDelimiter(), + csvFormatConfig.getEncoding(), + csvFormatConfig.getNullValue(), + dataFormatConfig.getDataFormat()); + } + else { + throw new IllegalArgumentException("Unsupported encoder format: " + dataFormatConfig.getDataFormat()); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableType.java new file mode 100644 index 000000000000..fd9b575ee19b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableType.java @@ -0,0 +1,20 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +public enum ExternalTableType +{ + READABLE, + WRITABLE +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableUtil.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableUtil.java new file mode 100644 index 000000000000..a7b8ba63055b --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/ExternalTableUtil.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import java.util.UUID; + +public final class ExternalTableUtil +{ + private ExternalTableUtil() + { + } + + public static String createExternalTableName(ExternalTableType tableType) + { + return String.format("trino_external_%s_%s", tableType.name().toLowerCase(), + UUID.randomUUID().toString().replace("-", "")); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadata.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadata.java new file mode 100644 index 000000000000..e27d77e4ef0a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadata.java @@ -0,0 +1,58 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.spi.connector.ConnectorPageSinkId; + +import java.util.List; +import java.util.Optional; + +public class GpfdistLoadMetadata + extends GpfdistMetadata +{ + private final String targetTable; + private final ConnectorPageSinkId pageSinkId; + private final Optional pageSinkIdColumn; + + public GpfdistLoadMetadata(String targetTable, + String sourceTable, + List columnNames, + List dataTypes, + ConnectorPageSinkId pageSinkId, + Optional pageSinkIdColumn, + ExternalTableFormatConfig externalTableFormatConfig, + String gpfdistLocation) + { + super(sourceTable, columnNames, dataTypes, externalTableFormatConfig, gpfdistLocation); + this.pageSinkId = pageSinkId; + this.pageSinkIdColumn = pageSinkIdColumn; + this.targetTable = targetTable; + } + + public ConnectorPageSinkId getPageSinkId() + { + return pageSinkId; + } + + public Optional getPageSinkIdColumn() + { + return pageSinkIdColumn; + } + + public String getTargetTable() + { + return targetTable; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactory.java new file mode 100644 index 000000000000..fdbb979626fa --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactory.java @@ -0,0 +1,26 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.spi.connector.ConnectorPageSinkId; +import io.trino.spi.connector.ConnectorSession; + +public interface GpfdistLoadMetadataFactory +{ + GpfdistLoadMetadata create(ConnectorSession session, + JdbcOutputTableHandle tableHandle, + ConnectorPageSinkId pageSinkId, + ExternalTableFormatConfig externalTableFormatConfig); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactoryImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactoryImpl.java new file mode 100644 index 000000000000..1a41e6d7dc29 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLoadMetadataFactoryImpl.java @@ -0,0 +1,80 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import com.google.inject.Inject; +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.spi.connector.ConnectorPageSinkId; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; +import java.util.Optional; + +import static io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableUtil.createExternalTableName; + +public class GpfdistLoadMetadataFactoryImpl + implements GpfdistLoadMetadataFactory +{ + private static final Logger log = Logger.get(GpfdistLoadMetadataFactoryImpl.class); + private final GpfdistLocationFactory gpfdistLocationFactory; + private final AdbSqlClient sqlClient; + + @Inject + public GpfdistLoadMetadataFactoryImpl( + @ForBaseJdbc JdbcClient sqlClient, + GpfdistLocationFactory gpfdistLocationFactory) + { + this.gpfdistLocationFactory = gpfdistLocationFactory; + this.sqlClient = (AdbSqlClient) sqlClient; + } + + @Override + public GpfdistLoadMetadata create(ConnectorSession session, + JdbcOutputTableHandle tableHandle, + ConnectorPageSinkId pageSinkId, + ExternalTableFormatConfig externalTableFormatConfig) + { + try { + ExternalTableType externalTableType = ExternalTableType.READABLE; + String targetTableName = sqlClient.getTargetTableName(tableHandle); + String externalTableName = createExternalTableName(externalTableType); + List columnNames = tableHandle.getColumnNames().stream() + .map(sqlClient::quoted) + .toList(); + List dataTypes = sqlClient.getColumnDataTypes(session, tableHandle); + Optional pageSinkIdColumn = tableHandle.getPageSinkIdColumnName() + .map(sqlClient::quoted); + String gpfdistLocation = gpfdistLocationFactory.create(externalTableName, externalTableType); + return new GpfdistLoadMetadata( + targetTableName, + externalTableName, + columnNames, + dataTypes, + pageSinkId, + pageSinkIdColumn, + externalTableFormatConfig, + gpfdistLocation); + } + catch (Exception e) { + String errMsg = "Failed to create gpfdistLoadMetadata: %s" + e.getMessage(); + log.error(errMsg, e); + throw new RuntimeException(errMsg, e); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactory.java new file mode 100644 index 000000000000..35970631b9af --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactory.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +public interface GpfdistLocationFactory +{ + String create(String externalTableName, ExternalTableType externalTableType); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactoryImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactoryImpl.java new file mode 100644 index 000000000000..6330da691c57 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistLocationFactoryImpl.java @@ -0,0 +1,62 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import com.google.inject.Inject; +import io.airlift.http.server.HttpServerInfo; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.GpfdistServerConfig; + +import java.net.URI; + +public class GpfdistLocationFactoryImpl + implements GpfdistLocationFactory +{ + private final GpfdistServerConfig config; + private final HttpServerInfo httpServerInfo; + + @Inject + public GpfdistLocationFactoryImpl(GpfdistServerConfig config, HttpServerInfo httpServerInfo) + { + this.config = config; + this.httpServerInfo = httpServerInfo; + } + + @Override + public String create(String externalTableName, ExternalTableType externalTableType) + { + String protocol = config.isServerSslEnabled() ? "gpfdists" : "gpfdist"; + URI uri = config.isServerSslEnabled() ? httpServerInfo.getHttpsExternalUri() : httpServerInfo.getHttpExternalUri(); + String host = uri.getHost(); + int port = config.getServerExternalPort() > 0 ? config.getServerExternalPort() : uri.getPort(); + return String.format("%s://%s:%d/gpfdist/%s/%s", + protocol, + host, + port, + getOperationPath(externalTableType), + externalTableName); + } + + private String getOperationPath(ExternalTableType externalTableType) + { + if (externalTableType == ExternalTableType.READABLE) { + return "read"; + } + else if (externalTableType == ExternalTableType.WRITABLE) { + return "write"; + } + else { + throw new IllegalArgumentException("Unsupported external table type: " + externalTableType); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistMetadata.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistMetadata.java new file mode 100644 index 000000000000..fd6463e3d77e --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistMetadata.java @@ -0,0 +1,86 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; + +import java.util.List; +import java.util.Objects; + +public class GpfdistMetadata +{ + private final String sourceTable; + private final List columnNames; + private final List dataTypes; + private final ExternalTableFormatConfig externalTableFormatConfig; + private final String gpfdistLocation; + + public GpfdistMetadata( + String sourceTable, + List columnNames, + List dataTypes, + ExternalTableFormatConfig externalTableFormatConfig, + String gpfdistLocation) + { + this.sourceTable = sourceTable; + this.columnNames = columnNames; + this.dataTypes = dataTypes; + this.externalTableFormatConfig = externalTableFormatConfig; + this.gpfdistLocation = gpfdistLocation; + } + + public String getSourceTable() + { + return sourceTable; + } + + public List getColumnNames() + { + return columnNames; + } + + public List getDataTypes() + { + return dataTypes; + } + + public ExternalTableFormatConfig getExternalTableFormatConfig() + { + return externalTableFormatConfig; + } + + public String getGpfdistLocation() + { + return gpfdistLocation; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + GpfdistMetadata that = (GpfdistMetadata) o; + return Objects.equals(sourceTable, that.sourceTable) && Objects.equals(columnNames, that.columnNames) && Objects.equals(dataTypes, that.dataTypes) && Objects.equals(externalTableFormatConfig, that.externalTableFormatConfig) && Objects.equals(gpfdistLocation, that.gpfdistLocation); + } + + @Override + public int hashCode() + { + return Objects.hash(sourceTable, columnNames, dataTypes, externalTableFormatConfig, gpfdistLocation); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadata.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadata.java new file mode 100644 index 000000000000..02f628dda698 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadata.java @@ -0,0 +1,59 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ColumnHandle; + +import java.util.List; + +public class GpfdistUnloadMetadata + extends GpfdistMetadata +{ + private final JdbcTableHandle targetTableHandle; + private final List columnHandles; + private final JdbcSplit split; + + public GpfdistUnloadMetadata(JdbcTableHandle targetTableHandle, + String sourceTable, + List columnNames, + List dataTypes, + List columnHandles, + JdbcSplit split, + ExternalTableFormatConfig externalTableFormatConfig, + String gpfdistLocation) + { + super(sourceTable, columnNames, dataTypes, externalTableFormatConfig, gpfdistLocation); + this.targetTableHandle = targetTableHandle; + this.columnHandles = columnHandles; + this.split = split; + } + + public JdbcTableHandle getTargetTableHandle() + { + return targetTableHandle; + } + + public JdbcSplit getSplit() + { + return split; + } + + public List getColumnHandles() + { + return columnHandles; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactory.java new file mode 100644 index 000000000000..9f3ab15d5ce0 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactory.java @@ -0,0 +1,30 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorSession; + +import java.util.List; + +public interface GpfdistUnloadMetadataFactory +{ + GpfdistUnloadMetadata create(ConnectorSession session, + JdbcTableHandle tableHandle, + JdbcSplit split, + List columnHandles, + ExternalTableFormatConfig externalTableFormatConfig); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactoryImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactoryImpl.java new file mode 100644 index 000000000000..386ea9460d8d --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/metadata/GpfdistUnloadMetadataFactoryImpl.java @@ -0,0 +1,94 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.metadata; + +import com.google.inject.Inject; +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.datatype.ColumnDataType; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorSession; + +import java.util.ArrayList; +import java.util.List; + +import static io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableUtil.createExternalTableName; + +public class GpfdistUnloadMetadataFactoryImpl + implements GpfdistUnloadMetadataFactory +{ + private static final Logger log = Logger.get(GpfdistUnloadMetadataFactoryImpl.class); + private static final ExternalTableType EXTERNAL_TABLE_TYPE = ExternalTableType.WRITABLE; + private final GpfdistLocationFactory gpfdistLocationFactory; + private final AdbSqlClient sqlClient; + + @Inject + public GpfdistUnloadMetadataFactoryImpl(@ForBaseJdbc JdbcClient sqlClient, + GpfdistLocationFactory gpfdistLocationFactory) + { + this.gpfdistLocationFactory = gpfdistLocationFactory; + this.sqlClient = (AdbSqlClient) sqlClient; + } + + @Override + public GpfdistUnloadMetadata create(ConnectorSession session, + JdbcTableHandle tableHandle, + JdbcSplit split, + List columnHandles, + ExternalTableFormatConfig externalTableFormatConfig) + { + try { + List columnNames = new ArrayList<>(); + List dataTypes = new ArrayList<>(); + String externalTableName = createExternalTableName(EXTERNAL_TABLE_TYPE); + JdbcTableHandle jdbcTableHandle = tableHandle.intersectedWithConstraint(split.getDynamicFilter() + .transformKeys(ColumnHandle.class::cast)); + initColumnsMetadata(session, columnHandles, columnNames, dataTypes); + String gpfdistLocation = gpfdistLocationFactory.create(externalTableName, EXTERNAL_TABLE_TYPE); + return new GpfdistUnloadMetadata( + jdbcTableHandle, + externalTableName, + columnNames, + dataTypes, + columnHandles, + split, + externalTableFormatConfig, + gpfdistLocation); + } + catch (Exception e) { + String errMsg = "Failed to create gpfdistUnloadMetadata: " + e.getMessage(); + log.error(errMsg, e); + throw new RuntimeException(errMsg, e); + } + } + + private void initColumnsMetadata(ConnectorSession session, + List columnHandles, + List columnNames, + List dataTypes) + { + List jdbcColumnHandles = columnHandles.stream() + .map(columnHandle -> ((JdbcColumnHandle) columnHandle)) + .toList(); + jdbcColumnHandles.forEach(columnHandle -> { + columnNames.add(columnHandle.getColumnName()); + dataTypes.add(sqlClient.getColumnDataType(session, columnHandle.getJdbcTypeHandle())); + }); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistResource.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistResource.java new file mode 100644 index 000000000000..6f558e643be1 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistResource.java @@ -0,0 +1,251 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server; + +import com.google.inject.Inject; +import io.airlift.log.Logger; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.ExecutorServiceProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.GpfdistPacketBuilder; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.context.WriteContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.process.GpfdistPageProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.load.process.GpfdistPageSerializer; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistReadableRequest; +import io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistWritableRequest; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessorFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistSegmentRequestProcessor; +import jakarta.ws.rs.Consumes; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.POST; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.container.AsyncResponse; +import jakarta.ws.rs.container.Suspended; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.MediaType; +import jakarta.ws.rs.core.Response; + +import java.io.InputStream; +import java.io.PipedInputStream; +import java.io.PipedOutputStream; +import java.util.Optional; +import java.util.concurrent.ExecutorService; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.adb.connector.protocol.gpfdist.GpfdistUtil.createGpfdistFileName; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_PROTO; +import static java.lang.String.format; + +@Path("/gpfdist") +public class GpfdistResource +{ + private static final Logger log = Logger.get(GpfdistResource.class); + private static final int GPFDIST_FOR_READ_PROTOCOL_VERSION = 1; + private static final int GPFDIST_FOR_WRITE_PROTOCOL_VERSION = 0; + private final ContextManager writeContextManager; + private final ContextManager readContextManager; + private final ExecutorService executorService; + private final AdbPluginConfig pluginConfig; + private final InputDataProcessorFactory inputDataProcessorFactory; + + @Inject + public GpfdistResource(ContextManager writeContextManager, + ContextManager readContextManager, + AdbPluginConfig pluginConfig, + InputDataProcessorFactory inputDataProcessorFactory) + { + this.writeContextManager = writeContextManager; + this.readContextManager = readContextManager; + this.pluginConfig = pluginConfig; + this.inputDataProcessorFactory = inputDataProcessorFactory; + this.executorService = ExecutorServiceProvider.GPFDIST_HTTP_REQUEST_EXECUTOR_SERVICE; + } + + @GET + @Produces("text/plain") + @Path("/read/{tableName}") + public void get(@PathParam("tableName") String tableName, @Context HttpHeaders headers, + @Suspended AsyncResponse asyncResponse) + { + GpfdistReadableRequest request = GpfdistReadableRequest.create(tableName, headers.getRequestHeaders()); + checkArgument(request.getGpProtocol() == GPFDIST_FOR_READ_PROTOCOL_VERSION, + format("Gpfdist protocol of version %d is only supported", GPFDIST_FOR_READ_PROTOCOL_VERSION)); + log.info("Input GET gpfdist request: %s", request); + Optional writeContextOptional = writeContextManager.get(new ContextId(tableName)); + if (writeContextOptional.isPresent()) { + WriteContext writeContext = writeContextOptional.get(); + executorService.submit(() -> processGetRequest(asyncResponse, request, writeContext)); + } + else { + log.info("There is no data for loading responded by request: %s", request); + asyncResponse.resume(createOkGetResponseBuilder(request) + .entity(new GpfdistPacketBuilder(createGpfdistFileName(tableName)) + .createSingleEmptyDataPacket()) + .build()); + } + } + + private void processGetRequest(AsyncResponse asyncResponse, GpfdistReadableRequest request, + WriteContext writeContext) + { + int bufferSizeInBytes = Long.valueOf(pluginConfig.getWriteBufferSize().toBytes()).intValue(); + try (PipedOutputStream outputStream = new PipedOutputStream(); + PipedInputStream inputStream = new PipedInputStream(outputStream, bufferSizeInBytes)) { + GpfdistPageProcessor gpfdistPageProcessor = new GpfdistPageProcessor(outputStream, + request, + writeContext, + new GpfdistPageSerializer(writeContext.getMetadata().getDataTypes(), writeContext.getRowEncoder())); + writeContext.getPageProcessorProvider().add(gpfdistPageProcessor); + asyncResponse.resume(createOkGetResponseBuilder(request) + .entity(inputStream) + .build()); + log.info("Request %s completed successfully", request); + } + catch (Exception e) { + asyncResponse.resume(Response.serverError() + .header(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_PLAIN) + .header(X_GP_PROTO, request.getGpProtocol()) + .entity(e.getMessage()) + .build()); + log.error("Request %s is failed. Error: %s", + request, + e.getMessage()); + } + } + + private Response.ResponseBuilder createOkGetResponseBuilder(GpfdistReadableRequest request) + { + return Response.ok() + .header(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_PLAIN) + .header(HttpHeaders.CACHE_CONTROL, "no-cache") + .header(X_GP_PROTO, request.getGpProtocol()); + } + + @POST + @Consumes("*/*") + @Path("/write/{tableName}") + public void post(@PathParam("tableName") String tableName, InputStream data, @Context HttpHeaders headers, + @Suspended AsyncResponse asyncResponse) + { + try { + GpfdistWritableRequest request = GpfdistWritableRequest.create(tableName, headers.getRequestHeaders()); + log.debug("Received POST request: %s", request); + checkArgument(request.getGpProtocol() == GPFDIST_FOR_WRITE_PROTOCOL_VERSION, + format("Gpfdist protocol version %s for write operation is supported", + GPFDIST_FOR_WRITE_PROTOCOL_VERSION)); + Optional readContextOptional = readContextManager.get(new ContextId(tableName)); + if (readContextOptional.isEmpty()) { + processNotFoundQueryRequest(tableName, asyncResponse, request); + } + else { + ReadContext readContext = readContextOptional.get(); + if (initialRequest(request)) { + processInitialRequest(asyncResponse, readContext, request); + } + else if (!isLast(request)) { + processDataRequest(data, asyncResponse, readContext, request); + } + else { + processTearDownRequest(asyncResponse, readContext, request); + } + } + } + catch (Exception e) { + failWriteResponse(asyncResponse, e); + } + } + + private static void processNotFoundQueryRequest(String tableName, AsyncResponse asyncResponse, + GpfdistWritableRequest request) + { + String errorMessage = "No active query for writeable table: " + tableName; + asyncResponse.resume(Response.status(Response.Status.BAD_REQUEST.getStatusCode(), errorMessage) + .header(X_GP_PROTO, request.getGpProtocol()) + .build()); + log.error("Failed to processed request: %s. " + errorMessage, request); + } + + private void processInitialRequest(AsyncResponse asyncResponse, ReadContext readContext, + GpfdistWritableRequest request) + { + InputDataProcessor dataProcessor = inputDataProcessorFactory.create(readContext.getRowDecoder(), + readContext.getRowProcessingService()); + readContext.getSegmentDataProcessors().putIfAbsent(request.getSegmentId(), + new GpfdistSegmentRequestProcessor(request.getSegmentId(), dataProcessor)); + asyncResponse.resume(Response.ok() + .header(X_GP_PROTO, request.getGpProtocol()) + .build()); + log.debug("Request for initial data transferring completed successfully: %s", request); + } + + private void processDataRequest(InputStream data, AsyncResponse asyncResponse, ReadContext readContext, + GpfdistWritableRequest request) + { + executorService.submit(() -> { + try { + GpfdistSegmentRequestProcessor processor = getSegmentProcessor(readContext, request.getSegmentId()); + processor.process(data); + asyncResponse.resume(Response.ok() + .header(X_GP_PROTO, request.getGpProtocol()) + .build()); + log.debug("Processing request for transferring data completed successfully: %s", request); + } + catch (Exception e) { + failWriteResponse(asyncResponse, e); + } + }); + } + + private void processTearDownRequest(AsyncResponse asyncResponse, ReadContext readContext, + GpfdistWritableRequest request) + { + GpfdistSegmentRequestProcessor processor = getSegmentProcessor(readContext, request.getSegmentId()); + processor.stop(); + asyncResponse.resume(Response.ok() + .header(X_GP_PROTO, request.getGpProtocol()) + .build()); + log.debug("Processing request for finishing data transferring completed successfully: %s", request); + } + + private GpfdistSegmentRequestProcessor getSegmentProcessor(ReadContext readContext, Integer segmentId) + { + return Optional.ofNullable(readContext.getSegmentDataProcessors().get(segmentId)) + .orElseThrow(() -> new IllegalStateException( + "Failed to get segment request processor by segmentId: " + segmentId)); + } + + private void failWriteResponse(AsyncResponse asyncResponse, Exception e) + { + asyncResponse.resume(Response.serverError() + .header(X_GP_PROTO, GPFDIST_FOR_WRITE_PROTOCOL_VERSION) + .build()); + log.error("Failed to process request: %s", e); + } + + private boolean initialRequest(GpfdistWritableRequest request) + { + return request.getGpSequence() == 1; + } + + private boolean isLast(GpfdistWritableRequest request) + { + return request.isLastChunk().isPresent() && request.isLastChunk().get(); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerConfig.java new file mode 100644 index 000000000000..67944c8c16dc --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerConfig.java @@ -0,0 +1,147 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server; + +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; + +public class GpfdistServerConfig +{ + private String serverHost; + private String serverExternalHost; + private int serverPort; + private int serverExternalPort; + private boolean serverSslEnabled; + private String serverSslKeystorePath; + private String serverSslKeystorePassword; + private String serverSslTrustStorePath; + private String serverSslTrustStorePassword; + + public String getServerHost() + { + return this.serverHost; + } + + @Config("adb.gpfdist.server.host") + @ConfigDescription("Host to bind the gpfdist server to. Defaults binds to node address") + public GpfdistServerConfig setServerHost(String serverHost) + { + this.serverHost = serverHost; + return this; + } + + public String getServerExternalHost() + { + return this.serverExternalHost; + } + + @Config("adb.gpfdist.server.external-host") + @ConfigDescription("Gpfdist server host that will be announced to Adb segments. Default is null (use adb.gpfdist.server.host)") + public GpfdistServerConfig setServerExternalHost(String serverExternalHost) + { + this.serverExternalHost = serverExternalHost; + return this; + } + + public int getServerPort() + { + return this.serverPort; + } + + @Config("adb.gpfdist.server.port") + @ConfigDescription("Port to bind the gpfdist server to. Default is 0 (pick random free port)") + public GpfdistServerConfig setServerPort(int serverPort) + { + this.serverPort = serverPort; + return this; + } + + public int getServerExternalPort() + { + return this.serverExternalPort; + } + + @Config("adb.gpfdist.server.external-port") + @ConfigDescription("Gpfdist server port that will be announced to Adb segments. Defaults is 0 (use adb.gpfdist.server.port)") + public GpfdistServerConfig setServerExternalPort(int serverExternalPort) + { + this.serverExternalPort = serverExternalPort; + return this; + } + + public boolean isServerSslEnabled() + { + return this.serverSslEnabled; + } + + @Config("adb.gpfdist.server.ssl.enabled") + @ConfigDescription("Whether to use SSL for communication with Adb segments") + public GpfdistServerConfig setServerSslEnabled(boolean serverSslEnabled) + { + this.serverSslEnabled = serverSslEnabled; + return this; + } + + public String getServerSslKeystorePath() + { + return this.serverSslKeystorePath; + } + + @Config("adb.gpfdist.server.ssl.keystore.path") + @ConfigDescription("Path to SSL keystore file") + public GpfdistServerConfig setServerSslKeystorePath(String serverSslKeystorePath) + { + this.serverSslKeystorePath = serverSslKeystorePath; + return this; + } + + public String getServerSslKeystorePassword() + { + return this.serverSslKeystorePassword; + } + + @Config("adb.gpfdist.server.ssl.keystore.password") + @ConfigDescription("SSL keystore password") + public GpfdistServerConfig setServerSslKeystorePassword(String serverSslKeystorePassword) + { + this.serverSslKeystorePassword = serverSslKeystorePassword; + return this; + } + + public String getServerSslTrustStorePath() + { + return this.serverSslTrustStorePath; + } + + @Config("adb.gpfdist.server.ssl.truststore.path") + @ConfigDescription("Path to SSL truststore file") + public GpfdistServerConfig setServerSslTrustStorePath(String serverSslTrustStorePath) + { + this.serverSslTrustStorePath = serverSslTrustStorePath; + return this; + } + + public String getServerSslTrustStorePassword() + { + return this.serverSslTrustStorePassword; + } + + @Config("adb.gpfdist.server.ssl.truststore.password") + @ConfigDescription("SSL truststore password") + public GpfdistServerConfig setServerSslTrustStorePassword(String serverSslTrustStorePassword) + { + this.serverSslTrustStorePassword = serverSslTrustStorePassword; + return this; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerModule.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerModule.java new file mode 100644 index 000000000000..784ea59e0a3e --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistServerModule.java @@ -0,0 +1,103 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server; + +import com.google.inject.Binder; +import com.google.inject.Provides; +import com.google.inject.Singleton; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.event.client.NullEventClient; +import io.airlift.http.server.HttpServer; +import io.airlift.http.server.HttpServerConfig; +import io.airlift.http.server.HttpServerInfo; +import io.airlift.http.server.HttpServerProvider; +import io.airlift.http.server.HttpsConfig; +import io.airlift.http.server.RequestStats; +import io.airlift.http.server.TheServlet; +import io.airlift.jaxrs.JaxrsBinder; +import io.airlift.jaxrs.JaxrsModule; +import io.airlift.json.JsonModule; +import io.airlift.node.NodeInfo; +import jakarta.servlet.Servlet; + +import java.util.Optional; +import java.util.Set; + +import static io.airlift.configuration.ConfigBinder.configBinder; + +public class GpfdistServerModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + super.install(new JsonModule()); + super.install(new JaxrsModule()); + configBinder(binder).bindConfig(GpfdistServerConfig.class); + JaxrsBinder.jaxrsBinder(binder).bind(GpfdistResource.class); + } + + @Provides + @Singleton + public static HttpServerConfig getHttpConfig(GpfdistServerConfig config) + { + return new HttpServerConfig() + .setHttpEnabled(!config.isServerSslEnabled()) + .setHttpsEnabled(config.isServerSslEnabled()) + .setLogEnabled(false) + .setHttpPort(config.getServerPort()); + } + + @Provides + @Singleton + public static Optional getHttpsConfig(GpfdistServerConfig config) + { + return !config.isServerSslEnabled() + ? Optional.empty() + : Optional.of(new HttpsConfig() + .setHttpsPort(config.getServerPort()) + .setKeystorePath(config.getServerSslKeystorePath()) + .setKeystorePassword(config.getServerSslKeystorePassword()) + .setTrustStorePath(config.getServerSslTrustStorePath()) + .setTrustStorePassword(config.getServerSslTrustStorePassword())); + } + + @Provides + @Singleton + public static HttpServerProvider getHttpServerProvider( + HttpServerConfig httpConfig, + Optional httpsConfig, + @TheServlet Servlet servlet, + NodeInfo nodeInfo, + HttpServerInfo httpServerInfo, + RequestStats requestStats) + { + return new HttpServerProvider( + httpServerInfo, + nodeInfo, + httpConfig, + httpsConfig, + servlet, + Set.of(), + Set.of(), + Set.of(), + false, + false, + false, + HttpServer.ClientCertificate.NONE, + requestStats, + new NullEventClient(), + Optional.empty()); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistReadableRequest.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistReadableRequest.java new file mode 100644 index 000000000000..abcec168d579 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistReadableRequest.java @@ -0,0 +1,157 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server.request; + +import jakarta.ws.rs.core.MultivaluedMap; + +import java.util.Optional; +import java.util.UUID; + +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_CID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_CSV_OPT; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_DATABASE; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_LINE_DELIM_LENGTH; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_MASTER_HOST; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_MASTER_PORT; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_PROTO; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEGMENT_COUNT; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEGMENT_ID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEG_DATADIR; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEG_PG_CONF; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SESSION_ID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SN; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_XID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_X_GP_USER; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_X_SEG_PORT; + +public class GpfdistReadableRequest +{ + private String requestId; + private String transactionId; + private String commandId; + private String scanId; + private Optional segmentId; + private Optional segmentsCount; + private Optional lineDelimiterLength; + private short gpProtocol; + private Optional gpMasterHost; + private Optional gpMasterPort; + private Optional gpcsvFormat; + private Optional gpSegmentConfigPath; + private Optional gpSegmentDataDirectory; + private Optional gpDatabase; + private Optional gpUser; + private Optional gpSegmentPort; + private Optional gpSessionId; + + public GpfdistReadableRequest( + String requestId, + String transactionId, + String commandId, + String scanId, + Optional segmentId, + Optional segmentsCount, + Optional lineDelimiterLength, + short gpProtocol, + Optional gpMasterHost, + Optional gpMasterPort, + Optional gpcsvFormat, + Optional gpSegmentConfigPath, + Optional gpSegmentDataDirectory, + Optional gpDatabase, + Optional gpUser, + Optional gpSegmentPort, + Optional gpSessionId) + { + this.requestId = requestId; + this.transactionId = transactionId; + this.commandId = commandId; + this.scanId = scanId; + this.segmentId = segmentId; + this.segmentsCount = segmentsCount; + this.lineDelimiterLength = lineDelimiterLength; + this.gpProtocol = gpProtocol; + this.gpMasterHost = gpMasterHost; + this.gpMasterPort = gpMasterPort; + this.gpcsvFormat = gpcsvFormat; + this.gpSegmentConfigPath = gpSegmentConfigPath; + this.gpSegmentDataDirectory = gpSegmentDataDirectory; + this.gpDatabase = gpDatabase; + this.gpUser = gpUser; + this.gpSegmentPort = gpSegmentPort; + this.gpSessionId = gpSessionId; + } + + public static GpfdistReadableRequest create(String tableName, MultivaluedMap values) + { + return new GpfdistReadableRequest( + createRequestId(tableName), + values.getFirst(X_GP_XID), + values.getFirst(X_GP_CID), + values.getFirst(X_GP_SN), + Optional.ofNullable(values.get(X_GP_SEGMENT_ID)) + .map(v -> Integer.parseInt(v.getFirst())), + Optional.ofNullable(values.get(X_GP_SEGMENT_COUNT)) + .map(v -> Integer.parseInt(v.getFirst())), + Optional.ofNullable(values.get(X_GP_LINE_DELIM_LENGTH)) + .map(v -> Integer.parseInt(v.getFirst())), + Short.parseShort(values.get(X_GP_PROTO).getFirst()), + Optional.ofNullable(values.getFirst(X_GP_MASTER_HOST)), + Optional.ofNullable(values.get(X_GP_MASTER_PORT)) + .map(v -> Integer.parseInt(v.getFirst())), + Optional.ofNullable(values.getFirst(X_GP_CSV_OPT)), + Optional.ofNullable(values.getFirst(X_GP_SEG_PG_CONF)), + Optional.ofNullable(values.getFirst(X_GP_SEG_DATADIR)), + Optional.ofNullable(values.getFirst(X_GP_DATABASE)), + Optional.ofNullable(values.getFirst(X_GP_X_GP_USER)), + Optional.ofNullable(values.get(X_GP_X_SEG_PORT)) + .map(v -> Integer.parseInt(v.getFirst())), + Optional.ofNullable(values.getFirst(X_GP_SESSION_ID)) + .map(Integer::parseInt)); + } + + private static String createRequestId(String tableName) + { + return tableName + "_" + UUID.randomUUID(); + } + + public short getGpProtocol() + { + return gpProtocol; + } + + @Override + public String toString() + { + return "GpfdistReadableRequest{" + + "requestId='" + requestId + '\'' + + ", transactionId='" + transactionId + '\'' + + ", commandId='" + commandId + '\'' + + ", scanId='" + scanId + '\'' + + ", segmentId=" + segmentId + + ", segmentsCount=" + segmentsCount + + ", lineDelimiterLength=" + lineDelimiterLength + + ", gpProtocol=" + gpProtocol + + ", gpMasterHost=" + gpMasterHost + + ", gpMasterPort=" + gpMasterPort + + ", gpcsvFormat=" + gpcsvFormat + + ", gpSegmentConfigPath=" + gpSegmentConfigPath + + ", gpSegmentDataDirectory=" + gpSegmentDataDirectory + + ", gpDatabase=" + gpDatabase + + ", gpUser=" + gpUser + + ", gpSegmentPort=" + gpSegmentPort + + ", gpSessionId=" + gpSessionId + + '}'; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistRequestHeader.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistRequestHeader.java new file mode 100644 index 000000000000..b7530cbf93e5 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistRequestHeader.java @@ -0,0 +1,41 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server.request; + +public final class GpfdistRequestHeader +{ + public static final String X_GP_XID = "X-GP-XID"; + public static final String X_GP_CID = "X-GP-CID"; + public static final String X_GP_SN = "X-GP-SN"; + public static final String X_GP_SEGMENT_ID = "X-GP-SEGMENT-ID"; + public static final String X_GP_SEGMENT_COUNT = "X-GP-SEGMENT-COUNT"; + public static final String X_GP_LINE_DELIM_LENGTH = "X-GP-LINE-DELIM-LENGTH"; + public static final String X_GP_PROTO = "X-GP-PROTO"; + public static final String X_GP_PROTOCOL_VERSION = "X-GPFDIST-VERSION"; + public static final String X_GP_SEQUENCE = "X-GP-SEQ"; + public static final String X_GP_DONE = "X-GP-DONE"; + public static final String X_GP_MASTER_HOST = "X-GP-MASTER_HOST"; + public static final String X_GP_MASTER_PORT = "X-GP-MASTER_PORT"; + public static final String X_GP_CSV_OPT = "X-GP-CSVOPT"; + public static final String X_GP_SEG_PG_CONF = "X-GP_SEG_PG_CONF"; + public static final String X_GP_SEG_DATADIR = "X-GP_SEG_DATADIR"; + public static final String X_GP_DATABASE = "X-GP-DATABASE"; + public static final String X_GP_X_GP_USER = "X-GP-X-GP-USER"; + public static final String X_GP_X_SEG_PORT = "X-GP-X-GP-SEG-PORT"; + public static final String X_GP_SESSION_ID = "x-gp-session-id"; + + private GpfdistRequestHeader() + { + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistWritableRequest.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistWritableRequest.java new file mode 100644 index 000000000000..c6406a823e2c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/request/GpfdistWritableRequest.java @@ -0,0 +1,142 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.server.request; + +import jakarta.ws.rs.core.MultivaluedMap; + +import java.util.Optional; +import java.util.UUID; + +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_CID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_DATABASE; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_DONE; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_LINE_DELIM_LENGTH; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_PROTO; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_PROTOCOL_VERSION; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEGMENT_COUNT; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEGMENT_ID; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SEQUENCE; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_SN; +import static io.trino.plugin.adb.connector.protocol.gpfdist.server.request.GpfdistRequestHeader.X_GP_XID; + +public class GpfdistWritableRequest +{ + private String requestId; + private String transactionId; + private String commandId; + private String scanId; + private Integer segmentId; + private Optional segmentsCount; + private Optional lineDelimiterLength; + private short gpProtocol; + private Optional gpProtocolVersion; + private int gpSequence; + private Optional isLastChunk; + private Optional gpDatabase; + + public GpfdistWritableRequest( + String requestId, + String transactionId, + String commandId, + String scanId, + Integer segmentId, + Optional segmentsCount, + Optional lineDelimiterLength, + short gpProtocol, + Optional gpProtocolVersion, + int gpSequence, + Optional isLastChunk, + Optional gpDatabase) + { + this.requestId = requestId; + this.transactionId = transactionId; + this.commandId = commandId; + this.scanId = scanId; + this.segmentId = segmentId; + this.segmentsCount = segmentsCount; + this.lineDelimiterLength = lineDelimiterLength; + this.gpProtocol = gpProtocol; + this.gpProtocolVersion = gpProtocolVersion; + this.gpSequence = gpSequence; + this.isLastChunk = isLastChunk; + this.gpDatabase = gpDatabase; + } + + public static GpfdistWritableRequest create(String tableName, MultivaluedMap values) + { + return new GpfdistWritableRequest( + createRequestId(tableName), + values.getFirst(X_GP_XID), + values.getFirst(X_GP_CID), + values.getFirst(X_GP_SN), + Optional.ofNullable(values.get(X_GP_SEGMENT_ID)) + .map(v -> Integer.parseInt(v.getFirst())) + .orElseThrow( + () -> new IllegalArgumentException("Request header not found: " + X_GP_SEGMENT_ID)), + Optional.ofNullable(values.get(X_GP_SEGMENT_COUNT)) + .map(v -> Integer.parseInt(v.getFirst())), + Optional.ofNullable(values.get(X_GP_LINE_DELIM_LENGTH)) + .map(v -> Integer.parseInt(v.getFirst())), + Short.parseShort(values.get(X_GP_PROTO).getFirst()), + Optional.ofNullable(values.getFirst(X_GP_PROTOCOL_VERSION)), + Integer.parseInt(values.get(X_GP_SEQUENCE).getFirst()), + Optional.ofNullable(values.get(X_GP_DONE)) + .map(v -> v.getFirst().equals("1")), + Optional.ofNullable(values.getFirst(X_GP_DATABASE))); + } + + private static String createRequestId(String tableName) + { + return tableName + "_" + UUID.randomUUID(); + } + + public Integer getSegmentId() + { + return segmentId; + } + + public short getGpProtocol() + { + return gpProtocol; + } + + public int getGpSequence() + { + return gpSequence; + } + + public Optional isLastChunk() + { + return isLastChunk; + } + + @Override + public String toString() + { + return "GpfdistWritableRequest{" + + "requestId='" + requestId + '\'' + + ", transactionId='" + transactionId + '\'' + + ", commandId='" + commandId + '\'' + + ", scanId='" + scanId + '\'' + + ", segmentId=" + segmentId + + ", segmentsCount=" + segmentsCount + + ", lineDelimiterLength=" + lineDelimiterLength + + ", gpProtocol=" + gpProtocol + + ", gpProtocolVersion=" + gpProtocolVersion + + ", gpSequence=" + gpSequence + + ", isLastChunk=" + isLastChunk + + ", gpDatabase=" + gpDatabase + + '}'; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/GpfdistConnectorRow.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/GpfdistConnectorRow.java new file mode 100644 index 000000000000..6430ce72cd46 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/GpfdistConnectorRow.java @@ -0,0 +1,33 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +import io.trino.plugin.adb.connector.decode.ColumnValue; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; + +public record GpfdistConnectorRow(long estimatedSize, ColumnValue[] columnValues) + implements ConnectorRow +{ + @Override + public ColumnValue[] getColumnValues() + { + return columnValues; + } + + @Override + public long getEstimatedSize() + { + return estimatedSize; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessor.java new file mode 100644 index 000000000000..02ddfd8a8cb6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessor.java @@ -0,0 +1,23 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +import java.io.InputStream; + +public interface InputDataProcessor +{ + void process(InputStream dataStream); + + ProcessingDataResult getProcessedDataResult(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessorFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessorFactory.java new file mode 100644 index 000000000000..b075ad036c99 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/InputDataProcessorFactory.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +import io.trino.plugin.adb.connector.decode.RowDecoder; + +public interface InputDataProcessorFactory +{ + InputDataProcessor create(RowDecoder rowDecoder, RowProcessingService rowProcessingService); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/ProcessingDataResult.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/ProcessingDataResult.java new file mode 100644 index 000000000000..f183d0ad4667 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/ProcessingDataResult.java @@ -0,0 +1,20 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +import java.util.concurrent.atomic.AtomicLong; + +public record ProcessingDataResult(AtomicLong rowCount, AtomicLong estimatedProcessedBytes) +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/RowProcessingService.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/RowProcessingService.java new file mode 100644 index 000000000000..69e2e927c0c8 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/RowProcessingService.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; + +public interface RowProcessingService +{ + void put(ConnectorRow row); + + ConnectorRow take(); + + boolean isEmpty(); + + void stop(); + + void clear(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/SegmentRequestStatus.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/SegmentRequestStatus.java new file mode 100644 index 000000000000..7c2f97252a5a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/SegmentRequestStatus.java @@ -0,0 +1,22 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +public enum SegmentRequestStatus +{ + INITIALIZED, + PROCESSING, + ERROR, + FINISHED; +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/WithEstimatedSize.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/WithEstimatedSize.java new file mode 100644 index 000000000000..b16445265e4a --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/WithEstimatedSize.java @@ -0,0 +1,19 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload; + +public interface WithEstimatedSize +{ + long getEstimatedSize(); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContext.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContext.java new file mode 100644 index 000000000000..5c4bf060bf71 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContext.java @@ -0,0 +1,97 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.context; + +import io.airlift.log.Logger; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.protocol.gpfdist.Context; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistBufferedRowProcessingService; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistSegmentRequestProcessor; + +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public class ReadContext + implements Context +{ + private static final Logger log = Logger.get(ReadContext.class); + private final ContextId id; + private final GpfdistUnloadMetadata metadata; + private final RowDecoder rowDecoder; + private final GpfdistBufferedRowProcessingService rowProcessingService; + private final AtomicLong completedBytes = new AtomicLong(); + + public ReadContext(GpfdistUnloadMetadata metadata, + RowDecoder rowDecoder, + GpfdistBufferedRowProcessingService rowProcessingService) + { + this(new ContextId(metadata.getSourceTable()), metadata, rowDecoder, rowProcessingService); + } + + public ReadContext(ContextId id, + GpfdistUnloadMetadata metadata, + RowDecoder rowDecoder, + GpfdistBufferedRowProcessingService rowProcessingService) + { + this.id = id; + this.metadata = metadata; + this.rowDecoder = rowDecoder; + this.rowProcessingService = rowProcessingService; + } + + @Override + public ContextId getId() + { + return id; + } + + public GpfdistUnloadMetadata getMetadata() + { + return metadata; + } + + public Map getSegmentDataProcessors() + { + return rowProcessingService.getSegmentDataProcessors(); + } + + public RowDecoder getRowDecoder() + { + return rowDecoder; + } + + public AtomicLong getCompletedBytes() + { + return completedBytes; + } + + public AtomicLong getMemoryUsage() + { + return rowProcessingService.getUsedBufferedRowsMemory(); + } + + public GpfdistBufferedRowProcessingService getRowProcessingService() + { + return rowProcessingService; + } + + @Override + public void close() + { + rowProcessingService.clear(); + log.debug("Closed read context %s", id); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContextManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContextManager.java new file mode 100644 index 000000000000..b483ed0cc504 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContextManager.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.context; + +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractContextManager; + +public class ReadContextManager + extends AbstractContextManager +{ +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistBufferedRowProcessingService.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistBufferedRowProcessingService.java new file mode 100644 index 000000000000..15a45daa2df6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistBufferedRowProcessingService.java @@ -0,0 +1,171 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.RowProcessingService; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.SegmentRequestStatus; + +import java.util.LinkedList; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; + +public class GpfdistBufferedRowProcessingService + implements RowProcessingService +{ + private final Queue rowsQueue = new LinkedList<>(); + private final AtomicLong usedBufferedRowsMemory = new AtomicLong(0); + private final AtomicBoolean isStopped = new AtomicBoolean(false); + private final Map segmentDataProcessors = new ConcurrentHashMap<>(); + private final Condition isFullCondition; + private final Condition isReadyForTransferCondition; + private final Lock lock; + private final long maxByteBufferSize; + private long currentBufferSize; + + public GpfdistBufferedRowProcessingService(AdbPluginConfig pluginConfig) + { + this.maxByteBufferSize = pluginConfig.getReadBufferSize().toBytes(); + lock = new ReentrantLock(); + isFullCondition = lock.newCondition(); + isReadyForTransferCondition = lock.newCondition(); + } + + @Override + public void put(ConnectorRow row) + { + lock.lock(); + try { + while (isFull()) { + try { + isFullCondition.await(); + } + catch (InterruptedException ex) { + throw new RuntimeException(ex); + } + } + rowsQueue.add(row); + currentBufferSize += row.getEstimatedSize(); + usedBufferedRowsMemory.addAndGet(row.getEstimatedSize()); + isReadyForTransferCondition.signalAll(); + } + finally { + lock.unlock(); + } + } + + @Override + public ConnectorRow take() + { + ConnectorRow element; + lock.lock(); + try { + while (isNotReady()) { + try { + isReadyForTransferCondition.await(); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + element = rowsQueue.poll(); + if (element != null) { + currentBufferSize -= element.getEstimatedSize(); + } + isFullCondition.signalAll(); + } + finally { + lock.unlock(); + } + return element; + } + + @Override + public boolean isEmpty() + { + return rowsQueue.isEmpty(); + } + + @Override + public void stop() + { + lock.lock(); + try { + isStopped.set(true); + isReadyForTransferCondition.signalAll(); + } + finally { + lock.unlock(); + } + } + + @Override + public void clear() + { + lock.lock(); + try { + segmentDataProcessors.clear(); + rowsQueue.clear(); + } + finally { + lock.unlock(); + } + } + + private boolean isNotReady() + { + return isDataTransferNotInitialized() || isNotAllDataProcessed(); + } + + private boolean isDataTransferNotInitialized() + { + //checked that insert adb query is not finished and there are no segment processors yet + return !isStopped.get() && segmentDataProcessors.isEmpty(); + } + + private boolean isNotAllDataProcessed() + { + //check that there are no rows in queue and segment processors are not finished yet + return isEmpty() + && segmentDataProcessors.values().stream() + .anyMatch(req -> req.getStatus() != SegmentRequestStatus.FINISHED); + } + + private boolean isFull() + { + return currentBufferSize >= maxByteBufferSize; + } + + public AtomicLong getUsedBufferedRowsMemory() + { + return usedBufferedRowsMemory; + } + + public AtomicBoolean getIsStopped() + { + return isStopped; + } + + public Map getSegmentDataProcessors() + { + return segmentDataProcessors; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistCsvDataProcessor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistCsvDataProcessor.java new file mode 100644 index 000000000000..845247dd1fa6 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistCsvDataProcessor.java @@ -0,0 +1,93 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import com.opencsv.CSVReader; +import com.opencsv.CSVReaderBuilder; +import com.opencsv.RFC4180ParserBuilder; +import com.opencsv.enums.CSVReaderNullFieldIndicator; +import com.opencsv.exceptions.CsvValidationException; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.encode.DataFormatConfig; +import io.trino.plugin.adb.connector.encode.csv.CsvFormatConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.ProcessingDataResult; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.RowProcessingService; +import io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; + +import java.io.BufferedReader; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.Charset; +import java.util.concurrent.atomic.AtomicLong; + +public class GpfdistCsvDataProcessor + implements InputDataProcessor +{ + private final CsvFormatConfig dataFormatConfig; + private final RowDecoder rowDecoder; + private final RowProcessingService rowProcessingService; + private final ProcessingDataResult processingDataResult; + + public GpfdistCsvDataProcessor(DataFormatConfig dataFormatConfig, + RowDecoder rowDecoder, + RowProcessingService rowProcessingService) + { + this.dataFormatConfig = (CsvFormatConfig) dataFormatConfig; + this.rowDecoder = rowDecoder; + this.rowProcessingService = rowProcessingService; + processingDataResult = new ProcessingDataResult(new AtomicLong(0), new AtomicLong(0)); + } + + @Override + public void process(InputStream dataStream) + { + try { + try (InputStreamReader streamReader = new InputStreamReader(dataStream, + Charset.forName(dataFormatConfig.getEncoding())); + BufferedReader bufferedStreamReader = new BufferedReader(streamReader); + CSVReader csvReader = new CSVReaderBuilder(bufferedStreamReader) + .withCSVParser(new RFC4180ParserBuilder() + .withSeparator(dataFormatConfig.getDelimiter()) + .withFieldAsNull(CSVReaderNullFieldIndicator.EMPTY_SEPARATORS) + .build()) + .build()) { + String[] rawRow; + try { + while ((rawRow = csvReader.readNext()) != null) { + ConnectorRow row = rowDecoder.decode(rawRow); + rowProcessingService.put(row); + processingDataResult.rowCount().incrementAndGet(); + processingDataResult.estimatedProcessedBytes().addAndGet(row.getEstimatedSize()); + } + } + catch (CsvValidationException e) { + throw new RuntimeException(e); + } + } + } + catch (Throwable e) { + throw new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, + "Failed to process input data: " + e.getMessage(), e); + } + } + + @Override + public ProcessingDataResult getProcessedDataResult() + { + return processingDataResult; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistInputDataProcessorFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistInputDataProcessorFactory.java new file mode 100644 index 000000000000..81ef36e59c0d --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistInputDataProcessorFactory.java @@ -0,0 +1,45 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import com.google.inject.Inject; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.encode.DataFormat; +import io.trino.plugin.adb.connector.encode.DataFormatConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessorFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.RowProcessingService; + +public class GpfdistInputDataProcessorFactory + implements InputDataProcessorFactory +{ + private final DataFormatConfig dataFormatConfig; + + @Inject + public GpfdistInputDataProcessorFactory(DataFormatConfig dataFormatConfig) + { + this.dataFormatConfig = dataFormatConfig; + } + + @Override + public InputDataProcessor create(RowDecoder rowDecoder, RowProcessingService rowProcessingService) + { + if (dataFormatConfig.getDataFormat() == DataFormat.CSV) { + return new GpfdistCsvDataProcessor(dataFormatConfig, rowDecoder, rowProcessingService); + } + else { + throw new UnsupportedOperationException("Unsupported data format: " + dataFormatConfig.getDataFormat()); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordCursor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordCursor.java new file mode 100644 index 000000000000..aea7cf529c76 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordCursor.java @@ -0,0 +1,207 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import io.airlift.log.Logger; +import io.airlift.slice.Slice; +import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow; +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.DataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContext; +import io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.RecordCursor; +import io.trino.spi.type.Type; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicLong; + +import static com.google.common.base.Preconditions.checkArgument; + +public class GpfdistRecordCursor + implements RecordCursor +{ + private static final Logger log = Logger.get(GpfdistRecordCursor.class); + private final AtomicLong readTimeNanos = new AtomicLong(); + private final ReadContext readContext; + private final ContextManager contextManager; + private final DataTransferQueryExecutor dataTransferQueryExecutor; + private final List columnTypes; + private final GpfdistBufferedRowProcessingService rowProcessingService; + private boolean isClosed; + private boolean isQueryStarted; + private ConnectorRow currentRow; + private long unloadedRows; + private CompletableFuture dataTransferQueryFuture; + private Throwable queryExecutionException; + + public GpfdistRecordCursor(ContextManager contextManager, + ReadContext readContext, + DataTransferQueryExecutor dataTransferQueryExecutor, + List columnTypes) + { + this.readContext = readContext; + this.contextManager = contextManager; + this.dataTransferQueryExecutor = dataTransferQueryExecutor; + this.columnTypes = columnTypes; + this.rowProcessingService = readContext.getRowProcessingService(); + } + + @Override + public boolean advanceNextPosition() + { + if (isClosed) { + return false; + } + else { + long startNanoSec = System.nanoTime(); + try { + if (!isQueryStarted) { + dataTransferQueryFuture = executeTransferDataQuery(); + isQueryStarted = true; + } + if (queryExecutionException != null) { + throw queryExecutionException; + } + currentRow = rowProcessingService.take(); + if (currentRow != null) { + unloadedRows++; + readContext.getCompletedBytes().addAndGet(currentRow.getEstimatedSize()); + readContext.getMemoryUsage().addAndGet(-currentRow.getEstimatedSize()); + return true; + } + else { + log.info("Data processing is finished. Unloaded rows %s. Processing result: %s", + unloadedRows, + readContext.getSegmentDataProcessors().values()); + closeCtx(); + return false; + } + } + catch (Throwable e) { + if (!readContext.getRowProcessingService().getIsStopped().get()) { + dataTransferQueryFuture.completeExceptionally(e); + } + log.warn("Processed rows to target: %d. Unloading from adb result: %s", + unloadedRows, + readContext.getSegmentDataProcessors().values()); + closeCtx(); + throw new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, + "Failed to unload data from adb: " + e.getMessage(), e); + } + finally { + readTimeNanos.addAndGet(System.nanoTime() - startNanoSec); + } + } + } + + private CompletableFuture executeTransferDataQuery() + { + return dataTransferQueryExecutor.execute().whenComplete((_, error) -> { + readContext.getRowProcessingService().stop(); + if (error != null) { + queryExecutionException = error; + throw new CompletionException(error); + } + }); + } + + @Override + public long getCompletedBytes() + { + return readContext.getCompletedBytes().get(); + } + + @Override + public long getReadTimeNanos() + { + return readTimeNanos.get(); + } + + @Override + public long getMemoryUsage() + { + return readContext.getMemoryUsage().get(); + } + + @Override + public Type getType(int field) + { + checkArgument(field < columnTypes.size()); + return columnTypes.get(field); + } + + @Override + public boolean getBoolean(int field) + { + return get(field); + } + + @Override + public long getLong(int field) + { + return get(field); + } + + @Override + public double getDouble(int field) + { + return get(field); + } + + @Override + public Slice getSlice(int field) + { + return get(field); + } + + @Override + public Object getObject(int field) + { + return get(field); + } + + @Override + public boolean isNull(int field) + { + return get(field) == null; + } + + @SuppressWarnings("unchecked") + private T get(int field) + { + return (T) currentRow.getColumnValues()[field].value(); + } + + @Override + public void close() + { + if (!this.isClosed) { + this.isClosed = true; + closeCtx(); + } + } + + private void closeCtx() + { + try { + readContext.close(); + } + finally { + contextManager.remove(readContext.getId()); + } + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSet.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSet.java new file mode 100644 index 000000000000..0246638be158 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSet.java @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.DataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContext; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.spi.connector.RecordCursor; +import io.trino.spi.connector.RecordSet; +import io.trino.spi.type.Type; + +import java.util.List; + +public class GpfdistRecordSet + implements RecordSet +{ + private final ContextManager contextManager; + private final ReadContext readContext; + private final DataTransferQueryExecutor dataTransferQueryExecutor; + private final List columnTypes; + + public GpfdistRecordSet(ContextManager contextManager, + ReadContext readContext, + DataTransferQueryExecutor dataTransferQueryExecutor) + { + this.contextManager = contextManager; + this.readContext = readContext; + this.dataTransferQueryExecutor = dataTransferQueryExecutor; + columnTypes = readContext.getMetadata().getColumnHandles().stream() + .map(columnHandle -> ((JdbcColumnHandle) columnHandle).getColumnType()) + .toList(); + } + + @Override + public List getColumnTypes() + { + return columnTypes; + } + + @Override + public RecordCursor cursor() + { + return new GpfdistRecordCursor(contextManager, readContext, dataTransferQueryExecutor, columnTypes); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSetProvider.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSetProvider.java new file mode 100644 index 000000000000..f74693b3d87c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordSetProvider.java @@ -0,0 +1,141 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import com.google.inject.Inject; +import io.trino.plugin.adb.AdbPluginConfig; +import io.trino.plugin.adb.connector.AdbJdbcSplit; +import io.trino.plugin.adb.connector.AdbSessionProperties; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.decode.RowDecoder; +import io.trino.plugin.adb.connector.decode.RowDecoderFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.ContextManager; +import io.trino.plugin.adb.connector.protocol.gpfdist.CreateExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.DataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.ExecutorServiceProvider; +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableFormatConfig; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableFormatConfigFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadataFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.context.ReadContext; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.query.GpfdistUnloadDataTransferQueryExecutor; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.ForRecordCursor; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcRecordSetProvider; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorRecordSetProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplit; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.RecordSet; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static com.google.common.base.Preconditions.checkArgument; + +public class GpfdistRecordSetProvider + implements ConnectorRecordSetProvider +{ + private static final ExternalTableType EXTERNAL_TABLE_TYPE = ExternalTableType.WRITABLE; + private final JdbcRecordSetProvider regularProvider; + private final AdbSqlClient client; + private final GpfdistUnloadMetadataFactory unloadMetadataFactory; + private final ExternalTableFormatConfigFactory externalTableFormatConfigFactory; + private final RowDecoderFactory rowDecoderFactory; + private final ExecutorService unloadQueryThreadExecutor; + private final CreateExternalTableQueryFactory externalTableQueryFactory; + private final ContextManager contextManager; + private final AdbPluginConfig pluginConfig; + private final InsertDataQueryFactory insertDataQueryFactory; + + @Inject + public GpfdistRecordSetProvider(@ForBaseJdbc JdbcClient client, + GpfdistUnloadMetadataFactory unloadMetadataFactory, + @ForRecordCursor ExecutorService recordSetExecutor, + ExternalTableFormatConfigFactory externalTableFormatConfigFactory, + RowDecoderFactory rowDecoderFactory, + Set createExternalTableQueryFactories, + ContextManager contextManager, + AdbPluginConfig pluginConfig, + Set insertDataQueryFactories) + { + this.regularProvider = new JdbcRecordSetProvider(client, recordSetExecutor); + this.client = (AdbSqlClient) client; + this.unloadMetadataFactory = unloadMetadataFactory; + this.externalTableFormatConfigFactory = externalTableFormatConfigFactory; + this.rowDecoderFactory = rowDecoderFactory; + this.contextManager = contextManager; + this.unloadQueryThreadExecutor = ExecutorServiceProvider.LOAD_DATA_QUERY_EXECUTOR_SERVICE; + Map externalTableQueryFactoryMap = + createExternalTableQueryFactories.stream() + .collect(Collectors.toMap(CreateExternalTableQueryFactory::getExternalTableType, + Function.identity())); + externalTableQueryFactory = externalTableQueryFactoryMap.get(EXTERNAL_TABLE_TYPE); + checkArgument(externalTableQueryFactory != null, + "failed to get writable table query factory by externalTableType %s", + EXTERNAL_TABLE_TYPE); + Map insertDataFactoryMap = insertDataQueryFactories.stream() + .collect(Collectors.toMap(InsertDataQueryFactory::getExternalTableType, Function.identity())); + insertDataQueryFactory = insertDataFactoryMap.get(EXTERNAL_TABLE_TYPE); + checkArgument(insertDataQueryFactory != null, + "failed to get insert data query factory by externalTableType %s", + EXTERNAL_TABLE_TYPE); + this.pluginConfig = pluginConfig; + } + + @Override + public RecordSet getRecordSet(ConnectorTransactionHandle transaction, + ConnectorSession session, + ConnectorSplit split, + ConnectorTableHandle table, + List columns) + { + if (split instanceof AdbJdbcSplit) { + ExternalTableFormatConfig externalTableFormatConfig = externalTableFormatConfigFactory.create(); + GpfdistUnloadMetadata gpfdistUnloadMetadata = unloadMetadataFactory.create(session, + (JdbcTableHandle) table, + (AdbJdbcSplit) split, + columns, + externalTableFormatConfig); + RowDecoder rowDecoder = rowDecoderFactory.create(session, gpfdistUnloadMetadata.getDataTypes()); + GpfdistBufferedRowProcessingService rowProcessingService = + new GpfdistBufferedRowProcessingService(pluginConfig); + ReadContext readContext = new ReadContext( + gpfdistUnloadMetadata, + rowDecoder, + rowProcessingService); + DataTransferQueryExecutor dataTransferQueryExecutor = new GpfdistUnloadDataTransferQueryExecutor( + client, + session, + unloadQueryThreadExecutor, + gpfdistUnloadMetadata, + AdbSessionProperties.getGpfdistRetryTimeout(session), + externalTableQueryFactory, + insertDataQueryFactory); + contextManager.add(readContext); + return new GpfdistRecordSet(contextManager, readContext, dataTransferQueryExecutor); + } + return regularProvider.getRecordSet(transaction, session, split, table, columns); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistSegmentRequestProcessor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistSegmentRequestProcessor.java new file mode 100644 index 000000000000..4c905ac1fa3c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistSegmentRequestProcessor.java @@ -0,0 +1,69 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.process; + +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.InputDataProcessor; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.ProcessingDataResult; +import io.trino.plugin.adb.connector.protocol.gpfdist.unload.SegmentRequestStatus; + +import java.io.InputStream; + +import static java.lang.String.format; + +public class GpfdistSegmentRequestProcessor +{ + private final InputDataProcessor delegate; + private final int segmentId; + private final ProcessingDataResult processingDataResult; + private SegmentRequestStatus status; + + public GpfdistSegmentRequestProcessor(int segmentId, InputDataProcessor delegate) + { + this.delegate = delegate; + this.segmentId = segmentId; + this.processingDataResult = delegate.getProcessedDataResult(); + status = SegmentRequestStatus.INITIALIZED; + } + + public void process(InputStream dataStream) + { + try { + if (status != SegmentRequestStatus.PROCESSING) { + status = SegmentRequestStatus.PROCESSING; + } + delegate.process(dataStream); + } + catch (Exception e) { + status = SegmentRequestStatus.ERROR; + throw new RuntimeException(format("Failed to process data: %s. error: %s", this, e.getMessage()), e); + } + } + + public void stop() + { + status = SegmentRequestStatus.FINISHED; + } + + public SegmentRequestStatus getStatus() + { + return status; + } + + @Override + public String toString() + { + return "Data processing for segmentId: " + segmentId + "; status: " + + status + "; result: " + processingDataResult; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/CreateWritableExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/CreateWritableExternalTableQueryFactory.java new file mode 100644 index 000000000000..adfb201d7b5e --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/CreateWritableExternalTableQueryFactory.java @@ -0,0 +1,56 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.query; + +import io.trino.plugin.adb.connector.AdbJdbcSplit; +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata; + +import java.util.HashSet; +import java.util.List; +import java.util.stream.Collectors; + +public class CreateWritableExternalTableQueryFactory + extends AbstractExternalTableQueryFactory +{ + @Override + public String createQuery(GpfdistMetadata metadata) + { + GpfdistUnloadMetadata unloadMetadata = (GpfdistUnloadMetadata) metadata; + String externalTableDistribution = createExternalTableDistribution((AdbJdbcSplit) unloadMetadata.getSplit(), + unloadMetadata.getColumnNames()); + return createCommonQuery(metadata) + " " + externalTableDistribution; + } + + private String createExternalTableDistribution(AdbJdbcSplit split, List columns) + { + List distribution = split.getDistribution(); + if (!distribution.isEmpty() && new HashSet<>(columns).containsAll(distribution)) { + return distribution.stream() + .map(columnName -> "\"" + columnName + "\"") + .collect(Collectors.joining(", ", "DISTRIBUTED BY (", ")")); + } + else { + return "DISTRIBUTED RANDOMLY"; + } + } + + @Override + public ExternalTableType getExternalTableType() + { + return ExternalTableType.WRITABLE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/GpfdistUnloadDataTransferQueryExecutor.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/GpfdistUnloadDataTransferQueryExecutor.java new file mode 100644 index 000000000000..0265f3c9d7cb --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/GpfdistUnloadDataTransferQueryExecutor.java @@ -0,0 +1,85 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.query; + +import io.airlift.log.Logger; +import io.airlift.units.Duration; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.protocol.gpfdist.AbstractDataTransferQueryExecutor; +import io.trino.plugin.adb.connector.protocol.gpfdist.CreateExternalTableQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.SQLException; +import java.util.Optional; +import java.util.concurrent.ExecutorService; + +public class GpfdistUnloadDataTransferQueryExecutor + extends AbstractDataTransferQueryExecutor +{ + private static final Logger log = Logger.get(GpfdistUnloadDataTransferQueryExecutor.class); + private final GpfdistUnloadMetadata unloadMetadata; + private final Optional gpfdistRetryTimeout; + + public GpfdistUnloadDataTransferQueryExecutor(AdbSqlClient client, + ConnectorSession session, + ExecutorService executor, + GpfdistUnloadMetadata unloadMetadata, + Optional gpfdistRetryTimeout, + CreateExternalTableQueryFactory externalTableQueryFactory, + InsertDataQueryFactory insertDataQueryFactory) + { + super(client, session, executor, externalTableQueryFactory, insertDataQueryFactory); + this.unloadMetadata = unloadMetadata; + this.gpfdistRetryTimeout = gpfdistRetryTimeout; + } + + @Override + protected void executeQueries() + throws SQLException + { + createReadableExternalTable(); + setGpfdistRetryTimeoutIfNeeded(); + insertIntoExternalTable(); + } + + private void createReadableExternalTable() + throws SQLException + { + String sql = externalTableQueryFactory.createQuery(unloadMetadata); + client.execute(session, connection, sql); + log.info("Executed create writable external table query: %s", sql); + } + + private void setGpfdistRetryTimeoutIfNeeded() + throws SQLException + { + if (gpfdistRetryTimeout.isPresent()) { + long retryInSeconds = gpfdistRetryTimeout.get().toJavaTime().toSeconds(); + String query = String.format("SET gpfdist_retry_timeout TO %d", retryInSeconds); + client.execute(session, connection, query); + log.info("Executed set gpfdist retry timeout query: %s", query); + } + } + + private void insertIntoExternalTable() + throws SQLException + { + PreparedQuery query = insertDataQueryFactory.create(session, connection, unloadMetadata); + client.executeAsPreparedStatement(session, connection, query); + log.info("Executed insert into writable external table from source table query: %s", query.query()); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/InsertDataToExternalTableQueryFactory.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/InsertDataToExternalTableQueryFactory.java new file mode 100644 index 000000000000..7b5b7c283724 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/query/InsertDataToExternalTableQueryFactory.java @@ -0,0 +1,74 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.protocol.gpfdist.unload.query; + +import com.google.common.collect.ImmutableMap; +import com.google.inject.Inject; +import io.trino.plugin.adb.connector.AdbSqlClient; +import io.trino.plugin.adb.connector.protocol.gpfdist.InsertDataQueryFactory; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ExternalTableType; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistMetadata; +import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata; +import io.trino.plugin.jdbc.ForBaseJdbc; +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.PreparedQuery; +import io.trino.spi.connector.ConnectorSession; + +import java.sql.Connection; +import java.util.List; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Collectors; + +public class InsertDataToExternalTableQueryFactory + implements InsertDataQueryFactory +{ + private final AdbSqlClient client; + + @Inject + public InsertDataToExternalTableQueryFactory(@ForBaseJdbc JdbcClient client) + { + this.client = (AdbSqlClient) client; + } + + @Override + public PreparedQuery create(ConnectorSession session, Connection connection, GpfdistMetadata metadata) + { + GpfdistUnloadMetadata unloadMetadata = (GpfdistUnloadMetadata) metadata; + PreparedQuery preparedSelectQuery = client.prepareQuery(session, + connection, + unloadMetadata.getTargetTableHandle(), + Optional.empty(), + (List) unloadMetadata.getColumnHandles(), + ImmutableMap.of(), + Optional.of(unloadMetadata.getSplit())); + Function insertTransformer = query -> unloadMetadata.getColumnHandles().isEmpty() + ? String.format("INSERT INTO %s %s", unloadMetadata.getSourceTable(), query) + : String.format( + "INSERT INTO %s (%s) %s", + unloadMetadata.getSourceTable(), + unloadMetadata.getColumnHandles().stream() + .map(column -> client.quoted(((JdbcColumnHandle) column).getColumnName())) + .collect(Collectors.joining(", ")), + query); + return preparedSelectQuery.transformQuery(insertTransformer); + } + + @Override + public ExternalTableType getExternalTableType() + { + return ExternalTableType.WRITABLE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbCreateTableStorageConfig.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbCreateTableStorageConfig.java new file mode 100644 index 000000000000..6a395708184f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbCreateTableStorageConfig.java @@ -0,0 +1,111 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import io.airlift.configuration.Config; + +public class AdbCreateTableStorageConfig +{ + private Boolean appendOptimized; + private Integer blockSize; + private AdbTableStorageOrientation orientation; + private Boolean checksum; + private AdbTableStorageCompressType compressType; + private Integer compressLevel; + private Integer fillFactor; + + public Boolean getAppendOptimized() + { + return this.appendOptimized; + } + + @Config("adb.create-table.appendoptimized") + public AdbCreateTableStorageConfig setAppendOptimized(Boolean appendOptimized) + { + this.appendOptimized = appendOptimized; + return this; + } + + public Integer getBlockSize() + { + return this.blockSize; + } + + @Config("adb.create-table.blocksize") + public AdbCreateTableStorageConfig setBlockSize(Integer blockSize) + { + this.blockSize = blockSize; + return this; + } + + public AdbTableStorageOrientation getOrientation() + { + return this.orientation; + } + + @Config("adb.create-table.orientation") + public AdbCreateTableStorageConfig setOrientation(AdbTableStorageOrientation orientation) + { + this.orientation = orientation; + return this; + } + + public Boolean getChecksum() + { + return this.checksum; + } + + @Config("adb.create-table.checksum") + public AdbCreateTableStorageConfig setChecksum(Boolean checksum) + { + this.checksum = checksum; + return this; + } + + public AdbTableStorageCompressType getCompressType() + { + return this.compressType; + } + + @Config("adb.create-table.compresstype") + public AdbCreateTableStorageConfig setCompressType(AdbTableStorageCompressType compressType) + { + this.compressType = compressType; + return this; + } + + public Integer getCompressLevel() + { + return this.compressLevel; + } + + @Config("adb.create-table.compresslevel") + public AdbCreateTableStorageConfig setCompressLevel(Integer compressLevel) + { + this.compressLevel = compressLevel; + return this; + } + + public Integer getFillFactor() + { + return this.fillFactor; + } + + @Config("adb.create-table.fillfactor") + public AdbCreateTableStorageConfig setFillFactor(Integer fillFactor) + { + this.fillFactor = fillFactor; + return this; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableDistributed.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableDistributed.java new file mode 100644 index 000000000000..86ad7ab90ce3 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableDistributed.java @@ -0,0 +1,20 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +public enum AdbTableDistributed +{ + RANDOMLY, + REPLICATED; +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableProperties.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableProperties.java new file mode 100644 index 000000000000..5c2d970eaf59 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableProperties.java @@ -0,0 +1,128 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.VarcharType; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +public class AdbTableProperties + implements TablePropertiesProvider +{ + public static final String DISTRIBUTED_PROPERTY = "distributed"; + public static final String DISTRIBUTED_BY_PROPERTY = "distributed_by"; + public static final String APPEND_OPTIMIZED_PROPERTY = "appendoptimized"; + public static final String BLOCK_SIZE_PROPERTY = "blocksize"; + public static final String ORIENTATION_PROPERTY = "orientation"; + public static final String CHECKSUM_PROPERTY = "checksum"; + public static final String COMPRESS_TYPE_PROPERTY = "compresstype"; + public static final String COMPRESS_LEVEL_PROPERTY = "compresslevel"; + public static final String FILL_FACTOR_PROPERTY = "fillfactor"; + private final List> tableProperties; + + @Inject + public AdbTableProperties(AdbCreateTableStorageConfig config) + { + ImmutableList.Builder> propertiesBuilder = ImmutableList.builder(); + this.tableProperties = propertiesBuilder + .add(PropertyMetadata.enumProperty( + DISTRIBUTED_PROPERTY, + "Distribution policy. Valid values: randomly, replicated. For hash distribution use distributed_by property", + AdbTableDistributed.class, + null, + false)) + .add(new PropertyMetadata( + DISTRIBUTED_BY_PROPERTY, + "Columns for hash distribution", + new ArrayType(VarcharType.VARCHAR), + List.class, + null, + false, + value -> value, + value -> value)) + .add(PropertyMetadata.booleanProperty(APPEND_OPTIMIZED_PROPERTY, "Whether table is append-optimized", config.getAppendOptimized(), false)) + .add(PropertyMetadata.integerProperty(BLOCK_SIZE_PROPERTY, "Block size for append-optimized tables", config.getBlockSize(), false)) + .add(PropertyMetadata.enumProperty( + ORIENTATION_PROPERTY, "Table orientation. Valid values: column, row", AdbTableStorageOrientation.class, config.getOrientation(), false)) + .add(PropertyMetadata.booleanProperty(CHECKSUM_PROPERTY, "Whether table is append-optimized", config.getChecksum(), false)) + .add(PropertyMetadata.enumProperty( + COMPRESS_TYPE_PROPERTY, + "Compression type. Valid values: zlib, zstd, rle_type, none", + AdbTableStorageCompressType.class, + config.getCompressType(), + false)) + .add(PropertyMetadata.integerProperty(COMPRESS_LEVEL_PROPERTY, "Compression level", config.getCompressLevel(), false)) + .add(PropertyMetadata.integerProperty(FILL_FACTOR_PROPERTY, "Fill factor", config.getFillFactor(), false)) + .build(); + } + + @Override + public List> getTableProperties() + { + return this.tableProperties; + } + + public static Optional getDistributed(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(DISTRIBUTED_PROPERTY)).map(AdbTableDistributed.class::cast); + } + + public static Optional> getDistributedBy(Map tableProperties) + { + List value = (List) tableProperties.get(DISTRIBUTED_BY_PROPERTY); + return value == null ? Optional.empty() : Optional.of(ImmutableList.copyOf(value)); + } + + public static Optional getAppendOptimized(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(APPEND_OPTIMIZED_PROPERTY)).map(Boolean.class::cast); + } + + public static Optional getBlockSize(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(BLOCK_SIZE_PROPERTY)).map(Integer.class::cast); + } + + public static Optional getOrientation(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(ORIENTATION_PROPERTY)).map(AdbTableStorageOrientation.class::cast); + } + + public static Optional getChecksum(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(CHECKSUM_PROPERTY)).map(Boolean.class::cast); + } + + public static Optional getCompressType(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(COMPRESS_TYPE_PROPERTY)).map(AdbTableStorageCompressType.class::cast); + } + + public static Optional getCompressLevel(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(COMPRESS_LEVEL_PROPERTY)).map(Integer.class::cast); + } + + public static Optional getFillFactor(Map tableProperties) + { + return Optional.ofNullable(tableProperties.get(FILL_FACTOR_PROPERTY)).map(Integer.class::cast); + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageCompressType.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageCompressType.java new file mode 100644 index 000000000000..10a60fac3104 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageCompressType.java @@ -0,0 +1,22 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +public enum AdbTableStorageCompressType +{ + ZLIB, + ZSTD, + RLE_TYPE, + NONE; +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageOrientation.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageOrientation.java new file mode 100644 index 000000000000..d033e2a68555 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/AdbTableStorageOrientation.java @@ -0,0 +1,20 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +public enum AdbTableStorageOrientation +{ + COLUMN, + ROW +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManager.java new file mode 100644 index 000000000000..d176b1500d31 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManager.java @@ -0,0 +1,23 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitSource; + +public interface SplitSourceManager +{ + ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle tableHandle); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManagerImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManagerImpl.java new file mode 100644 index 000000000000..a67f4135d88c --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/SplitSourceManagerImpl.java @@ -0,0 +1,132 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import com.google.inject.Inject; +import io.trino.plugin.adb.connector.AdbJdbcSplit; +import io.trino.plugin.adb.connector.AdbSessionProperties; +import io.trino.plugin.adb.connector.metadata.AdbMetadataDao; +import io.trino.plugin.base.mapping.IdentifierMapping; +import io.trino.plugin.jdbc.JdbcSplit; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitSource; +import io.trino.spi.connector.FixedSplitSource; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.IntStream; + +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.plugin.adb.AdbPluginConfig.IDENTIFIER_QUOTE; +import static java.lang.Math.min; +import static java.lang.String.format; + +public class SplitSourceManagerImpl + implements SplitSourceManager +{ + private static final int MINIMUM_PARALLELISM_VALUE = 1; + private final AdbMetadataDao metadata; + private final IdentifierMapping identifierMapping; + + @Inject + public SplitSourceManagerImpl(AdbMetadataDao metadata, IdentifierMapping identifierMapping) + { + this.metadata = metadata; + this.identifierMapping = identifierMapping; + } + + @Override + public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle tableHandle) + { + int parallelism = getSplitParallelism(session, tableHandle); + return new FixedSplitSource(createAdbSplits(session, tableHandle, segmentedSplits(parallelism))); + } + + private int getSplitParallelism(ConnectorSession session, JdbcTableHandle tableHandle) + { + if (!tableHandle.isNamedRelation()) { + return MINIMUM_PARALLELISM_VALUE; + } + int maxParallelism = AdbSessionProperties.getMaxScanParallelism(session); + if (maxParallelism <= MINIMUM_PARALLELISM_VALUE) { + return MINIMUM_PARALLELISM_VALUE; + } + boolean segmented = this.metadata.isSegmentedTable(session, + quoted(tableHandle.getRequiredNamedRelation().getRemoteTableName())); + if (!segmented) { + return MINIMUM_PARALLELISM_VALUE; + } + else { + int segmentCount = this.metadata.getSegmentCount(session); + if (segmentCount <= 0) { + segmentCount = MINIMUM_PARALLELISM_VALUE; + } + return min(segmentCount, maxParallelism); + } + } + + private List segmentedSplits(int parallelism) + { + return parallelism == MINIMUM_PARALLELISM_VALUE + ? List.of(new JdbcSplit(Optional.empty())) + : IntStream.range(0, parallelism) + .boxed() + .map(id -> new JdbcSplit(Optional.of(format("gp_segment_id %% %s = %s", parallelism, id)))) + .collect(toImmutableList()); + } + + private List createAdbSplits(ConnectorSession session, + JdbcTableHandle tableHandle, + List splits) + { + String objectName = quoted(tableHandle.getRequiredNamedRelation().getRemoteTableName()); + Map tableProperties = this.metadata.getTableProperties(session, objectName, identifierMapping); + List distributionInfo = AdbTableProperties.getDistributedBy(tableProperties).orElse(List.of()); + return splits.stream() + .map(split -> new AdbJdbcSplit(distributionInfo, split.getAdditionalPredicate(), + split.getDynamicFilter())) + .toList(); + } + + public String quoted(RemoteTableName remoteTableName) + { + return quoted( + remoteTableName.getCatalogName().orElse(null), + remoteTableName.getSchemaName().orElse(null), + remoteTableName.getTableName()); + } + + private String quoted(String catalog, String schema, String table) + { + StringBuilder sb = new StringBuilder(); + if (!isNullOrEmpty(catalog)) { + sb.append(quoted(catalog)).append("."); + } + if (!isNullOrEmpty(schema)) { + sb.append(quoted(schema)).append("."); + } + sb.append(quoted(table)); + return sb.toString(); + } + + public String quoted(String name) + { + name = name.replace(IDENTIFIER_QUOTE, IDENTIFIER_QUOTE + IDENTIFIER_QUOTE); + return IDENTIFIER_QUOTE + name + IDENTIFIER_QUOTE; + } +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManager.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManager.java new file mode 100644 index 000000000000..034ee4d8ff68 --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManager.java @@ -0,0 +1,26 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.statistics.TableStatistics; + +import java.util.List; + +public interface StatisticsManager +{ + TableStatistics getTableStatistics(ConnectorSession session, JdbcTableHandle handle, List columns); +} diff --git a/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManagerImpl.java b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManagerImpl.java new file mode 100644 index 000000000000..2f1e6393fe9f --- /dev/null +++ b/plugin/trino-adb/src/main/java/io/trino/plugin/adb/connector/table/StatisticsManagerImpl.java @@ -0,0 +1,262 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.adb.connector.table; + +import com.google.inject.Inject; +import io.trino.plugin.jdbc.ConnectionFactory; +import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcStatisticsConfig; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.statistics.ColumnStatistics; +import io.trino.spi.statistics.Estimate; +import io.trino.spi.statistics.TableStatistics; +import org.jdbi.v3.core.Handle; +import org.jdbi.v3.core.Jdbi; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Throwables.throwIfInstanceOf; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; +import static java.util.function.Function.identity; + +public class StatisticsManagerImpl + implements StatisticsManager +{ + private final ConnectionFactory connectionFactory; + private final boolean statisticsEnabled; + + @Inject + public StatisticsManagerImpl(ConnectionFactory connectionFactory, + JdbcStatisticsConfig statisticsConfig) + { + this.connectionFactory = connectionFactory; + statisticsEnabled = statisticsConfig.isEnabled(); + } + + @Override + public TableStatistics getTableStatistics(ConnectorSession session, + JdbcTableHandle handle, + List columns) + { + if (!statisticsEnabled) { + return TableStatistics.empty(); + } + if (!handle.isNamedRelation()) { + return TableStatistics.empty(); + } + try { + return readTableStatistics(session, handle, columns); + } + catch (SQLException | RuntimeException e) { + throwIfInstanceOf(e, TrinoException.class); + throw new TrinoException(JDBC_ERROR, "Failed fetching statistics for table: " + handle, e); + } + } + + private TableStatistics readTableStatistics(ConnectorSession session, JdbcTableHandle table, + List columns) + throws SQLException + { + checkArgument(table.isNamedRelation(), "Relation is not a table: %s", table); + try (Connection connection = connectionFactory.openConnection(session); + Handle handle = Jdbi.open(connection)) { + StatisticsDao statisticsDao = new StatisticsDao(handle); + + Optional optionalRowCount = readRowCountTableStat(statisticsDao, table); + if (optionalRowCount.isEmpty()) { + // Table not found + return TableStatistics.empty(); + } + long rowCount = optionalRowCount.get(); + if (rowCount == -1) { + // Table has never yet been vacuumed or analyzed + return TableStatistics.empty(); + } + TableStatistics.Builder tableStatistics = TableStatistics.builder(); + tableStatistics.setRowCount(Estimate.of(rowCount)); + + if (rowCount == 0) { + return tableStatistics.build(); + } + + RemoteTableName remoteTableName = table.getRequiredNamedRelation().getRemoteTableName(); + Map columnStatistics = + statisticsDao.getColumnStatistics(remoteTableName.getSchemaName().orElse(null), + remoteTableName.getTableName()).stream() + .collect(toImmutableMap(ColumnStatisticsResult::columnName, identity())); + + for (JdbcColumnHandle column : columns) { + ColumnStatisticsResult result = columnStatistics.get(column.getColumnName()); + if (result == null) { + continue; + } + + ColumnStatistics statistics = ColumnStatistics.builder() + .setNullsFraction(result.nullsFraction() + .map(Estimate::of) + .orElseGet(Estimate::unknown)) + .setDistinctValuesCount(result.distinctValuesIndicator() + .map(distinctValuesIndicator -> { + if (distinctValuesIndicator >= 0.0) { + return distinctValuesIndicator; + } + return -distinctValuesIndicator * rowCount; + }) + .map(Estimate::of) + .orElseGet(Estimate::unknown)) + .setDataSize(result.averageColumnLength() + .flatMap(averageColumnLength -> + result.nullsFraction().map(nullsFraction -> + Estimate.of( + 1.0 * averageColumnLength * rowCount * (1 - nullsFraction)))) + .orElseGet(Estimate::unknown)) + .build(); + + tableStatistics.setColumnStatistics(column, statistics); + } + + return tableStatistics.build(); + } + } + + private static Optional readRowCountTableStat(StatisticsDao statisticsDao, JdbcTableHandle table) + { + RemoteTableName remoteTableName = table.getRequiredNamedRelation().getRemoteTableName(); + String schemaName = remoteTableName.getSchemaName().orElse(null); + Optional rowCount = statisticsDao.getRowCountFromPgClass(schemaName, remoteTableName.getTableName()); + if (rowCount.isEmpty()) { + // Table not found + return Optional.empty(); + } + if (statisticsDao.isPartitionedTable(schemaName, remoteTableName.getTableName())) { + Optional partitionedTableRowCount = + statisticsDao.getRowCountPartitionedTableFromPgClass(schemaName, remoteTableName.getTableName()); + if (partitionedTableRowCount.isPresent()) { + return partitionedTableRowCount; + } + + return statisticsDao.getRowCountPartitionedTableFromPgStats(schemaName, remoteTableName.getTableName()); + } + if (rowCount.get() == 0) { + rowCount = statisticsDao.getRowCountFromPgStat(schemaName, remoteTableName.getTableName()); + } + return rowCount; + } + + private static class StatisticsDao + { + private final Handle handle; + + public StatisticsDao(Handle handle) + { + this.handle = requireNonNull(handle, "handle is null"); + } + + Optional getRowCountFromPgClass(String schema, String tableName) + { + return handle.createQuery("SELECT reltuples " + + "FROM pg_class " + + "WHERE relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = :schema) " + + "AND relname = :table_name") + .bind("schema", schema) + .bind("table_name", tableName) + .mapTo(Long.class) + .findOne(); + } + + Optional getRowCountFromPgStat(String schema, String tableName) + { + return handle.createQuery( + "SELECT n_live_tup FROM pg_stat_all_tables WHERE schemaname = :schema AND relname = :table_name") + .bind("schema", schema) + .bind("table_name", tableName) + .mapTo(Long.class) + .findOne(); + } + + Optional getRowCountPartitionedTableFromPgClass(String schema, String tableName) + { + return handle.createQuery("SELECT SUM(child.reltuples) " + + "FROM pg_inherits " + + "JOIN pg_class parent ON pg_inherits.inhparent = parent.oid " + + "JOIN pg_class child ON pg_inherits.inhrelid = child.oid " + + "JOIN pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace " + + "JOIN pg_namespace child_ns ON child_ns.oid = child.relnamespace " + + "WHERE parent.oid = :schema_table_name::regclass") + .bind("schema_table_name", format("%s.%s", schema, tableName)) + .mapTo(Long.class) + .findOne(); + } + + Optional getRowCountPartitionedTableFromPgStats(String schema, String tableName) + { + return handle.createQuery("SELECT SUM(stat.n_live_tup) " + + "FROM pg_inherits " + + "JOIN pg_class parent ON pg_inherits.inhparent = parent.oid " + + "JOIN pg_class child ON pg_inherits.inhrelid = child.oid " + + "JOIN pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace " + + "JOIN pg_namespace child_ns ON child_ns.oid = child.relnamespace " + + "JOIN pg_stat_all_tables stat ON stat.schemaname = child_ns.nspname AND stat.relname = child.relname " + + "WHERE parent.oid = :schema_table_name::regclass") + .bind("schema_table_name", format("%s.%s", schema, tableName)) + .mapTo(Long.class) + .findOne(); + } + + List getColumnStatistics(String schema, String tableName) + { + return handle.createQuery( + "SELECT attname, null_frac, n_distinct, avg_width FROM pg_stats WHERE schemaname = :schema AND tablename = :table_name") + .bind("schema", schema) + .bind("table_name", tableName) + .map((rs, ctx) -> new ColumnStatisticsResult( + requireNonNull(rs.getString("attname"), "attname is null"), + Optional.ofNullable(rs.getObject("null_frac", Float.class)), + Optional.ofNullable(rs.getObject("n_distinct", Float.class)), + Optional.ofNullable(rs.getObject("avg_width", Integer.class)))) + .list(); + } + + boolean isPartitionedTable(String schema, String tableName) + { + return handle.createQuery("SELECT true " + + "FROM pg_class " + + "WHERE relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = :schema) " + + "AND relname = :table_name " + + "AND relkind = 'p'") + .bind("schema", schema) + .bind("table_name", tableName) + .mapTo(Boolean.class) + .findOne() + .orElse(false); + } + } + + private record ColumnStatisticsResult(String columnName, Optional nullsFraction, + Optional distinctValuesIndicator, + Optional averageColumnLength) + { + } +} diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DynamicFilteringJdbcSplitSource.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DynamicFilteringJdbcSplitSource.java index 812979ed2bcc..cca6e851f00c 100644 --- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DynamicFilteringJdbcSplitSource.java +++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DynamicFilteringJdbcSplitSource.java @@ -64,7 +64,7 @@ public CompletableFuture getNextBatch(int maxSize) JdbcSplit jdbcSplit = (JdbcSplit) split; // If split was a subclass of JdbcSplit, there would be additional information // that we would need to pass further on. - verify(jdbcSplit.getClass() == JdbcSplit.class, "Unexpected split type %s", jdbcSplit); + verify(jdbcSplit.getClass() == JdbcSplit.class || jdbcSplit instanceof JdbcDynamicFilterAwareSplit, "Unexpected split type %s", jdbcSplit); return jdbcSplit.withDynamicFilter(dynamicFilterPredicate); }) .collect(toImmutableList()), diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcDynamicFilterAwareSplit.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcDynamicFilterAwareSplit.java new file mode 100644 index 000000000000..fbffd08dbe00 --- /dev/null +++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcDynamicFilterAwareSplit.java @@ -0,0 +1,18 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.jdbc; + +public interface JdbcDynamicFilterAwareSplit +{ +} diff --git a/pom.xml b/pom.xml index b0ec67e73889..137ddd6a18b0 100644 --- a/pom.xml +++ b/pom.xml @@ -60,6 +60,7 @@ lib/trino-record-decoder plugin/trino-accumulo plugin/trino-accumulo-iterators + plugin/trino-adb plugin/trino-atop plugin/trino-base-jdbc plugin/trino-bigquery diff --git a/testing/trino-server-dev/etc/catalog/adb.properties b/testing/trino-server-dev/etc/catalog/adb.properties new file mode 100644 index 000000000000..f381b3a8c1c1 --- /dev/null +++ b/testing/trino-server-dev/etc/catalog/adb.properties @@ -0,0 +1,7 @@ +connector.name=adb +connection-url=jdbc:postgresql://localhost:6000/postgres +connection-user=gpadmin +connection-password=gpadmin + +adb.gpfdist.server.port=8990 +adb.gpfdist.server.host=trino-server diff --git a/testing/trino-server-dev/etc/config.properties b/testing/trino-server-dev/etc/config.properties index 60d21a787243..ebd3b1198f49 100644 --- a/testing/trino-server-dev/etc/config.properties +++ b/testing/trino-server-dev/etc/config.properties @@ -26,6 +26,7 @@ query.client.timeout=5m query.min-expire-age=30m plugin.bundles=\ + ../../plugin/trino-adb/pom.xml, \ ../../plugin/trino-resource-group-managers/pom.xml,\ ../../plugin/trino-password-authenticators/pom.xml, \ ../../plugin/trino-iceberg/pom.xml,\ diff --git a/testing/trino-server-dev/etc/log.properties b/testing/trino-server-dev/etc/log.properties index b615d661c74a..40f0b939e26c 100644 --- a/testing/trino-server-dev/etc/log.properties +++ b/testing/trino-server-dev/etc/log.properties @@ -5,7 +5,7 @@ # in production. For example configuration, see the Trino documentation. # -io.trino=INFO +io.trino=DEBUG # show classpath for plugins io.trino.server.PluginManager=DEBUG