diff --git a/core/trino-server/src/main/provisio/trino.xml b/core/trino-server/src/main/provisio/trino.xml index 8fcbb0d30943..1da10a350243 100644 --- a/core/trino-server/src/main/provisio/trino.xml +++ b/core/trino-server/src/main/provisio/trino.xml @@ -40,6 +40,12 @@ + + + + + + diff --git a/plugin/trino-couchbase/pom.xml b/plugin/trino-couchbase/pom.xml new file mode 100644 index 000000000000..3c108bfc4c73 --- /dev/null +++ b/plugin/trino-couchbase/pom.xml @@ -0,0 +1,267 @@ + + + 4.0.0 + + + io.trino + trino-root + 480-SNAPSHOT + ../../pom.xml + + + trino-couchbase + trino-plugin + ${project.artifactId} + Trino - Couchbase Connector + + + true + + + + + + com.couchbase.client + core-io + 3.11.1 + + + + com.couchbase.client + java-client + 3.11.1 + + + + com.google.guava + guava + + + + com.google.inject + guice + classes + + + + io.airlift + bootstrap + + + + io.airlift + configuration + + + + io.airlift + json + + + + io.airlift + security + + + + io.trino + trino-plugin-toolkit + + + + jakarta.annotation + jakarta.annotation-api + + + + jakarta.inject + jakarta.inject-api + + + + jakarta.validation + jakarta.validation-api + + + + org.slf4j + slf4j-api + + + + org.weakref + jmxutils + + + + com.fasterxml.jackson.core + jackson-annotations + provided + + + + io.airlift + slice + provided + + + + io.opentelemetry + opentelemetry-api + provided + + + + io.opentelemetry + opentelemetry-api-incubator + provided + + + + io.opentelemetry + opentelemetry-context + provided + + + + io.trino + trino-spi + provided + + + + io.airlift + log + runtime + + + + io.airlift + log-manager + runtime + + + + io.airlift + units + runtime + + + + com.exasol + exasol-testcontainers + 7.2.2 + test + + + + io.airlift + configuration-testing + test + + + io.airlift + junit-extensions + test + + + + io.airlift + testing + test + + + + io.trino + trino-base-jdbc + test-jar + test + + + + io.trino + trino-main + test + + + + io.trino + trino-main + test-jar + test + + + + io.trino + trino-parser + test + + + + io.trino + trino-testing + test + + + + io.trino + trino-testing-services + test + + + + io.trino + trino-tpch + test + + + + io.trino.tpch + tpch + test + + + + org.assertj + assertj-core + test + + + + org.jetbrains + annotations + test + + + + org.junit.jupiter + junit-jupiter-api + test + + + + org.junit.jupiter + junit-jupiter-engine + test + + + + org.junit.jupiter + junit-jupiter-params + test + + + + org.testcontainers + testcontainers + test + + + org.testcontainers + testcontainers-couchbase + 2.0.1 + test + + + + diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseClient.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseClient.java new file mode 100644 index 000000000000..8fd0af09b7e6 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseClient.java @@ -0,0 +1,121 @@ +/* + * 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.couchbase; + +import com.couchbase.client.core.env.Authenticator; +import com.couchbase.client.core.env.CertificateAuthenticator; +import com.couchbase.client.java.Bucket; +import com.couchbase.client.java.Cluster; +import com.couchbase.client.java.ClusterOptions; +import com.couchbase.client.java.Scope; +import io.airlift.security.pem.PemReader; +import jakarta.inject.Inject; + +import java.io.File; +import java.nio.file.Path; +import java.security.KeyStore; +import java.security.KeyStoreException; +import java.security.PrivateKey; +import java.security.cert.Certificate; +import java.security.cert.X509Certificate; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; + +public class CouchbaseClient +{ + private final CouchbaseConfig config; + private final Cluster cluster; + + @Inject + public CouchbaseClient(CouchbaseConfig config) + { + this.config = config; + try { + if (config.getTlsKey() != null) { + PrivateKey key; + Optional password = Optional.ofNullable(config.getTlsKeyPassword()); + List keyCertChain = new ArrayList<>(); + if (new File(config.getTlsKey()).exists()) { + // load from file + key = PemReader.loadPrivateKey(new File(config.getTlsKey()), password); + } + else { + // try loading from string + key = PemReader.loadPrivateKey(config.getTlsKey(), password); + } + if (config.getTlsCertificate() != null) { + KeyStore tlsKeyStore = PemReader.loadTrustStore(new File(config.getTlsCertificate())); + tlsKeyStore.aliases().asIterator().forEachRemaining(alias -> { + try { + for (Certificate cert : tlsKeyStore.getCertificateChain(alias)) { + if (cert instanceof X509Certificate) { + keyCertChain.add((X509Certificate) cert); + } + } + } + catch (KeyStoreException e) { + throw new RuntimeException("Failed to load TLS certificates", e); + } + }); + } + Authenticator authenticator = CertificateAuthenticator.fromKey( + key, password.orElse(""), keyCertChain); + cluster = Cluster.connect( + config.getCluster(), + ClusterOptions.clusterOptions(authenticator) + .environment(env -> { + env.securityConfig(security -> { + if (config.getTlsCertificate() != null) { + security.trustCertificate(Path.of(config.getTlsCertificate())); + } + }); + env.timeoutConfig(timeout -> { + timeout.kvTimeout(config.getTimeouts()); + timeout.queryTimeout(config.getTimeouts()); + }); + })); + } + else { + cluster = Cluster.connect( + config.getCluster(), + ClusterOptions.clusterOptions(config.getUsername(), config.getPassword()) + .environment(env -> { + env.securityConfig(security -> { + if (config.getTlsCertificate() != null) { + security.trustCertificate(Path.of(config.getTlsCertificate())); + } + }); + env.timeoutConfig(timeout -> { + timeout.kvTimeout(config.getTimeouts()); + timeout.queryTimeout(config.getTimeouts()); + }); + })); + } + } + catch (Exception e) { + throw new RuntimeException("Failed to instantiate Couchbase client", e); + } + } + + public Bucket getBucket() + { + return cluster.bucket(config.getBucket()); + } + + public Scope getScope() + { + return getBucket().scope(config.getScope()); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseColumnHandle.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseColumnHandle.java new file mode 100644 index 000000000000..52dc76ae6043 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseColumnHandle.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.couchbase; + +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableList; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.type.Type; + +import java.util.List; + +public record CouchbaseColumnHandle(List path, List dereferenceNames, Type type, + boolean synthetic) implements ColumnHandle +{ + public CouchbaseColumnHandle + { + path = ImmutableList.copyOf(path); + } + + public String fullName() + { + return String.format("`%s`", Joiner.on("`.`").join(path)); + } + + public String name() + { + return path.getLast(); + } + + public boolean isSynthetic() + { + return synthetic; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConfig.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConfig.java new file mode 100644 index 000000000000..e0cfcfeeea7b --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConfig.java @@ -0,0 +1,181 @@ +/* + * 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.couchbase; + +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; +import io.airlift.configuration.ConfigSecuritySensitive; +import io.trino.spi.function.Description; + +import java.time.Duration; + +public class CouchbaseConfig +{ + private String cluster = "localhost"; + private String username = "Administrator"; + private String password = "password"; + private String tlsKey; + private String tlsKeyPassword; + private String tlsCertificate; + private String schemaFolder = "couchbase-schema"; + private String bucket = "default"; + private String scope = "_default"; + private Duration timeouts = Duration.ofSeconds(60); + private Long pageSize = 5000L; + + @Config("couchbase.cluster") + @ConfigDescription("Couchbase cluster connection string") + public CouchbaseConfig setCluster(String connstring) + { + this.cluster = connstring; + return this; + } + + public String getCluster() + { + return cluster; + } + + @Config("couchbase.username") + @ConfigDescription("Username for the cluster") + public CouchbaseConfig setUsername(String username) + { + this.username = username; + return this; + } + + public String getUsername() + { + return username; + } + + @Config("couchbase.password") + @ConfigDescription("Password for the cluster") + @ConfigSecuritySensitive + public CouchbaseConfig setPassword(String password) + { + this.password = password; + return this; + } + + public String getPassword() + { + return password; + } + + @Config("couchbase.tls-key") + @ConfigDescription("Key file address for mTls") + public CouchbaseConfig setTlsKey(String tlsKey) + { + this.tlsKey = tlsKey; + return this; + } + + public String getTlsKey() + { + return tlsKey; + } + + @Config("couchbase.tls-key-password") + @ConfigDescription("Key password") + @ConfigSecuritySensitive + public CouchbaseConfig setTlsKeyPassword(String password) + { + this.tlsKeyPassword = password; + return this; + } + + public String getTlsKeyPassword() + { + return tlsKeyPassword; + } + + @Config("couchbase.tls-certificate") + @ConfigDescription("Cluster root certificate file address") + public CouchbaseConfig setTlsCertificate(String certificate) + { + this.tlsCertificate = certificate; + return this; + } + + public String getTlsCertificate() + { + return tlsCertificate; + } + + @Config("couchbase.schema-folder") + @ConfigDescription("Path for folder with json files containing Trino schema mappings") + public CouchbaseConfig setSchemaFolder(String schemaFolder) + { + this.schemaFolder = schemaFolder; + return this; + } + + public String getSchemaFolder() + { + return schemaFolder; + } + + @Config("couchbase.bucket") + @ConfigDescription("Bucket to connect to") + public CouchbaseConfig setBucket(String bucket) + { + this.bucket = bucket; + return this; + } + + public String getBucket() + { + return bucket; + } + + @Config("couchbase.scope") + @ConfigDescription("Scope to connect to") + public CouchbaseConfig setScope(String scope) + { + this.scope = scope; + return this; + } + + public String getScope() + { + return scope; + } + + @Config("couchbase.timeouts") + @ConfigDescription("Operations timeout in seconds") + public CouchbaseConfig setTimeouts(String timeout) + { + this.timeouts = Duration.ofSeconds(Long.parseLong(timeout)); + return this; + } + + public Duration getTimeouts() + { + return timeouts; + } + + @Config("couchbase.page-size") + @Description("Maximum number of rows to be fetched in a single query") + public CouchbaseConfig setPageSize(String value) + { + this.pageSize = Long.valueOf(value); + return this; + } + + public Long getPageSize() + { + return pageSize; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnector.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnector.java new file mode 100644 index 000000000000..5ab60b614e1d --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnector.java @@ -0,0 +1,71 @@ +/* + * 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.couchbase; + +import com.google.inject.Inject; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorMetadata; +import io.trino.spi.connector.ConnectorPageSourceProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitManager; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.transaction.IsolationLevel; + +public class CouchbaseConnector + implements Connector +{ + private final CouchbaseMetadata metadata; + private final CouchbaseSplitManager splitManager; + private final CouchbasePageSourceProvider pageSourceProvider; + + @Inject + public CouchbaseConnector( + CouchbaseMetadata metadata, + CouchbaseSplitManager splitManager, + CouchbasePageSourceProvider pageSourceProvider) + { + this.metadata = metadata; + this.splitManager = splitManager; + this.pageSourceProvider = pageSourceProvider; + } + + @Override + public void shutdown() + { + } + + @Override + public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly, boolean autoCommit) + { + return CouchbaseTransactionHandle.INSTANCE; + } + + @Override + public ConnectorMetadata getMetadata(ConnectorSession session, ConnectorTransactionHandle transactionHandle) + { + return metadata; + } + + @Override + public ConnectorSplitManager getSplitManager() + { + return splitManager; + } + + @Override + public ConnectorPageSourceProvider getPageSourceProvider() + { + return pageSourceProvider; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorFactory.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorFactory.java new file mode 100644 index 000000000000..3bd25a48a4ec --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorFactory.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.couchbase; + +import com.google.inject.Injector; +import io.airlift.bootstrap.Bootstrap; +import io.airlift.json.JsonModule; +import io.trino.plugin.base.ConnectorContextModule; +import io.trino.plugin.base.TypeDeserializerModule; +import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule; +import io.trino.plugin.base.jmx.MBeanServerModule; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorContext; +import io.trino.spi.connector.ConnectorFactory; +import org.weakref.jmx.guice.MBeanModule; + +import java.util.Map; + +import static io.trino.plugin.base.Versions.checkStrictSpiVersionMatch; +import static java.util.Objects.requireNonNull; + +public class CouchbaseConnectorFactory + implements ConnectorFactory +{ + @Override + public String getName() + { + return "couchbase"; + } + + @Override + public Connector create(String catalogName, Map config, ConnectorContext context) + { + requireNonNull(catalogName, "catalogName is null"); + requireNonNull(config, "config is null"); + checkStrictSpiVersionMatch(context, this); + + Bootstrap app = new Bootstrap( + "io.trino.bootstrap.catalog." + catalogName, + new MBeanModule(), + new MBeanServerModule(), + new ConnectorObjectNameGeneratorModule("io.trino.plugin.couchbase", "trino.plugin.couchbase"), + new JsonModule(), + new TypeDeserializerModule(), + new CouchbaseConnectorModule(), + new ConnectorContextModule(catalogName, context)); + + Injector injector = app + .doNotInitializeLogging() + .disableSystemProperties() + .setRequiredConfigurationProperties(config) + .initialize(); + + return injector.getInstance(CouchbaseConnector.class); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorModule.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorModule.java new file mode 100644 index 000000000000..1bd2458c7543 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseConnectorModule.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.couchbase; + +import com.google.inject.Binder; +import com.google.inject.Scopes; +import io.airlift.configuration.AbstractConfigurationAwareModule; + +import static io.airlift.configuration.ConfigBinder.configBinder; +import static org.weakref.jmx.guice.ExportBinder.newExporter; + +public class CouchbaseConnectorModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + binder.bind(CouchbaseConnector.class).in(Scopes.SINGLETON); + binder.bind(CouchbaseMetadata.class).in(Scopes.SINGLETON); + binder.bind(CouchbaseSplitManager.class).in(Scopes.SINGLETON); + binder.bind(CouchbasePageSourceProvider.class).in(Scopes.SINGLETON); + binder.bind(CouchbaseClient.class).in(Scopes.SINGLETON); + + newExporter(binder).export(CouchbaseClient.class).withGeneratedName(); + + configBinder(binder).bindConfig(CouchbaseConfig.class); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseMetadata.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseMetadata.java new file mode 100644 index 000000000000..ba558165e536 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseMetadata.java @@ -0,0 +1,630 @@ +/* + * 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.couchbase; + +import com.couchbase.client.java.json.JsonArray; +import com.couchbase.client.java.json.JsonObject; +import com.couchbase.client.java.manager.collection.CollectionManager; +import com.couchbase.client.protostellar.admin.collection.v1.CollectionAdminServiceGrpc; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import io.trino.plugin.base.projection.ApplyProjectionUtil; +import io.trino.plugin.couchbase.translations.TrinoExpressionToCb; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.AggregateFunction; +import io.trino.spi.connector.AggregationApplicationResult; +import io.trino.spi.connector.Assignment; +import io.trino.spi.connector.CatalogSchemaTableName; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorMetadata; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.ConnectorTableVersion; +import io.trino.spi.connector.Constraint; +import io.trino.spi.connector.ConstraintApplicationResult; +import io.trino.spi.connector.JoinApplicationResult; +import io.trino.spi.connector.JoinStatistics; +import io.trino.spi.connector.JoinType; +import io.trino.spi.connector.LimitApplicationResult; +import io.trino.spi.connector.ProjectionApplicationResult; +import io.trino.spi.connector.SampleApplicationResult; +import io.trino.spi.connector.SampleType; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.SortItem; +import io.trino.spi.connector.TableFunctionApplicationResult; +import io.trino.spi.connector.TableScanRedirectApplicationResult; +import io.trino.spi.connector.TopNApplicationResult; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.expression.Constant; +import io.trino.spi.expression.Variable; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; +import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.BigintType; +import io.trino.spi.type.BooleanType; +import io.trino.spi.type.DateType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.IntegerType; +import io.trino.spi.type.RowType; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.VarcharType; +import jakarta.annotation.Nullable; +import jakarta.inject.Inject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.Reader; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; + +import static com.couchbase.client.core.deps.com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.plugin.base.projection.ApplyProjectionUtil.extractSupportedProjectedColumns; +import static io.trino.plugin.base.projection.ApplyProjectionUtil.replaceWithNewVariables; +import static io.trino.plugin.couchbase.translations.TrinoToCbType.isPushdownSupportedType; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static java.util.Objects.requireNonNull; +import static java.util.function.Function.identity; + +public class CouchbaseMetadata + implements ConnectorMetadata +{ + private static final Logger LOG = LoggerFactory.getLogger(CouchbaseMetadata.class); + private final TypeManager typeManager; + private final CouchbaseClient client; + private final CouchbaseConfig config; + private final Map mtimeCache = new ConcurrentHashMap<>(); + private final Map metaCache = new ConcurrentHashMap<>(); + + @Inject + public CouchbaseMetadata(TypeManager typeManager, CouchbaseClient client, CouchbaseConfig config) + { + this.typeManager = typeManager; + this.client = client; + this.config = config; + } + + @Override + public boolean schemaExists(ConnectorSession session, String schemaName) + { + return client.getBucket().collections().getAllScopes().stream() + .filter(scope -> scope.name().equals(schemaName)) + .findAny().isPresent(); + } + + @Nullable + @Override + public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName, Optional startVersion, Optional endVersion) + { + if (startVersion.isPresent() || endVersion.isPresent()) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support versioned tables"); + } + + CollectionManager cm = client.getBucket().collections(); + return cm.getAllScopes().stream() + .filter(scopeSpec -> scopeSpec.name().equals(tableName.getSchemaName())) + .flatMap(scopeSpec -> scopeSpec.collections().stream()) + .filter(collectionSpec -> collectionSpec.name().equals(tableName.getTableName())) + .findFirst().map(collectionSpec -> CouchbaseTableHandle.fromSchemaAndName(tableName.getSchemaName(), tableName.getTableName())) + .orElse(null); + } + + @Override + public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle table) + { + requireNonNull(config.getSchemaFolder(), "Couchbase schema folder is not set"); + if (!(table instanceof CouchbaseTableHandle)) { + throw new RuntimeException("Couchbase table handle is not an instance of CouchbaseTableHandle"); + } + + CouchbaseTableHandle handle = (CouchbaseTableHandle) table; + String tablePath = handle.path(); + checkMtime(handle); + if (metaCache.get(tablePath) == null) { + File schemaFile = getSchemaFile(handle); + if (!schemaFile.exists()) { + throw new RuntimeException(String.format("Couchbase schema file '%s' does not exist", schemaFile.getAbsolutePath())); + } + List columns = new LinkedList<>(); + try (Reader schemaReader = Files.newBufferedReader(Path.of(schemaFile.toURI()))) { + JsonObject schema = JsonObject.fromJson(schemaReader.readAllAsString()); + JsonObject properties = schema.getObject("properties"); + + ColumnMetadata[] orderedColumns = new ColumnMetadata[properties.size()]; + boolean unordered = false; + boolean ordered = false; + for (String propertyName : properties.getNames()) { + if (propertyName.equals("~meta")) { + // skip the meta column + continue; + } + try { + JsonObject property = properties.getObject(propertyName); + if (property.containsKey("order")) { + if (unordered) { + throw new RuntimeException(String.format("unable to mix ordered and unordered properties: %s", tablePath)); + } + int order = property.getInt("order"); + orderedColumns[order] = new ColumnMetadata(propertyName, deductType(property)); + } + else { + if (ordered) { + throw new RuntimeException(String.format("unable to mix ordered and unordered properties: %s", tablePath)); + } + unordered = true; + columns.add(new ColumnMetadata(propertyName, deductType(property))); + } + } + catch (Exception e) { + throw new RuntimeException(String.format("Failed to read schema for column '%s': %s", + propertyName, e.getMessage()), e); + } + } + + if (!unordered) { + columns = Arrays.asList(orderedColumns); + } + + LOG.debug("Loaded schema for table '{}': {}", tablePath, columns); + ConnectorTableMetadata result = new ConnectorTableMetadata(new SchemaTableName(handle.schema(), handle.name()), columns); + metaCache.put(tablePath, result); + mtimeCache.put(tablePath, Files.getLastModifiedTime(schemaFile.toPath()).toMillis()); + } + catch (Exception e) { + LOG.error(String.format("Failed to read schema for table '%s'", tablePath), e); + throw new RuntimeException(String.format("Failed to read schema for collection '%s.%s': %s", + ((CouchbaseTableHandle) table).schema(), ((CouchbaseTableHandle) table).name(), e.getMessage()), e); + } + } + + return metaCache.get(tablePath); + } + + @Override + public Map getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle) + { + CouchbaseTableHandle handle = (CouchbaseTableHandle) tableHandle; + ConnectorTableMetadata tableMetadata = getTableMetadata(session, tableHandle); + return tableMetadata.getColumns().stream() + .collect(Collectors.toMap(ColumnMetadata::getName, + column -> new CouchbaseColumnHandle( + Arrays.asList(handle.schema(), handle.name(), column.getName()), + new ArrayList<>(), + column.getType(), + false))); + } + + @Override + public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle) + { + ConnectorTableMetadata tableMeta = getTableMetadata(session, tableHandle); + return tableMeta.getColumns().stream() + .filter(column -> column.getName().equals(((CouchbaseColumnHandle) columnHandle).name())) + .findFirst() + .orElseGet(() -> { + if (columnHandle instanceof CouchbaseColumnHandle cbHandle) { + return new ColumnMetadata(cbHandle.name(), cbHandle.type()); + } + return null; + }); + } + + private File getSchemaFile(CouchbaseTableHandle handle) + { + return new File(new File(config.getSchemaFolder()), String.format("%s.%s.%s.json", client.getBucket().name(), handle.schema(), handle.name())); + } + + private void checkMtime(CouchbaseTableHandle handle) + { + final String path = handle.path(); + try { + long cached = mtimeCache.getOrDefault(path, 0L); + long actual = Files.getLastModifiedTime(getSchemaFile(handle).toPath()).toMillis(); + if (actual - cached > 1000) { + metaCache.remove(path); + } + } + catch (Exception e) { + metaCache.remove(path); + } + } + + @Override + public Optional> applyFilter(ConnectorSession session, ConnectorTableHandle handle, Constraint constraint) + { + if (!takeOrReject(constraint.getAssignments())) { + LOG.info("rejecting constraint {} for table {}: unsupported assignments", constraint, handle); + return Optional.empty(); + } + + if (handle instanceof CouchbaseTableHandle cbHandle) { + // don't apply where on aggregations or root query + if (cbHandle.subQuery().isEmpty() || cbHandle.isAggregated().get() || !cbHandle.groupings().isEmpty()) { + cbHandle = cbHandle.wrap(); + } + TupleDomain oldDomain = cbHandle.constraint(); + TupleDomain newDomain = oldDomain.intersect(constraint.getSummary()); + TupleDomain remainingFilter; + if (newDomain.isNone()) { + remainingFilter = TupleDomain.all(); + } + else { + Map domains = newDomain.getDomains().orElseThrow(); + + Map supported = new HashMap<>(); + Map unsupported = new HashMap<>(); + + for (Map.Entry entry : domains.entrySet()) { + CouchbaseColumnHandle columnHandle = (CouchbaseColumnHandle) entry.getKey(); + Domain domain = entry.getValue(); + Type columnType = columnHandle.type(); + if (isPushdownSupportedType(columnType)) { + supported.put(entry.getKey(), entry.getValue()); + } + else { + unsupported.put(columnHandle, domain); + } + } + newDomain = TupleDomain.withColumnDomains(supported); + remainingFilter = TupleDomain.withColumnDomains(unsupported); + } + + boolean modified = false; + if (!oldDomain.equals(newDomain)) { + handle = cbHandle = cbHandle.withConstraint(newDomain); + modified = true; + } + + if (constraint.getExpression() != Constant.TRUE && !cbHandle.containsConstraint(constraint)) { + cbHandle.whereClauses().add(TrinoExpressionToCb.convert(constraint.getExpression(), constraint.getAssignments())); + modified = true; + } + + if (!modified) { + return Optional.empty(); + } + + return Optional.of(new ConstraintApplicationResult<>(handle, remainingFilter, constraint.getExpression(), false)); + } + return ConnectorMetadata.super.applyFilter(session, handle, constraint); + } + + private Type deductType(JsonObject property) + { + Object type = property.get("type"); + if (type instanceof String) { + type = JsonArray.from(type); + } + + JsonArray types = (JsonArray) type; + List deductedTypes = new ArrayList<>(); + for (int i = 0; i < types.size(); i++) { + String cbType = types.getString(i); + if (cbType.equals("null")) { + continue; + } + else if (cbType.equals("string")) { + deductedTypes.add(VarcharType.VARCHAR); + } + else if (cbType.startsWith("varchar(")) { + int size = Integer.valueOf(cbType.replace("varchar(", "").replace(")", "")); + deductedTypes.add(VarcharType.createVarcharType(size)); + } + else if (cbType.equals("boolean")) { + deductedTypes.add(BooleanType.BOOLEAN); + } + else if (cbType.equals("number")) { + deductedTypes.add(DecimalType.createDecimalType()); + } + else if (cbType.equals("integer")) { + deductedTypes.add(IntegerType.INTEGER); + } + else if (cbType.equals("date")) { + deductedTypes.add(DateType.DATE); + } + else if (cbType.equals("bigint")) { + deductedTypes.add(BigintType.BIGINT); + } + else if (cbType.equals("double")) { + deductedTypes.add(DoubleType.DOUBLE); + } + else if (cbType.equals("array")) { + deductedTypes.add(new ArrayType(deductType(property.getObject("items")))); + } + else if (cbType.equals("object")) { + List fields = new ArrayList<>(); + JsonObject properties = property.getObject("properties"); + for (String propertyName : properties.getNames()) { + JsonObject subProperty = properties.getObject(propertyName); + fields.add(new RowType.Field(Optional.of(propertyName), deductType(subProperty))); + } + deductedTypes.add(RowType.from(fields)); + } + else { + throw new RuntimeException("Unsupported couchbase type: " + cbType); + } + } + if (deductedTypes.size() != 1) { + throw new RuntimeException("Ambiguous couchbase type: " + deductedTypes); + } + return deductedTypes.get(0); + } + + @Override + public List listTables(ConnectorSession session, Optional schemaName) + { + CollectionManager cm = client.getBucket().collections(); + return cm.getAllScopes().stream() + .filter(scopeSpec -> scopeSpec.name().equals(client.getScope().name())) + .flatMap(scopeSpec -> scopeSpec.collections().stream()) + .map(collectionSpec -> new SchemaTableName(client.getScope().name(), collectionSpec.name())) + .toList(); + } + + @Override + public List listSchemaNames(ConnectorSession session) + { + return Arrays.asList(client.getScope().name()); + } + + @Override + public Optional> applyTopN(ConnectorSession session, ConnectorTableHandle handle, long topNCount, List sortItems, Map assignments) + { + if (!takeOrReject(assignments)) { + LOG.info("Rejecting topN assignments: {}", assignments); + return Optional.empty(); + } + + if (handle instanceof CouchbaseTableHandle cbHandle) { + if (cbHandle.topNCount().longValue() != -1 || !cbHandle.orderClauses().isEmpty()) { + if (cbHandle.topNCount().longValue() <= topNCount && cbHandle.compareSortItems(sortItems, + assignments)) { + LOG.info("Rejecting topN: no effect"); + return Optional.empty(); + } + LOG.info("Wrapping table handle: already got topN or non-matching order"); + cbHandle = cbHandle.wrap(); + handle = cbHandle; + } + + + + cbHandle.setTopNCount(topNCount); + cbHandle.addSortItems(sortItems, assignments); + } + else { + LOG.warn("Rejecting topN assignments: handle is not couchbase"); + return Optional.empty(); + } + + LOG.info("Accepted topN assignments: {}", handle); + return Optional.of(new TopNApplicationResult(handle, true, true)); + } + + private boolean takeOrReject(Map assignments) + { + return assignments.keySet().stream() + .allMatch(key -> assignments.get(key) instanceof CouchbaseColumnHandle); + } + + @Override + public Optional> applyLimit(ConnectorSession session, ConnectorTableHandle handle, long limit) + { + if (handle instanceof CouchbaseTableHandle cbHandle) { + if (cbHandle.topNCount().longValue() != -1) { + if (cbHandle.topNCount().longValue() <= limit) { + return Optional.empty(); + } + } + cbHandle.setTopNCount(limit); + } + else { + return Optional.empty(); + } + + return Optional.of(new LimitApplicationResult<>(handle, true, false)); + } + + @Override + public Optional> applyProjection(ConnectorSession session, ConnectorTableHandle handle, List projections, Map assignments) + { + if (!(handle instanceof CouchbaseTableHandle)) { + return Optional.empty(); + } + + CouchbaseTableHandle cbTable = (CouchbaseTableHandle) handle; + + try { + if (cbTable.containsProjections(projections, assignments)) { + if (cbTable.selectClauses().size() == projections.size()) { + return Optional.empty(); + } + cbTable = cbTable.wrap(); + } + } + catch (IllegalArgumentException e) { + LOG.warn(String.format("Exception while applying projections to couchbase table: %s", cbTable), e); + return Optional.empty(); + } + + Set projectedExpressions = projections.stream() + .flatMap(expression -> extractSupportedProjectedColumns(expression, ex -> true).stream()) + .collect(toImmutableSet()); + + Map columnProjections = projectedExpressions.stream() + .collect(toImmutableMap(identity(), ApplyProjectionUtil::createProjectedColumnRepresentation)); + + Map newAssignments = new HashMap<>(); + Map newColumnAssignmentMap = new HashMap<>(); + ImmutableMap.Builder newVariablesBuilder = ImmutableMap.builder(); + ImmutableSet.Builder projectedColumnsBuilder = ImmutableSet.builder(); + + for (Map.Entry entry : columnProjections.entrySet()) { + ConnectorExpression expression = entry.getKey(); + ApplyProjectionUtil.ProjectedColumnRepresentation projectedColumn = entry.getValue(); + + CouchbaseColumnHandle baseColumnHandle = (CouchbaseColumnHandle) assignments.get(projectedColumn.getVariable().getName()); + CouchbaseColumnHandle projectedColumnHandle = projectColumn(baseColumnHandle, projectedColumn.getDereferenceIndices(), expression.getType()); + String projectedColumnName = projectedColumnHandle.name(); + + Variable projectedColumnVariable = new Variable(projectedColumnName, expression.getType()); + Assignment newAssignment = new Assignment(projectedColumnName, projectedColumnHandle, expression.getType()); + newAssignments.putIfAbsent(projectedColumnName, newAssignment); + newColumnAssignmentMap.putIfAbsent(projectedColumnName, projectedColumnHandle); + + newVariablesBuilder.put(expression, projectedColumnVariable); + projectedColumnsBuilder.add(projectedColumnHandle); + } + + Map newVariables = newVariablesBuilder.buildOrThrow(); + List newProjections = projections.stream() + .map(expression -> replaceWithNewVariables(expression, newVariables)) + .collect(toImmutableList()); + + if (cbTable.containsProjections(newProjections, newColumnAssignmentMap)) { + if (newProjections.size() == cbTable.selectClauses().size()) { + return Optional.empty(); + } + cbTable = cbTable.wrap(); + } + + List outputAssignments = newAssignments.values().stream().collect(toImmutableList()); + List projectionAssignments = cbTable.addProjections(newProjections, newColumnAssignmentMap); + for (int i = 0; i < newProjections.size(); i++) { + String assignedName = projectionAssignments.get(i); + if (!newColumnAssignmentMap.containsKey(assignedName)) { + ConnectorExpression projection = newProjections.get(i); + newColumnAssignmentMap.put(assignedName, new CouchbaseColumnHandle( + Arrays.asList(cbTable.schema(), cbTable.name(), assignedName), + Collections.emptyList(), + projection.getType(), + !(projection instanceof Variable))); + } + } + + return Optional.of(new ProjectionApplicationResult<>( + cbTable, + newProjections, + outputAssignments, + false)); + } + + private static CouchbaseColumnHandle projectColumn(CouchbaseColumnHandle baseColumn, List indices, Type projectedColumnType) + { + if (indices.isEmpty()) { + return baseColumn; + } + ImmutableList.Builder dereferenceNamesBuilder = ImmutableList.builder(); + dereferenceNamesBuilder.addAll(baseColumn.dereferenceNames()); + + Type type = baseColumn.type(); + for (int index : indices) { + checkArgument(type instanceof RowType, "type should be Row type"); + RowType rowType = (RowType) type; + RowType.Field field = rowType.getFields().get(index); + dereferenceNamesBuilder.add(field.getName() + .orElseThrow(() -> new TrinoException(NOT_SUPPORTED, "ROW type does not have field names declared: " + rowType))); + type = field.getType(); + } + return new CouchbaseColumnHandle( + baseColumn.path(), + dereferenceNamesBuilder.build(), + projectedColumnType, + baseColumn.isSynthetic()); + } + + @Override + public Optional> applySample(ConnectorSession session, ConnectorTableHandle handle, SampleType sampleType, double sampleRatio) + { + return ConnectorMetadata.super.applySample(session, handle, sampleType, sampleRatio); + } + + @Override + public Optional> applyAggregation(ConnectorSession session, ConnectorTableHandle handle, List aggregates, Map assignments, List> groupingSets) + { + if (handle instanceof CouchbaseTableHandle cbTable) { + if (cbTable.containsAllAggregations(aggregates, assignments) && cbTable.containsAllGroupings(groupingSets)) { + return Optional.empty(); + } + ImmutableList.Builder newAssignmentsBuilder = ImmutableList.builder(); + ImmutableList.Builder projectionsBuilder = ImmutableList.builder(); + + cbTable = cbTable.wrap(); + + for (int i = 0; i < aggregates.size(); i++) { + AggregateFunction aggregateFunction = aggregates.get(i); + NamedParametrizedString result = cbTable.addAggregateFunction(aggregateFunction, assignments); + newAssignmentsBuilder.add( + new Assignment( + result.name(), + new CouchbaseColumnHandle( + Arrays.asList(cbTable.schema(), cbTable.name(), result.name()), + new ArrayList<>(), + aggregateFunction.getOutputType(), + true), + aggregateFunction.getOutputType())); + projectionsBuilder.add(new Variable(result.name(), aggregateFunction.getOutputType())); + } + + cbTable.addGroupings(groupingSets); + + return Optional.of(new AggregationApplicationResult<>( + cbTable, + projectionsBuilder.build(), + newAssignmentsBuilder.build(), + ImmutableMap.of(), + true + )); + } + return Optional.empty(); + } + + @Override + public Optional> applyJoin(ConnectorSession session, JoinType joinType, ConnectorTableHandle left, ConnectorTableHandle right, ConnectorExpression joinCondition, Map leftAssignments, Map rightAssignments, JoinStatistics statistics) + { + return ConnectorMetadata.super.applyJoin(session, joinType, left, right, joinCondition, leftAssignments, rightAssignments, statistics); + } + + @Override + public Optional> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle) + { + return ConnectorMetadata.super.applyTableFunction(session, handle); + } + + @Override + public Optional applyTableScanRedirect(ConnectorSession session, ConnectorTableHandle tableHandle) + { + return ConnectorMetadata.super.applyTableScanRedirect(session, tableHandle); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSource.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSource.java new file mode 100644 index 000000000000..67401bc184ea --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSource.java @@ -0,0 +1,449 @@ +/* + * 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.couchbase; + +import com.couchbase.client.java.json.JsonArray; +import com.couchbase.client.java.json.JsonObject; +import com.couchbase.client.java.query.QueryOptions; +import com.couchbase.client.java.query.QueryResult; +import com.google.common.collect.ImmutableMap; +import com.google.common.primitives.Shorts; +import com.google.common.primitives.SignedBytes; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.plugin.base.metrics.LongCount; +import io.trino.spi.Page; +import io.trino.spi.PageBuilder; +import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.block.SqlMap; +import io.trino.spi.block.SqlRow; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorPageSource; +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.DynamicFilter; +import io.trino.spi.connector.SourcePage; +import io.trino.spi.metrics.Metrics; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.BigintType; +import io.trino.spi.type.BooleanType; +import io.trino.spi.type.DateType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.Decimals; +import io.trino.spi.type.Int128; +import io.trino.spi.type.IntegerType; +import io.trino.spi.type.MapType; +import io.trino.spi.type.RealType; +import io.trino.spi.type.RowType; +import io.trino.spi.type.SmallintType; +import io.trino.spi.type.TinyintType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; + +import static com.couchbase.client.core.cnc.tracing.TracingAttribute.COLLECTION_NAME; +import static com.google.common.base.Verify.verify; +import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; + +public final class CouchbasePageSource + implements ConnectorPageSource +{ + private static final Logger LOG = LoggerFactory.getLogger(CouchbasePageSource.class); + private final ConnectorTransactionHandle transaction; + private final ConnectorSession session; + private final ConnectorSplit split; + private final CouchbaseTableHandle table; + private final PageBuilder pageBuilder; + private final CouchbaseClient client; + private final String queryString; + private final Long pageSize; + private final List types = new LinkedList<>(); + private final List names = new LinkedList<>(); + private long offset; + private long total, pageCount, duration, errors; + private boolean finished; + + public CouchbasePageSource(CouchbaseClient client, CouchbaseTransactionHandle transaction, ConnectorSession session, CouchbaseSplit split, CouchbaseTableHandle table, List columns, DynamicFilter dynamicFilter, Long pageSize) + { + this.client = client; + this.transaction = transaction; + this.session = session; + this.split = split; + this.pageSize = pageSize; + + if (columns != null && !columns.isEmpty()) { + CouchbaseTableHandle finalTable = table; + columns.forEach(column -> { + if (!finalTable.coversColumn(column)) { + finalTable.addColumn(column); + } + }); +// table = table.wrap(); +// table.addColumns(columns); + } +// else if (columns != null) { +// table.clearSelectElements(); +// } + + this.table = table; + + TupleDomain predicate = dynamicFilter.getCurrentPredicate(); + + if (!predicate.isAll()) { + table = table.wrap(); + table.addPredicate(predicate); + } + + long limit = table.topNCount().get(); + if (limit < 0) { + limit = pageSize; + } + columns.forEach(column -> { + types.add(column.type()); + names.add(column.name()); + }); + this.pageBuilder = new PageBuilder((int) Math.min(limit, pageSize), types); + queryString = String.format("SELECT %s FROM (%s) data OFFSET %%d LIMIT %%d", + names.stream().collect(Collectors.joining("`, `", "`", "`")), + table.toSql().replaceAll("%", "%%")); + } + + @Override + public SourcePage getNextSourcePage() + { + if (finished) { + return null; + } + final long started = System.currentTimeMillis(); + verify(pageBuilder.isEmpty()); + JsonArray queryArgs = JsonArray.create(); + table.getParameters().forEach(queryArgs::add); + QueryOptions options = QueryOptions.queryOptions().parameters(queryArgs); + + try { + final String query = String.format(queryString, offset, pageSize); + QueryResult result = client.getScope().query(query, options); + List rows = result.rowsAsObject(); + LOG.info("Couchbase query ({} result rows): {}; arguments: {}", rows.size(), query, queryArgs); + + for (int j = 0; j < rows.size(); j++) { + JsonObject row = rows.get(j); + pageBuilder.declarePosition(); + for (int i = 0; i < names.size(); i++) { + Type type = types.get(i); + BlockBuilder output = pageBuilder.getBlockBuilder(i); + appendValue(output, type, row.get(names.get(i))); + } + } + + // update metrics + offset += pageSize; + total += rows.size(); + pageCount++; + duration += System.currentTimeMillis() - started; + + if (rows.size() != pageSize) { + finished = true; + } + else if (rows.isEmpty()) { + finished = true; + return null; + } + } catch (Throwable t) { + errors++; + throw t; + } + + + Page page = pageBuilder.build(); + pageBuilder.reset(); + return SourcePage.create(page); + } + + private void appendValue(BlockBuilder output, Type type, Object value) + { + if (value == null) { + output.appendNull(); + return; + } + + Class javaType = type.getJavaType(); + + try { + if (type == BooleanType.BOOLEAN) { + type.writeBoolean(output, Boolean.valueOf(String.valueOf(value))); + return; + } + else if (type == VarcharType.VARCHAR || javaType == Slice.class) { + Slice slice = Slices.utf8Slice(String.valueOf(value)); + type.writeSlice(output, slice); + } + else if (javaType == long.class) { + if (type.equals(BigintType.BIGINT)) { + type.writeLong(output, ((Number) value).longValue()); + } + else if (type.equals(IntegerType.INTEGER)) { + type.writeLong(output, ((Number) value).intValue()); + } + else if (type.equals(SmallintType.SMALLINT)) { + type.writeLong(output, Shorts.checkedCast(((Number) value).longValue())); + } + else if (type.equals(TinyintType.TINYINT)) { + type.writeLong(output, SignedBytes.checkedCast(((Number) value).longValue())); + } + else if (type.equals(RealType.REAL)) { + type.writeLong(output, Float.floatToIntBits(((Number) value).floatValue())); + } + else if (type instanceof DecimalType decimalType) { + LOG.info("test"); + throw new RuntimeException("test"); +// Decimal128 decimal = (Decimal128) value; +// if (decimal.compareTo(Decimal128.) == 0) { +// type.writeLong(output, encodeShortScaledValue(BigDecimal.ZERO, decimalType.getScale())); +// } +// else { +// type.writeLong(output, encodeShortScaledValue(decimal.bigDecimalValue(), decimalType.getScale())); +// } + } + else if (type.equals(DateType.DATE)) { + type.writeLong(output, Long.valueOf(value.toString())); + } + else { + throw new RuntimeException("Unsupported type: " + type); + } + } + else if (javaType == Int128.class) { + DecimalType decimalType = (DecimalType) type; + if (value instanceof Integer intValue) { + if (intValue == 0) { + type.writeObject(output, Decimals.encodeScaledValue(BigDecimal.ZERO, decimalType.getScale())); + } + else { + type.writeObject(output, Decimals.encodeScaledValue(BigDecimal.valueOf(intValue), decimalType.getScale())); + } + } + else if (value instanceof Double doubleValue) { + if (doubleValue == 0.0d) { + type.writeObject(output, Decimals.encodeScaledValue(BigDecimal.ZERO, decimalType.getScale())); + } + else { + BigDecimal result = new BigDecimal(BigInteger.valueOf(doubleValue.longValue())); + type.writeObject(output, Decimals.encodeScaledValue(result, decimalType.getScale())); + } + } + else { + throw new RuntimeException("Unsupported type: " + value.getClass()); + } + } + else if (javaType == Double.class || javaType == double.class) { + type.writeDouble(output, ((Number) value).doubleValue()); + } + else if (javaType == Block.class || javaType == SqlMap.class || javaType == SqlRow.class) { + writeBlock(output, type, value); + } + else { + throw new RuntimeException("Unsupported type " + javaType); + } + return; + } + catch (Exception e) { + throw new RuntimeException(String.format("Failed to append value '%s' of type %s from object type %s", String.valueOf(value), type, value.getClass()), e); + } + } + + private void writeBlock(BlockBuilder output, Type type, Object valueArg) + { + final Object value; + if (valueArg instanceof JsonObject document) { + value = document.toMap(); + } + else if (valueArg instanceof JsonArray arr) { + value = arr.toList(); + } else { + value = valueArg; + } + if (type instanceof ArrayType arrayType) { + if (value instanceof List list) { + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> list.forEach(element -> appendTo(arrayType.getElementType(), element, elementBuilder))); + return; + } + } + else if (type instanceof MapType mapType) { + if (value instanceof List) { + ((MapBlockBuilder) output).buildEntry((keyBuilder, valueBuilder) -> { + for (Object element : (List) value) { + if (!(element instanceof Map document)) { + continue; + } + + if (document.containsKey("key") && document.containsKey("value")) { + appendTo(mapType.getKeyType(), document.get("key"), keyBuilder); + appendTo(mapType.getValueType(), document.get("value"), valueBuilder); + } + } + }); + return; + } + if (value instanceof Map document) { + ((MapBlockBuilder) output).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : document.entrySet()) { + appendTo(mapType.getKeyType(), entry.getKey(), keyBuilder); + appendTo(mapType.getValueType(), entry.getValue(), valueBuilder); + } + }); + return; + } + } + else if (type instanceof RowType rowType) { + List fields = rowType.getFields(); + if (value instanceof Map mapValue) { + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int i = 0; i < fields.size(); i++) { + RowType.Field field = fields.get(i); + String fieldName = field.getName().orElse("field" + i); + appendTo(field.getType(), mapValue.get(fieldName), fieldBuilders.get(i)); + } + }); + return; + } + if (value instanceof List listValue) { + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int index = 0; index < fields.size(); index++) { + if (index < listValue.size()) { + appendTo(fields.get(index).getType(), listValue.get(index), fieldBuilders.get(index)); + } + else { + fieldBuilders.get(index).appendNull(); + } + } + }); + return; + } + } + else { + throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Block: " + type.getDisplayName()); + } + + // not a convertible value + output.appendNull(); + } + + private void appendTo(Type elementType, Object element, BlockBuilder elementBuilder) + { + appendValue(elementBuilder, elementType, element); + } + + @Override + public long getCompletedBytes() + { + return 0; + } + + @Override + public long getReadTimeNanos() + { + return 0; + } + + @Override + public boolean isFinished() + { + return finished; + } + + @Override + public long getMemoryUsage() + { + return 0; + } + + @Override + public void close() + throws IOException + { + } + + public ConnectorTransactionHandle transaction() + { + return transaction; + } + + public ConnectorSession session() + { + return session; + } + + public ConnectorSplit split() + { + return split; + } + + public ConnectorTableHandle table() + { + return table; + } + + @Override + public boolean equals(Object obj) + { + if (obj == this) { + return true; + } + if (obj == null || obj.getClass() != this.getClass()) { + return false; + } + var that = (CouchbasePageSource) obj; + return Objects.equals(this.transaction, that.transaction) && Objects.equals(this.session, that.session) && Objects.equals(this.split, that.split) && Objects.equals(this.table, that.table); + } + + @Override + public int hashCode() + { + return Objects.hash(transaction, session, split, table); + } + + @Override + public String toString() + { + return "CouchbasePageSource[" + "transaction=" + transaction + ", " + "session=" + session + ", " + "split=" + split + ", " + "table=" + table + ']'; + } + + @Override + public Metrics getMetrics() + { + return new Metrics(ImmutableMap.of( + "rows", new LongCount(total), + "duration", new LongCount(duration), + "pages", new LongCount(pageCount), + "errors", new LongCount(errors))); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSourceProvider.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSourceProvider.java new file mode 100644 index 000000000000..9763df3ae646 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePageSourceProvider.java @@ -0,0 +1,54 @@ +/* + * 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.couchbase; + +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.ConnectorPageSourceProvider; +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.DynamicFilter; +import jakarta.inject.Inject; + +import java.util.List; + +public class CouchbasePageSourceProvider + implements ConnectorPageSourceProvider +{ + private final CouchbaseClient client; + private final CouchbaseConfig config; + + @Inject + public CouchbasePageSourceProvider(CouchbaseConfig config, CouchbaseClient client) + { + this.config = config; + this.client = client; + } + + @Override + public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List columns, DynamicFilter dynamicFilter) + { + return new CouchbasePageSource( + client, + (CouchbaseTransactionHandle) transaction, + session, + (CouchbaseSplit) split, + (CouchbaseTableHandle) table, + columns.stream().map(CouchbaseColumnHandle.class::cast).toList(), + dynamicFilter, + config.getPageSize()); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePlugin.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePlugin.java new file mode 100644 index 000000000000..99204c662c57 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbasePlugin.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.couchbase; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; +import io.trino.spi.Plugin; +import io.trino.spi.connector.ConnectorFactory; + +import static java.util.Objects.requireNonNull; + +public class CouchbasePlugin + implements Plugin +{ + private ConnectorFactory connectorFactory; + + public CouchbasePlugin() + { + connectorFactory = new CouchbaseConnectorFactory(); + } + + @VisibleForTesting + CouchbasePlugin(ConnectorFactory connectorFactory) + { + connectorFactory = requireNonNull(connectorFactory, "factory is null"); + } + + @Override + public Iterable getConnectorFactories() + { + return ImmutableList.of(connectorFactory); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplit.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplit.java new file mode 100644 index 000000000000..b36c68027ac0 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplit.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.couchbase; + +import io.trino.spi.connector.ConnectorSplit; + +public record CouchbaseSplit( +) implements ConnectorSplit +{ +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplitManager.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplitManager.java new file mode 100644 index 000000000000..bd2fff866acd --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseSplitManager.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.couchbase; + +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplitManager; +import io.trino.spi.connector.ConnectorSplitSource; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.Constraint; +import io.trino.spi.connector.DynamicFilter; +import io.trino.spi.connector.FixedSplitSource; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; + +public class CouchbaseSplitManager + implements ConnectorSplitManager +{ + @Override + public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableHandle table, DynamicFilter dynamicFilter, Constraint constraint) + { + return new FixedSplitSource(new CouchbaseSplit()); + } + + @Override + public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableFunctionHandle function) + { + return new FixedSplitSource(new CouchbaseSplit()); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTableHandle.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTableHandle.java new file mode 100644 index 000000000000..69b33a7446d7 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTableHandle.java @@ -0,0 +1,531 @@ +/* + * 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.couchbase; + +import com.google.common.collect.Streams; +import io.airlift.slice.Slice; +import io.trino.plugin.couchbase.translations.TrinoExpressionToCb; +import io.trino.plugin.couchbase.translations.TrinoToCbType; +import io.trino.spi.connector.AggregateFunction; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.Constraint; +import io.trino.spi.connector.SortItem; +import io.trino.spi.connector.SortOrder; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.expression.Variable; +import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.SortedRangeSet; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.Type; +import jakarta.validation.constraints.NotNull; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static java.util.stream.Collectors.joining; + +public record CouchbaseTableHandle(String schema, String name, Optional subQuery, + List selectClauses, List selectTypes, + List selectNames, List whereClauses, + TupleDomain constraint, + LinkedHashMap orderClauses, + Set groupings, + AtomicBoolean isAggregated, + AtomicLong topNCount) implements ConnectorTableHandle +{ + public String path() + { + return String.format("`%s`.`%s`", schema, name); + } + + public static CouchbaseTableHandle fromSchemaAndName(String schema, String name) + { + return new CouchbaseTableHandle( + schema, + name, + Optional.empty(), + new ArrayList<>(), + new ArrayList<>(), + new ArrayList<>(), + new ArrayList<>(), + TupleDomain.all(), + new LinkedHashMap<>(), + new HashSet<>(), + new AtomicBoolean(false), + new AtomicLong(-1L)); + } + + public void addSortItems(List sortItems, Map assignments) + { + CouchbaseTableHandle sq = subQuery.orElse(null); + List pushdown = new ArrayList<>(); + sortItems.forEach(sortItem -> { + CouchbaseColumnHandle sourceColumn = (CouchbaseColumnHandle) assignments.get(sortItem.getName()); + orderClauses.put(transformSortItem(sortItem, assignments), sourceColumn); + }); + if (!pushdown.isEmpty()) { + sq.addSortItems(pushdown, assignments); + } + } + + protected String transformSortItem(SortItem sortItem, Map assignments) + { + CouchbaseColumnHandle column = (CouchbaseColumnHandle) assignments.get(sortItem.getName()); + return String.format("%s %s", column.name(), sortItem.getSortOrder().toString()); + } + + public boolean compareSortItems(List sortItems, Map assignments) + { + if (this.orderClauses.size() != sortItems.size()) { + return false; + } + return sortItems.stream().map(si -> transformSortItem(si, assignments)).allMatch(orderClauses.keySet()::contains); + } + + public boolean hasVariable(String name) + { + return selectClauses.stream() + .anyMatch(c -> { + if (c.name() == null) { + String text = c.value().text(); + return c.value().params().isEmpty() && + text.startsWith(name, 1) && + text.length() == name.length() + 2 && + text.endsWith("`"); + } + return c.name().equals(name); + }); + } + + public void setTopNCount(long topNCount) + { + this.topNCount.set(topNCount); + } + + public List addProjections(List projections, Map assignments) + { + return projections.stream().map(p -> addProjection(p, assignments)).toList(); + } + + private String addProjection(ConnectorExpression projection, Map assignments) + { + NamedParametrizedString compiled = compileProjection(projection, assignments); + if (!selectClauses.contains(compiled)) { + String otherName = findName(compiled.value()).orElse(null); + if (otherName == null) { + if (compiled.name() == null) { + compiled = new NamedParametrizedString(generateColumnName(), compiled.value()); + } else if (hasVariable(compiled.name())) { + return compiled.name(); + } + selectClauses.add(compiled); + selectTypes.add(projection.getType()); + selectNames.add(compiled.name()); + } + else { + return otherName; + } + } + return compiled.name(); + } + + private Optional findName(ParametrizedString value) + { + return selectClauses.stream().filter(nps -> nps.value().equals(value)).findFirst().map(nps -> nps.name()); + } + + private NamedParametrizedString compileProjection(ConnectorExpression projection, Map assignments) + { + if (projection instanceof Variable variable) { + ParametrizedString compiled = TrinoExpressionToCb.convert(projection, assignments); + return new NamedParametrizedString(variable.getName(), compiled); + } + else { + ParametrizedString compiled = TrinoExpressionToCb.convert(projection, assignments); + return new NamedParametrizedString(null, compiled); + } + } + + public String toSql() + { + List fromClause = new ArrayList<>(); + boolean fromSubQuery = false; + if (subQuery.isPresent()) { + CouchbaseTableHandle sq = subQuery.get(); + if (this.topNCount.get() < 0 && this.whereClauses().isEmpty() && this.orderClauses.isEmpty() && sq.selectClauses().containsAll(this.selectClauses) && this.groupings.isEmpty()) { + return sq.toSql(); + } + if (sq != this && (sq.schema().equals(schema) && sq.name().equals(name))) { + fromClause.add(String.format("(%s) `%s`", sq.toSql(), "data")); +// selectClauses.add(new NamedParametrizedString("data", ParametrizedString.from(String.format("`%s`.*", "data")))); + fromSubQuery = true; + } + } + if (fromClause.isEmpty()) { + fromClause.add(String.format("`%s`", name)); + } + + StringBuilder groupByClause = new StringBuilder(); + if (!groupings.isEmpty()) { + groupByClause.append(groupings.stream() + .map(CouchbaseColumnHandle::name) + .collect(Collectors.joining("`, `", " GROUP BY `", "`")) + ); + } + + StringBuilder orderByClause = new StringBuilder(); + if (!orderClauses.isEmpty()) { + orderByClause.append(String.format(" ORDER BY %s", String.join(", ", orderClauses.keySet()))); + if (!fromSubQuery) { + orderByClause.append(", META().id"); + } + } else if (!fromSubQuery) { + orderByClause.append(" ORDER BY META().id"); + } + + StringBuilder whereClause = new StringBuilder(); + if (!whereClauses.isEmpty()) { + whereClause.append(String.format(" WHERE %s", + whereClauses.stream().map(ParametrizedString::toString).collect(joining(" AND ")))); + } + + + String query = String.format("SELECT %s FROM %s%s%s%s", + selectClauses.isEmpty() ? String.format("`%s`.*", subQuery.isPresent() ? "data" : name()): + selectClauses.stream().map(NamedParametrizedString::toString).collect(joining(", ")), + String.join(", ", fromClause), + whereClause.toString(), + groupByClause.toString(), + orderByClause.toString()); + + if (topNCount.get() > -1) { + query = String.format("%s LIMIT %d", query, topNCount.get()); + } + + return query; + } + + private Stream getParametrizedStrings() + { + return Streams.concat( + selectClauses.stream().map(NamedParametrizedString::value), + subQuery.stream().flatMap(CouchbaseTableHandle::getParametrizedStrings), + whereClauses.stream()); + } + + public List getParameters() + { + return getParametrizedStrings().flatMap(p -> p.params().stream()).collect(Collectors.toList()); + } + + public boolean isEmpty() + { + return (subQuery.isEmpty() || (subQuery.get() != this && subQuery.get().isEmpty())) && + topNCount.get() == -1 && selectClauses.isEmpty() && whereClauses.isEmpty() && orderClauses.isEmpty() && groupings.isEmpty(); + } + + private ParametrizedString compileDomain(String left, Domain domain) + { + if (domain.isOnlyNull()) { + return ParametrizedString.from(String.format("%s IS NULL", left)); + } + else if (domain.isSingleValue()) { + return ParametrizedString.from(String.format("%s = ?", left), Arrays.asList(TrinoToCbType.serialize(domain.getType(), domain.getSingleValue()))); + } + else if (domain.getValues() instanceof SortedRangeSet rangeSet) { + List ranges = new ArrayList<>(); + if (rangeSet.isDiscreteSet()) { + List values = rangeSet.getDiscreteSet(); + List include = new ArrayList<>(); + List exclude = new ArrayList<>(); + boolean[] inclusives = rangeSet.getInclusive(); + + for (int i = 0; i < values.size(); i++) { + Object value = values.get(i); + if (value instanceof Slice slice) { + value = slice.toStringUtf8(); + } + if (inclusives[i]) { + include.add(value); + } + else { + exclude.add(value); + } + } + + if (!include.isEmpty()) { + ParametrizedString includePs = ParametrizedString.join(include.stream().map(v -> new ParametrizedString("?", Arrays.asList(v))).collect(Collectors.toUnmodifiableList()), ", ", String.format("%s IN [", left), "]"); + if (domain.isNullableDiscreteSet()) { + includePs = ParametrizedString.join( + Arrays.asList(includePs, ParametrizedString.from(String.format("%s IS NULL", left))), + ") OR (", "(", ")"); + } + ranges.add(includePs); + } + if (!exclude.isEmpty()) { + ranges.add(ParametrizedString.join(include.stream().map(v -> new ParametrizedString("?", Arrays.asList(v))).collect(Collectors.toUnmodifiableList()), ", ", String.format("%s NOT IN [", left), "]")); + } + return ParametrizedString.join(ranges, ") AND (", "(", ")"); + } + else { + rangeSet.getRanges().getOrderedRanges() + .stream().filter(range -> !range.isAll()) + .forEach(range -> { + List converted = new ArrayList<>(); + if (!range.isLowUnbounded()) { + String op = "%s > ?"; + if (range.isLowInclusive()) { + op = "%s >= ?"; + } + converted.add(ParametrizedString.from(String.format(op, left), Arrays.asList(TrinoToCbType.serialize(range.getType(), range.getLowValue())))); + } + if (!range.isHighUnbounded()) { + String op = "%s < ?"; + if (range.isHighInclusive()) { + op = "%s <= ?"; + } + converted.add(ParametrizedString.from(String.format(op, left), Arrays.asList(TrinoToCbType.serialize(range.getType(), range.getHighValue())))); + } + + ranges.add(ParametrizedString.join(converted, ") AND (", "(", ")")); + }); + if (domain.isNullAllowed()) { + ranges.add(ParametrizedString.from(String.format("%s IS NULL", left))); + return ParametrizedString.join(ranges, ") OR (", "(", ")"); + } else { + ParametrizedString nonNull = ParametrizedString.from(String.format("%s IS NOT NULL", left)); + if (ranges.isEmpty()) { + return nonNull; + } + + return ParametrizedString.join(List.of( + nonNull, + ParametrizedString.join(ranges, ") OR (", "(", ")") + ), " AND "); + } + } + } + else { + throw new RuntimeException("Unsupported domain type: " + domain.getClass().getName()); + } + } + + public boolean containsConstraint(Constraint constraint) + { + return whereClauses.contains(TrinoExpressionToCb.convert(constraint.getExpression(), constraint.getAssignments())); + } + + @Override + @NotNull + public String toString() + { + StringBuilder builder = new StringBuilder(); + builder.append(name); + if (subQuery().isPresent()) { + builder.append(String.format(" subQuery=[%s]", subQuery())); + } + if (!whereClauses.isEmpty()) { + builder.append(" filterPredicate=") + .append(ParametrizedString.join(whereClauses, ") AND (", "(", ")").text()); + } + if (constraint.isNone()) { + builder.append(" constraint=FALSE"); + } + else if (!constraint.isAll()) { + builder.append(" constraint on "); + builder.append(constraint.getDomains().orElseThrow().keySet().stream() + .map(columnHandle -> ((CouchbaseColumnHandle) columnHandle).name()) + .collect(joining(", ", "[", "]"))); + } + if (!selectNames.isEmpty()) { + builder.append(" columns=").append('[') + .append(String.join(", ", selectNames)) + .append(']'); + } + if (!orderClauses.isEmpty()) { + builder.append(" sortOrder=") + .append(orderClauses.keySet().stream() + .map(s -> { + int space = s.indexOf(' '); + if (space > -1) { + return String.format("%s:%s", s.substring(0, space), s.substring(space)); + } + return String.format("%s:", s); + }) + .collect(joining(", ", "[", "]"))); + } + if (topNCount.get() > -1) { + builder.append(" limit=").append(topNCount.get()); + } + return builder.toString(); + } + + public void addPredicate(TupleDomain predicate) + { + whereClauses.add(compilePredicate(predicate)); + } + + public ParametrizedString compilePredicate(TupleDomain predicate) + { + List clauses = new ArrayList<>(); + if (predicate.getDomains().isPresent() && !predicate.isAll()) { + Map domains = predicate.getDomains().get(); + if (!domains.isEmpty()) { + domains.forEach((column, domain) -> { + if (column instanceof CouchbaseColumnHandle cbcolumn) { + clauses.add(compileDomain(cbcolumn.name(), domain)); + } + else { + throw new IllegalArgumentException("Invalid column type: " + column.getClass().getName()); + } + }); + } + } + + if (clauses.isEmpty()) { + return ParametrizedString.from("TRUE"); + } + if (clauses.size() == 1) { + return clauses.get(0); + } + return ParametrizedString.join(clauses, ") AND (", "(", ")"); + } + + public void addColumns(Collection columns) + { + columns.forEach(this::addColumn); + } + + public void addColumn(CouchbaseColumnHandle column) + { + if (!coversColumn(column)) { + NamedParametrizedString nps = new NamedParametrizedString(column.name(), ParametrizedString.from(String.format("`%s`", column.name()))); + selectClauses.add(nps); + selectTypes.add(column.type()); + selectNames.add(column.name()); + subQuery.ifPresent(sq -> { + if (!sq.coversColumn(column)) { + sq.addColumn(column); + } + }); + } + } + + public CouchbaseTableHandle wrap() + { + return new CouchbaseTableHandle(schema(), name(), Optional.of(this), new ArrayList<>(), new ArrayList<>(), + new ArrayList<>(), new ArrayList<>(), TupleDomain.all(), new LinkedHashMap<>(), new HashSet<>(), + new AtomicBoolean(false), new AtomicLong(-1)); + } + + public boolean containsProjections(List projections, Map assignments) + { + return projections.stream() + .map(projection -> compileProjection(projection, assignments)) + .allMatch(projection -> selectClauses.stream() + .anyMatch(clause -> + (projection.name() != null && Objects.equals(clause.name(), projection.name())) || + (projection.name() == null && Objects.equals(clause.value(), projection.value())))); + } + + public CouchbaseTableHandle withConstraint(TupleDomain newDomain) + { + whereClauses.add(compilePredicate(newDomain)); + return new CouchbaseTableHandle(schema(), name(), subQuery, selectClauses, selectTypes, selectNames, whereClauses, + newDomain, orderClauses, groupings, new AtomicBoolean(false), topNCount); + } + + public boolean coversColumn(CouchbaseColumnHandle column) + { + return hasVariable(column.name()); + } + + public boolean containsAllAggregations(List aggregates, Map assignments) { + return aggregates.stream().allMatch(agg -> containsAggregation(agg, assignments)); + } + + public boolean containsAggregation(AggregateFunction aggregateFunction, Map assignments) { + return findAggregation(aggregateFunction, assignments).isPresent(); + } + + public Optional findAggregation(AggregateFunction aggregateFunction, + Map assignments) + { + ParametrizedString converted = TrinoExpressionToCb.convert(aggregateFunction, assignments); + return selectClauses.stream().filter(nps -> nps.value().equals(converted)).findFirst(); + } + + public NamedParametrizedString addAggregateFunction(AggregateFunction aggregate, Map assignments) + { + NamedParametrizedString result = new NamedParametrizedString( + generateColumnName(), + TrinoExpressionToCb.convert(aggregate, assignments)); + selectClauses.add(result); + selectNames.add(result.name()); + selectTypes.add(aggregate.getOutputType()); + isAggregated.set(true); + return result; + } + + private String generateColumnName() + { + return String.format("syn_column_%d", selectClauses.size()); + } + + public void clearSelectElements() + { + selectClauses.clear(); + selectNames.clear(); + selectTypes.clear(); + } + + public boolean containsAllGroupings(Collection> groupingSets) + { + return groupingSets.stream() + .flatMap(group -> group.stream().map(CouchbaseColumnHandle.class::cast)) + .allMatch(this.groupings::contains); + } + + public void addGroupings(Collection> groupingSets) + { + groupingSets.stream() + .flatMap(group -> group.stream().map(CouchbaseColumnHandle.class::cast)) + .forEach(chandle -> { + groupings.add(chandle); + if (!this.hasSortItemOn(chandle.name())) { + addSortItems(List.of(new SortItem(chandle.name(), SortOrder.ASC_NULLS_LAST)), Map.of(chandle.name(), chandle)); + } + }); + } + + public boolean hasSortItemOn(String columnName) + { + return this.orderClauses.values().stream() + .filter(Objects::nonNull) + .anyMatch(target -> target.name().equals(columnName)); + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTransactionHandle.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTransactionHandle.java new file mode 100644 index 000000000000..04154b2fb5b3 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/CouchbaseTransactionHandle.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.couchbase; + +import io.trino.spi.connector.ConnectorTransactionHandle; + +public enum CouchbaseTransactionHandle + implements ConnectorTransactionHandle +{ + INSTANCE +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/NamedParametrizedString.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/NamedParametrizedString.java new file mode 100644 index 000000000000..3cd9b14c0d0c --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/NamedParametrizedString.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.couchbase; + +import jakarta.annotation.Nullable; + +public record NamedParametrizedString(@Nullable String name, ParametrizedString value) +{ + @Override + public String toString() + { + if (name == null) { + return value.toString(); + } + else { + return String.format("%s `%s`", value, name); + } + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/ParametrizedString.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/ParametrizedString.java new file mode 100644 index 000000000000..8ec547135521 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/ParametrizedString.java @@ -0,0 +1,60 @@ +/* + * 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.couchbase; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +public record ParametrizedString(String text, List params) +{ + public static ParametrizedString from(String text) + { + return new ParametrizedString(text, Collections.emptyList()); + } + + public static ParametrizedString join(List others) { + return join(others, "", "", ""); + } + + public static ParametrizedString join(List others, String delimeter) { + return join(others, delimeter, "", ""); + } + + public static ParametrizedString join(List others, String delimeter, String before, String after) + { + List params = new ArrayList<>(); + return new ParametrizedString( + others.stream() + .peek(ps -> { + params.addAll(ps.params()); + }) + .map(ParametrizedString::toString) + .collect(Collectors.joining(delimeter, before, after)), + params); + } + + public static ParametrizedString from(String s, List params) + { + return new ParametrizedString( + s, params); + } + + @Override + public String toString() + { + return text; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/FunctionMappings.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/FunctionMappings.java new file mode 100644 index 000000000000..efc40684860f --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/FunctionMappings.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.couchbase.translations; + +public class FunctionMappings +{ + public static String trinoToCb(String functionName) { + // todo: validate all mappings + return functionName; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoExpressionToCb.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoExpressionToCb.java new file mode 100644 index 000000000000..05d6eaaf50c3 --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoExpressionToCb.java @@ -0,0 +1,196 @@ +/* + * 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.couchbase.translations; + +import ch.qos.logback.core.boolex.EvaluationException; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.plugin.couchbase.CouchbaseColumnHandle; +import io.trino.plugin.couchbase.NamedParametrizedString; +import io.trino.plugin.couchbase.ParametrizedString; +import io.trino.spi.connector.AggregateFunction; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.expression.Call; +import io.trino.spi.expression.ConnectorExpression; +import io.trino.spi.expression.Constant; +import io.trino.spi.expression.Variable; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.Int128; +import io.trino.spi.type.Type; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +public class TrinoExpressionToCb +{ + private TrinoExpressionToCb() + { + } + + public static List convert(List expressions, Map assignments) + { + return expressions.stream().map(e -> convert(e, assignments)).toList(); + } + + public static ParametrizedString convert(ConnectorExpression expression, Map assignments) + { + if (expression instanceof Call call) { + return processCall(call, assignments); + } + else if (expression instanceof Constant constant) { + // todo: support type casts + if (constant.getValue() instanceof Slice slice) { + return ParametrizedString.from("?", Arrays.asList(slice.toStringUtf8())); + } + return ParametrizedString.from("?", Arrays.asList(constant.getValue())); + } + else if (expression instanceof Variable variable) { + final String name = variable.getName(); + ColumnHandle assignment = assignments.get(name); + if (assignment == null) { + return ParametrizedString.from(String.format("`%s`", variable.getName())); + } + else if (assignment instanceof CouchbaseColumnHandle column) { + return ParametrizedString.from(String.format("`%s`", column.name())); + } + else { + throw new IllegalArgumentException("Unsupported column type: " + assignment); + } + } + throw new IllegalArgumentException("Unsupported expression type: " + expression.getClass()); + } + + private static ParametrizedString processCall(Call call, Map assignments) + { + final String fname = call.getFunctionName().getName(); + if ("$in".equals(fname)) { + ParametrizedString source = convert(call.getArguments().get(0), assignments); + ParametrizedString list = convert(call.getArguments().get(1), assignments); + return ParametrizedString.join( + Arrays.asList(source, list), " IN ", "", ""); + } + else if ("$array".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), ", ", "[", "]"); + } + else if ("$cast".equals(fname)) { + final String castFn = String.format("%s(", getCastFn(call.getType())); + return ParametrizedString.join( + convert(call.getArguments(), assignments), + String.format(", %s", castFn), + castFn, + ")"); + } + else if ("$equal".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), " = ", "", ""); + } + else if ("$modulus".equals(fname) || "mod".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " % ", "", ""); + } + else if ("$add".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " + ", "", ""); + } + else if ("$negate".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " + ", "-(", ")"); + } + else if ("$greater_than_or_equal".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " >= ", "", ""); + } + else if ("$less_than_or_equal".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " <= ", "", ""); + } + else if ("$less_than".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " < ", "", ""); + } + else if ("$greater_than".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " > ", "", ""); + } + else if ("$not".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + ") OR (", "NOT ((", "))"); + } + else if ("$not_equal".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " != ", "", ""); + } + else if ("$nullif".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + ", ", "NULLIF(", ")"); + } + else if ("$like".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + " LIKE ", "", ""); + } + else if ("random".equals(fname)) { + return ParametrizedString.join( + convert(call.getArguments(), assignments), + ", ", "RANDOM(", ")"); + } + else { + throw new IllegalArgumentException("Unsupported function: " + call.getFunctionName()); + } + } + + private static String getCastFn(Type type) + { + if (type.getJavaType() == Int128.class + || type instanceof DoubleType) { + return "to_number"; + } + else { + throw new IllegalArgumentException("Unsupported type: " + type); + } + } + + public static ParametrizedString convert(AggregateFunction aggregate, Map assignments) + { + ParametrizedString result; + if (aggregate.getFunctionName().equals("count") && aggregate.getArguments().isEmpty()) { + result = ParametrizedString.from("count(*)"); + } else { + result = ParametrizedString.join(Arrays.asList( + ParametrizedString.from(FunctionMappings.trinoToCb(aggregate.getFunctionName())), + ParametrizedString.join(convert(aggregate.getArguments(), assignments), ",", "(", ")") + )); + } + if (aggregate.isDistinct()) { + result = ParametrizedString.join(Arrays.asList( + ParametrizedString.from("DISTINCT"), + result + ), " "); + } + + return result; + } +} diff --git a/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoToCbType.java b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoToCbType.java new file mode 100644 index 000000000000..4209bedd39de --- /dev/null +++ b/plugin/trino-couchbase/src/main/java/io/trino/plugin/couchbase/translations/TrinoToCbType.java @@ -0,0 +1,90 @@ +/* + * 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.couchbase.translations; + +import com.google.common.collect.ImmutableSet; +import io.airlift.slice.Slice; +import io.trino.spi.type.BigintType; +import io.trino.spi.type.CharType; +import io.trino.spi.type.DateType; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.IntegerType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import jakarta.annotation.Nullable; + +import java.util.Optional; +import java.util.Set; + +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +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.TimeType.TIME_MILLIS; +import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS; +import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; +import static io.trino.spi.type.TinyintType.TINYINT; + +public final class TrinoToCbType +{ + private TrinoToCbType() + { + } + + private static final Set PUSHDOWN_SUPPORTED_PRIMITIVE_TYPES = ImmutableSet.of( + BOOLEAN, + TINYINT, + SMALLINT, + INTEGER, + BIGINT, + REAL, + DOUBLE, + DATE, + TIME_MILLIS, + TIMESTAMP_MILLIS, + TIMESTAMP_TZ_MILLIS); + + @Nullable + public static Object serialize(Type type, Object value) + { + if (value instanceof Optional optional) { + value = optional.orElse(null); + } + if (value == null) { + return null; + } + if (type == DateType.DATE || type.equals(BigintType.BIGINT) || type.equals(DoubleType.DOUBLE) || type.equals(INTEGER)) { + return value; + } + else if (type instanceof VarcharType) { + Slice slice = (Slice) value; + return slice.toStringUtf8(); + } + else { + throw new RuntimeException("Unsupported domain value type: " + type); + } + } + + public static boolean isPushdownSupportedType(Type type) + { + return type instanceof CharType + || type instanceof VarcharType + || type instanceof DecimalType + || PUSHDOWN_SUPPORTED_PRIMITIVE_TYPES.contains(type); + } +} diff --git a/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseConnectorTest.java b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseConnectorTest.java new file mode 100644 index 000000000000..0e9cfb8a4626 --- /dev/null +++ b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseConnectorTest.java @@ -0,0 +1,106 @@ +/* + * 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.couchbase; + +import io.trino.testing.BaseConnectorTest; +import io.trino.testing.QueryRunner; +import io.trino.testing.TestingConnectorBehavior; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; + +@TestInstance(PER_CLASS) +public class CouchbaseConnectorTest + extends BaseConnectorTest +{ + public static final String CBBUCKET = "trino-test"; + + private CouchbaseServer server; + + public CouchbaseConnectorTest() + { + this.server = new CouchbaseServer(CBBUCKET); + } + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return CouchbaseQueryRunner.builder(server).addInitialTables(REQUIRED_TPCH_TABLES).build(); + } + + @Override + protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) + { + return switch (connectorBehavior) { + case SUPPORTS_ADD_COLUMN, SUPPORTS_ARRAY, SUPPORTS_COMMENT_ON_TABLE, SUPPORTS_CREATE_SCHEMA, + SUPPORTS_CREATE_TABLE, SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_DELETE, SUPPORTS_INSERT, + SUPPORTS_MAP_TYPE, SUPPORTS_ROW_TYPE, SUPPORTS_NEGATIVE_DATE, SUPPORTS_JOIN_PUSHDOWN, + SUPPORTS_RENAME_COLUMN, SUPPORTS_RENAME_TABLE, SUPPORTS_SET_COLUMN_TYPE, + SUPPORTS_PREDICATE_EXPRESSION_PUSHDOWN, SUPPORTS_UPDATE, SUPPORTS_CREATE_VIEW, SUPPORTS_MERGE -> false; + + default -> super.hasBehavior(connectorBehavior); + }; + } + + @Test + @Override + public void testTopNPushdown() + { + super.testTopNPushdown(); + } + + @Test + public void testFailingItems() + { + assertQuery("SELECT regionkey, max(name) FROM nation GROUP BY regionkey LIMIT 5"); + } + + @Override + protected List largeInValuesCountData() + { + return List.of(20, 50, 100); + } + + @Test + @Override + public void testSelectInformationSchemaColumns() + { + Assumptions.abort("Skipping testSelectInformationSchemaColumns"); + } + + @Test + public void testInClauseDecoding() { + query("SELECT name from nation where name = 'FRANCE'") + .assertThat() + .isFullyPushedDown().result() + .rowCount().isEqualTo(1); + + query("SELECT name from nation where name = 'CANADA'") + .assertThat() + .isFullyPushedDown().result() + .rowCount().isEqualTo(1); + + query("SELECT name from nation where name = 'FRANCE' OR name = 'CANADA'") + .assertThat() + .isFullyPushedDown().result() + .rowCount().isEqualTo(2); + } + +} diff --git a/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseQueryRunner.java b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseQueryRunner.java new file mode 100644 index 000000000000..be6d8ead5aad --- /dev/null +++ b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseQueryRunner.java @@ -0,0 +1,261 @@ +/* + * 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.couchbase; + +import com.couchbase.client.core.error.CollectionExistsException; +import com.couchbase.client.java.Bucket; +import com.couchbase.client.java.Cluster; +import com.couchbase.client.java.Collection; +import com.couchbase.client.java.Scope; +import com.couchbase.client.java.json.JsonObject; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.airlift.log.Logger; +import io.airlift.log.Logging; +import io.trino.execution.QueryIdGenerator; +import io.trino.plugin.tpch.TpchPlugin; +import io.trino.spi.type.BigintType; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.IntegerType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.MaterializedResult; +import io.trino.testing.MaterializedRow; +import io.trino.testing.QueryRunner; +import io.trino.testing.TestingSession; +import io.trino.tpch.TpchTable; + +import java.io.File; +import java.io.IOException; +import java.io.Writer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.LocalDate; +import java.time.temporal.ChronoUnit; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static io.airlift.testing.Closeables.closeAllSuppress; +import static io.trino.plugin.couchbase.CouchbaseConnectorTest.CBBUCKET; +import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.DateType.DATE; +import static io.trino.tpch.TpchTable.CUSTOMER; +import static io.trino.tpch.TpchTable.NATION; +import static io.trino.tpch.TpchTable.ORDERS; +import static io.trino.tpch.TpchTable.REGION; +import static java.lang.String.format; + +public class CouchbaseQueryRunner +{ + private static final Logger log = Logger.get(CouchbaseQueryRunner.class); + private static final QueryIdGenerator queryIdGenerator = new QueryIdGenerator(); + public static final String TEST_SCHEMA = "tpch"; + private static final Path SCHEMA_DIR; + protected static final List> TPCH_TABLES = ImmutableList.of(NATION, ORDERS, REGION, CUSTOMER, TpchTable.LINE_ITEM, TpchTable.PART, TpchTable.PART_SUPPLIER, TpchTable.SUPPLIER); + + static { + try { + SCHEMA_DIR = Files.createTempDirectory("cbtestschema-"); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + private CouchbaseQueryRunner() + { + } + + public static Builder builder(CouchbaseServer server) + { + return new Builder(server) + .addConnectorProperty("couchbase.cluster", server.getConnectionString()) + .addConnectorProperty("couchbase.bucket", "trino-test") + .addConnectorProperty("couchbase.scope", "tpch") + .addConnectorProperty("couchbase.username", server.getUsername()) + .addConnectorProperty("couchbase.password", server.getPassword()) + .addConnectorProperty("couchbase.schema-folder", SCHEMA_DIR.toString()); + } + + public static final class Builder + extends DistributedQueryRunner.Builder + { + private final Map connectorProperties = new HashMap<>(); + private List> initialTables = ImmutableList.of(); + private CouchbaseServer server; + + public Builder(CouchbaseServer server) + { + super( + TestingSession.testSessionBuilder() + .setCatalog("couchbase") + .setSchema(TEST_SCHEMA) + .build()); + this.server = server; + } + + public Builder addConnectorProperty(String key, String value) + { + connectorProperties.put(key, value); + return this; + } + + public Builder addInitialTables(List> initialTables) + { + this.initialTables = initialTables; + return this; + } + + @Override + public DistributedQueryRunner build() + throws Exception + { + DistributedQueryRunner queryRunner = super.build(); + try { + queryRunner.installPlugin(new TpchPlugin()); + queryRunner.createCatalog("tpch", "tpch"); + + queryRunner.installPlugin(new CouchbasePlugin()); + queryRunner.createCatalog("couchbase", "couchbase", connectorProperties); + + log.info("Loading data from %s...", TEST_SCHEMA); + try (Cluster cluster = Cluster.connect(server.getConnectionString(), server.getUsername(), server.getPassword())) { + for (TpchTable table : initialTables) { + log.info("Running import for %s", table.getTableName()); + String tpchTableName = table.getTableName(); + MaterializedResult rows = queryRunner.execute(format("SELECT * FROM tpch.%s.%s", TINY_SCHEMA_NAME, tpchTableName)); + copyAndIngestTpchData(cluster, rows, server, table.getTableName(), connectorProperties); + generateTypeMappingFile(table.getTableName(), rows); + log.info("Imported %s rows for %s", rows.getRowCount(), table.getTableName()); + } + } + log.info("Loading into couchbase.%s complete", TEST_SCHEMA); + return queryRunner; + } + catch (Throwable e) { + closeAllSuppress(e, queryRunner); + throw e; + } + } + } + + private static void generateTypeMappingFile(String tableName, MaterializedResult rows) + { + try (Writer fw = Files.newBufferedWriter(Path.of(new File(SCHEMA_DIR.toFile(), String.format("%s.%s.%s.json", + CBBUCKET, + TEST_SCHEMA, tableName)).toURI()))) { + HashMap mappings = new HashMap<>(); + for (int i = 0; i < rows.getColumnNames().size(); i++) { + JsonObject mapping = JsonObject.create(); + mapping.put("type", rows.getTypes().get(i).getTypeSignature().jsonValue()); + mapping.put("order", i); + mappings.put(rows.getColumnNames().get(i), mapping); + } + JsonObject infer = JsonObject.from(mappings); + JsonObject propHolder = JsonObject.create(); + propHolder.put("properties", infer); + fw.write(propHolder.toString()); + log.info("Inferred JSON file for column %s", propHolder); + } + catch (Exception ex) { + throw new RuntimeException("Failed to generate INFER file", ex); + } + } + + private static void copyAndIngestTpchData(Cluster cluster, MaterializedResult rows, CouchbaseServer server, String tableName, Map connectorProperties) + { + String bucketName = CBBUCKET; + String scopeName = TEST_SCHEMA; + + Bucket bucket = cluster.bucket(bucketName); + try { + bucket.collections().createScope(scopeName); + } + catch (Exception _) { + // noop + } + Scope scope = bucket.scope(scopeName); + + try { + bucket.collections().createCollection(scopeName, tableName); + } + catch (CollectionExistsException _) { + // noop + } + + Collection target = scope.collection(tableName); + + List columns = rows.getColumnNames(); + List types = rows.getTypes(); + Long id = 0L; + for (MaterializedRow row : rows) { + JsonObject document = JsonObject.create(); + for (int i = 0; i < columns.size(); i++) { + String columnName = columns.get(i); + document.put(columnName, convertType(row.getField(i), types.get(i))); + } + target.upsert(String.valueOf(id++), document); + } + // let the dust settle + try { + Thread.sleep(1000L); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + + private static Object convertType(Object value, Type type) + { + if (value == null) { + return null; + } + + if (type == BOOLEAN + || type instanceof VarcharType + || type == IntegerType.INTEGER + || type == BigintType.BIGINT + || type == DoubleType.DOUBLE) { + return value; + } + else if (type == DATE) { + return ChronoUnit.DAYS.between(LocalDate.ofEpochDay(0), (LocalDate) value); + } + else { + throw new RuntimeException(String.format("Unsupported type: %s -- class: %s", type, value.getClass())); + } + } + + static void main() + throws Exception + { + Logging.initialize(); + QueryRunner queryRunner = builder(new CouchbaseServer("trino-test")) + .addInitialTables(TPCH_TABLES) + .addCoordinatorProperty("http-server.http.port", "8080") + .setExtraProperties(ImmutableMap.builder() + .put("sql.default-catalog", "tpch") + .put("sql.default-schema", "tiny") + .buildOrThrow()) + .withProtocolSpooling("json+zstd") + .build(); + Logger log = Logger.get(CouchbaseQueryRunner.class); + log.info("======== SERVER STARTED ========"); + log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); + } +} diff --git a/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseServer.java b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseServer.java new file mode 100644 index 000000000000..dd2354bc4236 --- /dev/null +++ b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/CouchbaseServer.java @@ -0,0 +1,52 @@ +/* + * 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.couchbase; + +import org.testcontainers.couchbase.BucketDefinition; +import org.testcontainers.couchbase.CouchbaseContainer; +import org.testcontainers.couchbase.CouchbaseService; +import org.testcontainers.utility.DockerImageName; + +public class CouchbaseServer + implements AutoCloseable +{ + private CouchbaseContainer container; + + public CouchbaseServer(String bucketName) + { + DockerImageName cbImage = DockerImageName.parse("couchbase:enterprise-8.0.0").asCompatibleSubstituteFor("couchbase/server"); + this.container = new CouchbaseContainer(cbImage).withBucket(new BucketDefinition(bucketName)).withEnabledServices(CouchbaseService.KV, CouchbaseService.INDEX, CouchbaseService.QUERY, CouchbaseService.SEARCH).withStartupAttempts(3); + this.container.start(); + } + + public String getConnectionString() + { + return container.getConnectionString(); + } + + public String getUsername() + { + return container.getUsername(); + } + + public String getPassword() + { + return container.getPassword(); + } + + public void close() + { + container.close(); + } +} diff --git a/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/TestCouchbaseConfig.java b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/TestCouchbaseConfig.java new file mode 100644 index 000000000000..099c1a30bfec --- /dev/null +++ b/plugin/trino-couchbase/src/test/java/io/trino/plugin/couchbase/TestCouchbaseConfig.java @@ -0,0 +1,84 @@ +/* + * 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.couchbase; + +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Map; + +import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; + +public class TestCouchbaseConfig +{ + @Test + void testDefaults() + { + assertRecordedDefaults( + recordDefaults(CouchbaseConfig.class) + .setCluster("localhost") + .setUsername("Administrator") + .setPassword("password") + .setBucket("default") + .setScope("_default") + .setTlsCertificate(null) + .setTlsKeyPassword(null) + .setTlsKey(null) + .setTimeouts("60") + .setSchemaFolder("couchbase-schema") + .setPageSize("10000")); + } + + @Test + void testExplicitPropertyMappings() + throws IOException + { + Path tls = Files.createTempFile(null, null); + Path keystoreFile = Files.createTempFile(null, null); + + Map properties = ImmutableMap.builder() + .put("couchbase.cluster", "remote-host") + .put("couchbase.username", "some-user") + .put("couchbase.password", "some-password") + .put("couchbase.bucket", "some-bucket") + .put("couchbase.scope", "some-scope") + .put("couchbase.tls-certificate", tls.toString()) + .put("couchbase.tls-key", keystoreFile.toString()) + .put("couchbase.tls-key-password", "some-keystore-password") + .put("couchbase.schema-folder", "some-folder") + .put("couchbase.timeouts", "10") + .put("couchbase.page-size", "1") + .buildOrThrow(); + + CouchbaseConfig expected = new CouchbaseConfig() + .setCluster("remote-host") + .setUsername("some-user") + .setPassword("some-password") + .setBucket("some-bucket") + .setScope("some-scope") + .setTlsCertificate(tls.toString()) + .setTlsKey(keystoreFile.toString()) + .setTlsKeyPassword("some-keystore-password") + .setSchemaFolder("some-folder") + .setTimeouts("10") + .setPageSize("1"); + + assertFullMapping(properties, expected); + } +} diff --git a/pom.xml b/pom.xml index 3393d27f16d6..5397f3cb6b4c 100644 --- a/pom.xml +++ b/pom.xml @@ -65,6 +65,7 @@ plugin/trino-blackhole plugin/trino-cassandra plugin/trino-clickhouse + plugin/trino-couchbase plugin/trino-datasketches plugin/trino-delta-lake plugin/trino-druid