Load Tests - Cassandra Reverse Replication

it/cassandra/pom.xml

@@ -0,0 +1,64 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+  ~ Copyright (C) 2025 Google Inc.
+  ~
+  ~ 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.
+  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~-->
+<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns="http://maven.apache.org/POM/4.0.0"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <groupId>com.google.cloud.teleport</groupId>
+        <artifactId>integration-testing-lib</artifactId>
+        <version>1.0-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>it-cassandra</artifactId>
+
+    <properties>
+        <parquet-avro.version>1.12.0</parquet-avro.version>
+    </properties>
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-it-conditions</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.auto.value</groupId>
+            <artifactId>auto-value</artifactId>
+            <version>${autovalue.version}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.google.auto.value</groupId>
+            <artifactId>auto-value-annotations</artifactId>
+            <version>${autovalue.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-it-cassandra</artifactId>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.datastax.cassandra</groupId>
+                    <artifactId>cassandra-driver-core</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+        <dependency>
+            <groupId>com.datastax.oss</groupId>
+            <artifactId>java-driver-core</artifactId>
+            <version>4.17.0</version>
+        </dependency>
+    </dependencies>
+</project>

it/cassandra/src/main/java/org/apache/beam/it/cassandra/conditions/CassandraRowsCheck.java

@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.beam.it.cassandra.conditions;
+
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.cql.ResultSet;
+import com.datastax.oss.driver.api.core.cql.Row;
+import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.google.auto.value.AutoValue;
+import java.net.InetSocketAddress;
+import java.time.Duration;
+import javax.annotation.Nullable;
+import org.apache.beam.it.cassandra.CassandraResourceManager;
+import org.apache.beam.it.conditions.ConditionCheck;
+
+/**
+ * Condition check for verifying the number of rows in a Cassandra table. This class is generic,
+ * allowing any type of Cassandra resource manager to be used at runtime.
+ */
+@AutoValue
+public abstract class CassandraRowsCheck extends ConditionCheck {
+
+  @Nullable
+  abstract CassandraResourceManager resourceManager();
+
+  abstract String tableName();
+
+  abstract Integer minRows();
+
+  @Nullable
+  abstract Integer maxRows();
+
+  @Override
+  public String getDescription() {
+    if (maxRows() != null) {
+      return String.format(
+          "Cassandra table check if table %s has between %d and %d rows",
+          tableName(), minRows(), maxRows());
+    }
+    return String.format(
+        "Cassandra table check if table %s has at least %d rows", tableName(), minRows());
+  }
+
+  /**
+   * Gets the row count for the specified table using the given CassandraResourceManager.
+   *
+   * @param resourceManager The CassandraResourceManager to use for the query.
+   * @param tableName The name of the table to count rows from.
+   * @return The number of rows in the table.
+   */
+  private long getRowCount(CassandraResourceManager resourceManager, String tableName) {
+    if (resourceManager == null) {
+      throw new IllegalArgumentException("CassandraResourceManager must not be null.");
+    }
+    try (CqlSession session =
+        CqlSession.builder()
+            .addContactPoint(
+                new InetSocketAddress(resourceManager.getHost(), resourceManager.getPort()))
+            .withLocalDatacenter("datacenter1")
+            .build()) {
+
+      String query =
+          String.format("SELECT COUNT(*) FROM %s.%s", resourceManager.getKeyspaceName(), tableName);
+      SimpleStatement statement =
+          SimpleStatement.builder(query).setTimeout(Duration.ofSeconds(20)).build();
+      ResultSet resultSet = session.execute(statement);
+      Row row = resultSet.one();
+      if (row != null) {
+        return row.getLong(0);
+      } else {
+        throw new RuntimeException("Query did not return a result for table: " + tableName);
+      }
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to execute query on CassandraResourceManager", e);
+    }
+  }
+
+  @Override
+  public CheckResult check() {
+    long totalRows = getRowCount(resourceManager(), tableName());
+    if (totalRows < minRows()) {
+      return new CheckResult(
+          false,
+          String.format("Expected at least %d rows but found only %d", minRows(), totalRows));
+    }
+    if (maxRows() != null && totalRows > maxRows()) {
+      return new CheckResult(
+          false, String.format("Expected up to %d rows but found %d", maxRows(), totalRows));
+    }
+
+    if (maxRows() != null) {
+      return new CheckResult(
+          true,
+          String.format(
+              "Expected between %d and %d rows and found %d", minRows(), maxRows(), totalRows));
+    }
+
+    return new CheckResult(
+        true, String.format("Expected at least %d rows and found %d", minRows(), totalRows));
+  }
+
+  /**
+   * Builder for {@link CassandraRowsCheck}. Now allows setting the CassandraResourceManager at
+   * runtime.
+   */
+  public static Builder builder(String tableName) {
+    return new AutoValue_CassandraRowsCheck.Builder().setTableName(tableName);
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    public abstract Builder setResourceManager(CassandraResourceManager resourceManager);
+
+    public abstract Builder setTableName(String tableName);
+
+    public abstract Builder setMinRows(Integer minRows);
+
+    public abstract Builder setMaxRows(Integer maxRows);
+
+    abstract CassandraRowsCheck autoBuild();
+
+    public CassandraRowsCheck build() {
+      return autoBuild();
+    }
+  }
+}

it/cassandra/src/main/java/org/apache/beam/it/cassandra/conditions/package-info.java

@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 that contains reusable Cassandra conditions. */
+package org.apache.beam.it.cassandra.conditions;

pom.xml

@@ -831,5 +831,6 @@
     <module>plaintext-logging</module>
     <module>python</module>
     <module>yaml</module>
+      <module>it/cassandra</module>
   </modules>
 </project>

v2/spanner-to-sourcedb/pom.xml

@@ -106,6 +106,12 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.google.cloud.teleport</groupId>
+            <artifactId>it-cassandra</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.beam</groupId>
             <artifactId>beam-it-cassandra</artifactId>