Add DataStax Enterprise as a Cassandra test environment

plugin/trino-cassandra/src/test/java/io/trino/plugin/cassandra/CassandraServer.java

@@ -20,6 +20,7 @@
 import com.datastax.oss.driver.api.core.config.ProgrammaticDriverConfigLoaderBuilder;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.Resources;
 import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
@@ -32,6 +33,7 @@
 import java.net.InetSocketAddress;
 import java.nio.file.Path;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeoutException;
 
 import static com.datastax.oss.driver.api.core.config.DefaultDriverOption.CONTROL_CONNECTION_AGREEMENT_TIMEOUT;
@@ -67,21 +69,29 @@ public class CassandraServer
     public CassandraServer()
             throws Exception
     {
-        this("2.2");
+        this("cassandra:2.2");
     }
 
-    public CassandraServer(String cassandraVersion)
+    public CassandraServer(String imageName)
+            throws Exception
+    {
+        this(imageName, ImmutableMap.of(), "/etc/cassandra/cassandra.yaml", "cu-cassandra.yaml");
+    }
+
+    public CassandraServer(String imageName, Map<String, String> environmentVariables, String configPath, String configFileName)
             throws Exception
     {
         log.info("Starting cassandra...");
 
-        this.dockerContainer = new GenericContainer<>("cassandra:" + cassandraVersion)
+        this.dockerContainer = new GenericContainer<>(imageName)
                 .withExposedPorts(PORT)
-                .withCopyFileToContainer(forHostPath(prepareCassandraYaml()), "/etc/cassandra/cassandra.yaml");
+                .withCopyFileToContainer(forHostPath(prepareCassandraYaml(configFileName)), configPath)
+                .withEnv(environmentVariables)
+                .withStartupTimeout(java.time.Duration.ofMinutes(10));
         this.dockerContainer.start();
 
         ProgrammaticDriverConfigLoaderBuilder driverConfigLoaderBuilder = DriverConfigLoader.programmaticBuilder();
-        driverConfigLoaderBuilder.withDuration(REQUEST_TIMEOUT, java.time.Duration.ofSeconds(12));
+        driverConfigLoaderBuilder.withDuration(REQUEST_TIMEOUT, java.time.Duration.ofSeconds(30));
         driverConfigLoaderBuilder.withString(PROTOCOL_VERSION, ProtocolVersion.V3.name());
         driverConfigLoaderBuilder.withDuration(CONTROL_CONNECTION_AGREEMENT_TIMEOUT, java.time.Duration.ofSeconds(30));
         // allow the retrieval of metadata for the system keyspaces
@@ -111,18 +121,18 @@ public CassandraServer(String cassandraVersion)
         this.session = session;
     }
 
-    private static String prepareCassandraYaml()
+    private static String prepareCassandraYaml(String fileName)
             throws IOException
     {
-        String original = Resources.toString(getResource("cu-cassandra.yaml"), UTF_8);
+        String original = Resources.toString(getResource(fileName), UTF_8);
 
         Path tmpDirPath = createTempDirectory(null);
         Path dataDir = tmpDirPath.resolve("data");
         createDirectory(dataDir);
 
         String modified = original.replaceAll("\\$\\{data_directory\\}", dataDir.toAbsolutePath().toString());
 
-        File yamlFile = tmpDirPath.resolve("cu-cassandra.yaml").toFile();
+        File yamlFile = tmpDirPath.resolve(fileName).toFile();
         yamlFile.deleteOnExit();
         write(modified, yamlFile, UTF_8);
 

plugin/trino-cassandra/src/test/java/io/trino/plugin/cassandra/TestCassandraLatestConnectorSmokeTest.java

@@ -28,7 +28,7 @@ public class TestCassandraLatestConnectorSmokeTest
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        CassandraServer server = closeAfterClass(new CassandraServer("3.11.10"));
+        CassandraServer server = closeAfterClass(new CassandraServer("cassandra:3.11.10"));
         CassandraSession session = server.getSession();
         createTestTables(session, KEYSPACE, Timestamp.from(TIMESTAMP_VALUE.toInstant()));
         return createCassandraQueryRunner(server, ImmutableMap.of(), ImmutableMap.of(), REQUIRED_TPCH_TABLES);

plugin/trino-cassandra/src/test/java/io/trino/plugin/cassandra/TestDatastaxConnectorSmokeTest.java

@@ -0,0 +1,43 @@
+/*
+ * 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.cassandra;
+
+import com.google.common.collect.ImmutableMap;
+import io.trino.testing.QueryRunner;
+
+import java.sql.Timestamp;
+import java.util.Map;
+
+import static io.trino.plugin.cassandra.CassandraQueryRunner.createCassandraQueryRunner;
+import static io.trino.plugin.cassandra.CassandraTestingUtils.createTestTables;
+
+public class TestDatastaxConnectorSmokeTest
+        extends BaseCassandraConnectorSmokeTest
+{
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        CassandraServer server = closeAfterClass(new CassandraServer(
+                "datastax/dse-server:6.8.25",
+                Map.of(
+                        "DS_LICENSE", "accept",
+                        "DC", "datacenter1"),
+                "/config/cassandra.yaml",
+                "cassandra-dse.yaml"));
+        CassandraSession session = server.getSession();
+        createTestTables(session, KEYSPACE, Timestamp.from(TIMESTAMP_VALUE.toInstant()));
+        return createCassandraQueryRunner(server, ImmutableMap.of(), ImmutableMap.of(), REQUIRED_TPCH_TABLES);
+    }
+}

plugin/trino-cassandra/src/test/resources/cassandra-dse.yaml

@@ -0,0 +1,20 @@
+endpoint_snitch: SimpleSnitch
+
+commitlog_sync: periodic
+commitlog_sync_period_in_ms: 10000
+
+partitioner: org.apache.cassandra.dht.Murmur3Partitioner
+
+seed_provider:
+    - class_name: org.apache.cassandra.locator.SimpleSeedProvider
+      parameters:
+          - seeds: "127.0.0.1"
+
+native_transport_port: 9142
+
+read_request_timeout_in_ms: 30000
+range_request_timeout_in_ms: 30000
+write_request_timeout_in_ms: 30000
+cas_contention_timeout_in_ms: 30000
+truncate_request_timeout_in_ms: 60000
+request_timeout_in_ms: 30000