More CqlSession Replacements

client-java/instrumentation/src/main/java/org/evomaster/client/java/instrumentation/ExecutedCqlCommand.java

@@ -18,6 +18,16 @@ public class ExecutedCqlCommand implements Serializable {
      */
     private final String cqlCommand;
 
+    /**
+     * The keyspace the command targets, if it could be determined from the query text; null otherwise
+     */
+    private final String keyspaceName;
+
+    /**
+     * The table the command targets, if it could be determined from the query text; null otherwise
+     */
+    private final String tableName;
+
     /**
      * Whether the CQL command failed, for any reason
      */
@@ -28,8 +38,10 @@ public class ExecutedCqlCommand implements Serializable {
      */
     private final long executionTime;
 
-    public ExecutedCqlCommand(String cqlCommand, boolean threwCqlException, long executionTime) {
+    public ExecutedCqlCommand(String cqlCommand, String keyspaceName, String tableName, boolean threwCqlException, long executionTime) {
         this.cqlCommand = cqlCommand;
+        this.keyspaceName = keyspaceName;
+        this.tableName = tableName;
         this.threwCqlException = threwCqlException;
         this.executionTime = executionTime;
     }
@@ -38,6 +50,14 @@ public String getCqlCommand() {
         return cqlCommand;
     }
 
+    public String getKeyspaceName() {
+        return keyspaceName;
+    }
+
+    public String getTableName() {
+        return tableName;
+    }
+
     public boolean hasThrownCqlException() {
         return threwCqlException;
     }

client-java/instrumentation/src/main/java/org/evomaster/client/java/instrumentation/coverage/methodreplacement/thirdpartyclasses/CqlSessionClassReplacement.java

@@ -2,6 +2,7 @@
 
 import org.evomaster.client.java.instrumentation.ExecutedCqlCommand;
 import org.evomaster.client.java.instrumentation.coverage.methodreplacement.Replacement;
+import org.evomaster.client.java.instrumentation.coverage.methodreplacement.ThirdPartyCast;
 import org.evomaster.client.java.instrumentation.coverage.methodreplacement.ThirdPartyMethodReplacementClass;
 import org.evomaster.client.java.instrumentation.coverage.methodreplacement.UsageFilter;
 import org.evomaster.client.java.instrumentation.shared.ReplacementCategory;
@@ -10,30 +11,55 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 public class CqlSessionClassReplacement extends ThirdPartyMethodReplacementClass {
     private static final CqlSessionClassReplacement singleton = new CqlSessionClassReplacement();
 
-    public static final String CASSANDRA_FIND_STRING_SYNC = "cassandraExecuteStringSync";
+    public static final String CASSANDRA_EXECUTE_STRING_SYNC = "cassandraExecuteStringSync";
+    public static final String CASSANDRA_EXECUTE_STRING_POSITIONAL_VALUES_SYNC = "cassandraExecuteStringPositionalValuesSync";
+    public static final String CASSANDRA_EXECUTE_STRING_NAMED_VALUES_SYNC = "cassandraExecuteStringNamedValuesSync";
+    public static final String CASSANDRA_EXECUTE_STATEMENT_SYNC = "cassandraExecuteStatementSync";
+
+    private static final String RESULT_SET_CLASS = "com.datastax.oss.driver.api.core.cql.ResultSet";
+    private static final String STATEMENT_CLASS = "com.datastax.oss.driver.api.core.cql.Statement";
 
     @Override
     protected String getNameOfThirdPartyTargetClass() {
         return "com.datastax.oss.driver.api.core.CqlSession";
     }
 
-    @Replacement(type = ReplacementType.TRACKER, id = CASSANDRA_FIND_STRING_SYNC, usageFilter = UsageFilter.ANY, category = ReplacementCategory.CASSANDRA, castTo = "com.datastax.oss.driver.api.core.cql.ResultSet")
+    @Replacement(type = ReplacementType.TRACKER, id = CASSANDRA_EXECUTE_STRING_SYNC, usageFilter = UsageFilter.ANY, category = ReplacementCategory.CASSANDRA, castTo = RESULT_SET_CLASS)
     public static Object execute(Object cqlSession, String query) {
-        return handleCqlExecute(CASSANDRA_FIND_STRING_SYNC, cqlSession, query);
+        return handleCqlExecute(CASSANDRA_EXECUTE_STRING_SYNC, cqlSession, query, query);
+    }
+
+    @Replacement(type = ReplacementType.TRACKER, id = CASSANDRA_EXECUTE_STRING_POSITIONAL_VALUES_SYNC, usageFilter = UsageFilter.ANY, category = ReplacementCategory.CASSANDRA, castTo = RESULT_SET_CLASS)
+    public static Object execute(Object cqlSession, String query, Object... values) {
+        return handleCqlExecute(CASSANDRA_EXECUTE_STRING_POSITIONAL_VALUES_SYNC, cqlSession, query, query, values);
+    }
+
+    @Replacement(type = ReplacementType.TRACKER, id = CASSANDRA_EXECUTE_STRING_NAMED_VALUES_SYNC, usageFilter = UsageFilter.ANY, category = ReplacementCategory.CASSANDRA, castTo = RESULT_SET_CLASS)
+    public static Object execute(Object cqlSession, String query, Map<String, Object> values) {
+        return handleCqlExecute(CASSANDRA_EXECUTE_STRING_NAMED_VALUES_SYNC, cqlSession, query, query, values);
+    }
+
+    @Replacement(type = ReplacementType.TRACKER, id = CASSANDRA_EXECUTE_STATEMENT_SYNC, usageFilter = UsageFilter.ANY, category = ReplacementCategory.CASSANDRA, castTo = RESULT_SET_CLASS)
+    public static Object execute(Object cqlSession, @ThirdPartyCast(actualType = STATEMENT_CLASS) Object statement) {
+        return handleCqlExecute(CASSANDRA_EXECUTE_STATEMENT_SYNC, cqlSession, extractQueryText(statement), statement);
     }
 
-    private static Object handleCqlExecute(String id, Object cqlSession, String query) {
+    private static Object handleCqlExecute(String id, Object cqlSession, String queryForTracking, Object... invokeArgs) {
         long start = System.currentTimeMillis();
         try {
             Method executeMethod = retrieveExecuteMethod(id, cqlSession);
-            Object result = executeMethod.invoke(cqlSession, query);
+            Object result = executeMethod.invoke(cqlSession, invokeArgs);
             long end = System.currentTimeMillis();
             long executionTime = end - start;
-            ExecutedCqlCommand info = new ExecutedCqlCommand(query, false, executionTime);
+            TableReference ref = extractTableReference(queryForTracking);
+            ExecutedCqlCommand info = new ExecutedCqlCommand(queryForTracking, ref.keyspaceName, ref.tableName, false, executionTime);
             ExecutionTracer.addCqlInfo(info);
             return result;
         } catch (IllegalAccessException e) {
@@ -43,6 +69,72 @@ private static Object handleCqlExecute(String id, Object cqlSession, String quer
         }
     }
 
+    /**
+     * Matches the keyspace/table reference after FROM (SELECT/DELETE), INTO (INSERT), or
+     * UPDATE, e.g. "FROM ks.tbl", "INTO tbl", "UPDATE \"Ks\".\"Tbl\"".
+     */
+    private static final Pattern TABLE_REFERENCE_PATTERN = Pattern.compile(
+            "(?i)\\b(?:FROM|INTO|UPDATE)\\s+(\"[^\"]+\"|[A-Za-z_]\\w*)(?:\\s*\\.\\s*(\"[^\"]+\"|[A-Za-z_]\\w*))?"
+    );
+
+    /**
+     * Best-effort extraction of keyspace/table from the CQL text. Returns both fields as
+     * null when the query doesn't match a recognised SELECT/INSERT/UPDATE/DELETE shape
+     * (eg DDL, batches).
+     */
+    private static TableReference extractTableReference(String query) {
+        if (query == null) {
+            return new TableReference(null, null);
+        }
+        Matcher matcher = TABLE_REFERENCE_PATTERN.matcher(query);
+        if (!matcher.find()) {
+            return new TableReference(null, null);
+        }
+        String first = stripQuotes(matcher.group(1));
+        String second = matcher.group(2) != null ? stripQuotes(matcher.group(2)) : null;
+        // if there is a "a.b" qualifier, a is the keyspace and b is the table;
+        // otherwise the single identifier is the table, and the keyspace is the session default
+        return second != null ? new TableReference(first, second) : new TableReference(null, first);
+    }
+
+    private static String stripQuotes(String identifier) {
+        if (identifier.length() >= 2 && identifier.charAt(0) == '"' && identifier.charAt(identifier.length() - 1) == '"') {
+            return identifier.substring(1, identifier.length() - 1);
+        }
+        return identifier;
+    }
+
+    private static class TableReference {
+        final String keyspaceName;
+        final String tableName;
+
+        TableReference(String keyspaceName, String tableName) {
+            this.keyspaceName = keyspaceName;
+            this.tableName = tableName;
+        }
+    }
+
+    /**
+     * Statement is a generic driver type; only SimpleStatement exposes the original
+     * CQL text directly, while BoundStatement requires going through its PreparedStatement.
+     */
+    private static String extractQueryText(Object statement) {
+        try {
+            Method getQuery = statement.getClass().getMethod("getQuery");
+            return (String) getQuery.invoke(statement);
+        } catch (ReflectiveOperationException e) {
+            // not a SimpleStatement (eg BoundStatement) -- fall through
+        }
+        try {
+            Method getPreparedStatement = statement.getClass().getMethod("getPreparedStatement");
+            Object prepared = getPreparedStatement.invoke(statement);
+            Method getQuery = prepared.getClass().getMethod("getQuery");
+            return (String) getQuery.invoke(prepared);
+        } catch (ReflectiveOperationException e) {
+            return statement.toString(); // eg BatchStatement -- best effort
+        }
+    }
+
     private static Method retrieveExecuteMethod(String id, Object cqlSession){
         return getOriginal(singleton, id, cqlSession);
     }

client-java/instrumentation/src/test/java/org/evomaster/client/java/instrumentation/coverage/methodreplacement/thirdpartyclasses/CqlSessionClassReplacementTest.java

@@ -1,6 +1,9 @@
 package org.evomaster.client.java.instrumentation.coverage.methodreplacement.thirdpartyclasses;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.cql.BoundStatement;
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
+import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import org.evomaster.client.java.instrumentation.AdditionalInfo;
 import org.evomaster.client.java.instrumentation.ExecutedCqlCommand;
 import org.evomaster.client.java.instrumentation.staticstate.ExecutionTracer;
@@ -13,8 +16,11 @@
 
 import java.net.InetSocketAddress;
 import java.time.Duration;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -31,7 +37,8 @@ public class CqlSessionClassReplacementTest {
                     .withStartupTimeout(Duration.ofMinutes(2)));
 
     private static final String KEYSPACE = "testks";
-    private static final String TABLE = KEYSPACE + ".users";
+    private static final String TABLE_NAME = "users";
+    private static final String TABLE = KEYSPACE + "." + TABLE_NAME;
 
     @BeforeAll
     static void startCassandra() {
@@ -80,6 +87,8 @@ void testExecuteSelectIsTracked() {
 
         ExecutedCqlCommand cmd = commands.iterator().next();
         assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
         assertFalse(cmd.hasThrownCqlException());
         assertTrue(cmd.getExecutionTime() >= 0);
     }
@@ -98,6 +107,8 @@ void testExecuteInsertIsTracked() {
 
         ExecutedCqlCommand cmd = commands.iterator().next();
         assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
         assertFalse(cmd.hasThrownCqlException());
         assertTrue(cmd.getExecutionTime() >= 0);
     }
@@ -140,4 +151,92 @@ void testExecutingInitCassandraFlagSuppressesTracking() {
         assertEquals(1, additionalInfoList.size());
         assertTrue(additionalInfoList.get(0).getCqlInfoData().isEmpty());
     }
+
+    @Test
+    void testExecuteWithPositionalValuesIsTracked() {
+        String query = "INSERT INTO " + TABLE + " (id, name, age) VALUES (?, ?, ?)";
+
+        CqlSessionClassReplacement.execute(cqlSession, query, UUID.randomUUID(), "Dave", 40);
+
+        List<AdditionalInfo> additionalInfoList = ExecutionTracer.exposeAdditionalInfoList();
+        assertEquals(1, additionalInfoList.size());
+
+        Set<ExecutedCqlCommand> commands = additionalInfoList.get(0).getCqlInfoData();
+        assertEquals(1, commands.size());
+
+        ExecutedCqlCommand cmd = commands.iterator().next();
+        assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
+        assertFalse(cmd.hasThrownCqlException());
+        assertTrue(cmd.getExecutionTime() >= 0);
+    }
+
+    @Test
+    void testExecuteWithNamedValuesIsTracked() {
+        String query = "INSERT INTO " + TABLE + " (id, name, age) VALUES (:id, :name, :age)";
+        Map<String, Object> values = new HashMap<>();
+        values.put("id", UUID.randomUUID());
+        values.put("name", "Erin");
+        values.put("age", 22);
+
+        CqlSessionClassReplacement.execute(cqlSession, query, values);
+
+        List<AdditionalInfo> additionalInfoList = ExecutionTracer.exposeAdditionalInfoList();
+        assertEquals(1, additionalInfoList.size());
+
+        Set<ExecutedCqlCommand> commands = additionalInfoList.get(0).getCqlInfoData();
+        assertEquals(1, commands.size());
+
+        ExecutedCqlCommand cmd = commands.iterator().next();
+        assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
+        assertFalse(cmd.hasThrownCqlException());
+        assertTrue(cmd.getExecutionTime() >= 0);
+    }
+
+    @Test
+    void testExecuteWithSimpleStatementIsTracked() {
+        String query = "SELECT * FROM " + TABLE;
+        SimpleStatement statement = SimpleStatement.newInstance(query);
+
+        CqlSessionClassReplacement.execute(cqlSession, statement);
+
+        List<AdditionalInfo> additionalInfoList = ExecutionTracer.exposeAdditionalInfoList();
+        assertEquals(1, additionalInfoList.size());
+
+        Set<ExecutedCqlCommand> commands = additionalInfoList.get(0).getCqlInfoData();
+        assertEquals(1, commands.size());
+
+        ExecutedCqlCommand cmd = commands.iterator().next();
+        assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
+        assertFalse(cmd.hasThrownCqlException());
+        assertTrue(cmd.getExecutionTime() >= 0);
+    }
+
+    @Test
+    void testExecuteWithBoundStatementIsTracked() {
+        String query = "INSERT INTO " + TABLE + " (id, name, age) VALUES (?, ?, ?)";
+        // Preparing directly on the session so it is NOT intercepted by the replacement
+        PreparedStatement prepared = cqlSession.prepare(query);
+        BoundStatement bound = prepared.bind(UUID.randomUUID(), "Frank", 33);
+
+        CqlSessionClassReplacement.execute(cqlSession, bound);
+
+        List<AdditionalInfo> additionalInfoList = ExecutionTracer.exposeAdditionalInfoList();
+        assertEquals(1, additionalInfoList.size());
+
+        Set<ExecutedCqlCommand> commands = additionalInfoList.get(0).getCqlInfoData();
+        assertEquals(1, commands.size());
+
+        ExecutedCqlCommand cmd = commands.iterator().next();
+        assertEquals(query, cmd.getCqlCommand());
+        assertEquals(KEYSPACE, cmd.getKeyspaceName());
+        assertEquals(TABLE_NAME, cmd.getTableName());
+        assertFalse(cmd.hasThrownCqlException());
+        assertTrue(cmd.getExecutionTime() >= 0);
+    }
 }