feat(mysql-dual-conn, tidb-stmt-cache): exercise COM_STMT_RESET + prepared-stmt orphan-EXECUTE

mysql-dual-conn/src/main/java/com/example/mysqlreplicate/QueryController.java

@@ -3,8 +3,12 @@
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PathVariable;
 import org.springframework.web.bind.annotation.RestController;
 
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -56,4 +60,33 @@ public List<Map<String, Object>> queryOms() {
     public List<Map<String, Object>> queryCamunda() {
         return camundaJdbc.queryForList("SELECT 1 AS camunda_check");
     }
+
+    /**
+     * Re-executes a server-prepared statement {@code n} times on the SAME
+     * JDBC connection. With useServerPrepStmts=true + useCursorFetch=true,
+     * Connector/J 8.x opportunistically emits COM_STMT_RESET before each
+     * COM_STMT_EXECUTE after the first, to clear cursor / long-data state.
+     *
+     * During Keploy replay this exercises the synthetic-OK fallback added
+     * in keploy/keploy#4217 — without it, the unmocked COM_STMT_RESET would
+     * cascade into "Connection closing due to no matching mock found" and
+     * tear down the TCP connection.
+     */
+    @GetMapping("/api/oms/stmt-reset/{n}")
+    public List<Integer> stmtReset(@PathVariable("n") int n) {
+        return omsJdbc.execute((java.sql.Connection conn) -> {
+            List<Integer> values = new ArrayList<>(n);
+            try (PreparedStatement ps = conn.prepareStatement("SELECT ? AS v")) {
+                for (int i = 0; i < n; i++) {
+                    ps.setInt(1, i);
+                    try (ResultSet rs = ps.executeQuery()) {
+                        if (rs.next()) {
+                            values.add(rs.getInt(1));
+                        }
+                    }
+                }
+            }
+            return values;
+        });
+    }
 }

mysql-dual-conn/src/main/resources/application.properties

@@ -1,7 +1,12 @@
 server.port=8080
 
 # --- OMS DataSource (primary) ---
-datasource.oms.jdbc-url=jdbc:mysql://localhost:3306/myntra_oms?useSSL=false&allowPublicKeyRetrieval=true
+# useServerPrepStmts + cachePrepStmts + useCursorFetch force Connector/J 8.x
+# to issue COM_STMT_PREPARE / COM_STMT_EXECUTE (and COM_STMT_RESET between
+# re-executions on the same connection) instead of plain COM_QUERY.
+# This lets /api/oms/stmt-reset/{n} exercise the COM_STMT_RESET synthetic-OK
+# fallback added in keploy/keploy#4217.
+datasource.oms.jdbc-url=jdbc:mysql://localhost:3306/myntra_oms?useSSL=false&allowPublicKeyRetrieval=true&useServerPrepStmts=true&cachePrepStmts=true&useCursorFetch=true
 datasource.oms.username=omsAppUser
 datasource.oms.password=omsPassword
 datasource.oms.driver-class-name=com.mysql.cj.jdbc.Driver

tidb-stmt-cache/docker-compose.yml

@@ -0,0 +1,32 @@
+version: '3.8'
+
+# Minimal single-node TiDB for keploy e2e.
+#
+# Why this stack and not the full pingcap/tidb-docker-compose with PD + TiKV?
+# - The sample only exercises the SQL layer (MySQL wire protocol on :4000).
+#   Coordination / storage flakiness in a 3-container topology adds CI noise
+#   without buying us any matcher coverage.
+# - `--store=unistore` keeps everything in a single process backed by an
+#   in-memory storage engine. Boot time is ~5s vs ~30-45s for a full PD+TiKV
+#   stack. Data is volatile, which is exactly what we want for keploy CI.
+#
+# Pin: v8.5.x is the LTS line current at the time this sample was added.
+# Bump as new LTS lines ship; the matcher behaviour we're testing has been
+# stable across TiDB versions because it depends on the MySQL wire protocol.
+services:
+  tidb:
+    image: pingcap/tidb:v8.5.6
+    command:
+      - --store=unistore
+      - --path=""
+      - --host=0.0.0.0
+      - --advertise-address=tidb
+      - --log-level=error
+    ports:
+      - "4000:4000"   # MySQL wire protocol
+      - "10080:10080" # status / readiness
+    healthcheck:
+      test: ["CMD-SHELL", "wget -qO- http://localhost:10080/status || exit 1"]
+      interval: 2s
+      timeout: 2s
+      retries: 30

tidb-stmt-cache/pom.xml

@@ -0,0 +1,55 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>org.springframework.boot</groupId>
+        <artifactId>spring-boot-starter-parent</artifactId>
+        <version>3.2.5</version>
+        <relativePath/>
+    </parent>
+
+    <groupId>com.example</groupId>
+    <artifactId>tidb-stmt-cache</artifactId>
+    <version>0.0.1-SNAPSHOT</version>
+    <name>tidb-stmt-cache</name>
+    <description>E2E sample exercising keploy's MySQL prepared-statement orphan-EXECUTE matching against TiDB</description>
+
+    <properties>
+        <java.version>17</java.version>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-web</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-jdbc</artifactId>
+        </dependency>
+        <!--
+            mysql-connector-j 8.x speaks the MySQL wire protocol that TiDB
+            implements (TiDB :4000 is MySQL-protocol-compatible). The same
+            cachePrepStmts + useServerPrepStmts flags that work against MySQL
+            8.0 are what cause this sample to exercise the orphan-EXECUTE
+            code path against TiDB — see application.properties.
+        -->
+        <dependency>
+            <groupId>com.mysql</groupId>
+            <artifactId>mysql-connector-j</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.springframework.boot</groupId>
+                <artifactId>spring-boot-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
+</project>

tidb-stmt-cache/src/main/java/com/example/tidbstmtcache/DataSourceConfig.java

@@ -0,0 +1,77 @@
+package com.example.tidbstmtcache;
+
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.core.JdbcTemplate;
+
+/**
+ * Single HikariCP pool against TiDB :4000 with MySQL Connector/J flags
+ * that force the orphan-EXECUTE scenario this sample is designed around:
+ *
+ *   useServerPrepStmts=true   -- server-side prepared statements (stmtIDs)
+ *   cachePrepStmts=true       -- per-Connection client-side PS cache
+ *   prepStmtCacheSize >= 1    -- the cache must actually retain entries
+ *
+ * Pool sizing > 1 with HikariCP's LIFO eviction means sequential HTTP
+ * requests to /api/kv/{i} often land on the same physical connection.
+ * On a cache hit, Connector/J skips COM_STMT_PREPARE and emits only
+ * COM_STMT_EXECUTE using the cached server-side stmtID. The recorder's
+ * mock for that second EXECUTE is the orphan case keploy/keploy@b2e68adb
+ * is designed to handle (recordedPrepByConn miss -> expectedQuery="" ->
+ * param-alone fallback).
+ *
+ * TiDB is preferred over MySQL here because TiDB's prepared-statement
+ * cache semantics diverge subtly from MySQL across COM_RESET_CONNECTION,
+ * which is what surfaced this matcher bug downstream. MySQL 8 alone is
+ * unlikely to reproduce the orphan condition reliably in one record cycle.
+ */
+@Configuration
+public class DataSourceConfig {
+
+    @Value("${datasource.tidb.jdbc-url}")
+    private String jdbcUrl;
+
+    @Value("${datasource.tidb.username}")
+    private String username;
+
+    @Value("${datasource.tidb.password}")
+    private String password;
+
+    @Value("${datasource.tidb.driver-class-name}")
+    private String driverClass;
+
+    @Bean(destroyMethod = "close")
+    public HikariDataSource tidbDataSource() {
+        HikariConfig config = new HikariConfig();
+        config.setPoolName("tidb-dataSource");
+        config.setUsername(username);
+        config.setPassword(password);
+        config.setJdbcUrl(jdbcUrl);
+        config.setDriverClassName(driverClass);
+
+        // Small pool: enough to be realistic (not 1), small enough that
+        // sequential curls reliably hit the same physical connection and
+        // therefore the same Connector/J prepared-statement cache.
+        config.setMaximumPoolSize(3);
+        config.setMinimumIdle(1);
+
+        // Keep connections alive long enough to span the whole record
+        // window so HikariCP doesn't churn the pool mid-test and flush
+        // the prep cache out from under us.
+        config.setKeepaliveTime(30_000);
+        config.setIdleTimeout(60_000);
+        config.setMaxLifetime(7_200_000);
+        config.setConnectionTimeout(10_000);
+        config.setValidationTimeout(5_000);
+
+        return new HikariDataSource(config);
+    }
+
+    @Bean
+    public JdbcTemplate tidbJdbcTemplate(HikariDataSource tidbDataSource) {
+        return new JdbcTemplate(tidbDataSource);
+    }
+}

tidb-stmt-cache/src/main/java/com/example/tidbstmtcache/QueryController.java

@@ -0,0 +1,89 @@
+package com.example.tidbstmtcache;
+
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PathVariable;
+import org.springframework.web.bind.annotation.RestController;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Endpoints that drive prepared-statement traffic against TiDB :4000.
+ *
+ * The orphan-EXECUTE scenario being exercised:
+ *
+ *   1. /api/kv/{v} prepares "SELECT ? AS v" with useServerPrepStmts +
+ *      cachePrepStmts on.
+ *   2. First call on Connection-A: Connector/J emits COM_STMT_PREPARE
+ *      then COM_STMT_EXECUTE; recorder captures both (PREPARE mock +
+ *      EXECUTE mock with the SAME connID/stmtID pair).
+ *   3. Subsequent call on Connection-A (HikariCP LIFO): Connector/J
+ *      finds the PreparedStatement in its client cache and emits ONLY
+ *      COM_STMT_EXECUTE using the cached stmtID. Recorder captures an
+ *      EXECUTE-only mock.
+ *   4. At replay time, the matcher tries to pair the "EXECUTE-only" mock
+ *      against the incoming COM_STMT_EXECUTE. If the recorder's connID
+ *      attribution or HikariCP's pool rotation makes the PREPARE entry
+ *      invisible to buildRecordedPrepIndex for this stmtID, expectedQuery
+ *      comes back empty -- which is the case keploy/keploy@b2e68adb
+ *      handles by accepting the EXECUTE on parameters alone instead of
+ *      crashing the connection with "no matching mock".
+ *
+ * Two endpoints with different SQL shapes are exposed so the matcher
+ * gets nontrivial work to do (it cannot just memoize one stmtID).
+ */
+@RestController
+public class QueryController {
+
+    private final JdbcTemplate jdbc;
+
+    public QueryController(JdbcTemplate jdbc) {
+        this.jdbc = jdbc;
+    }
+
+    /**
+     * Lightweight liveness probe. Plain query, no prepared statement --
+     * used by the CI script's wait_for_app loop so app readiness is not
+     * coupled to TiDB prep-cache behaviour.
+     */
+    @GetMapping("/api/health")
+    public Map<String, Object> health() {
+        Integer one = jdbc.queryForObject("SELECT 1", Integer.class);
+        Map<String, Object> out = new HashMap<>();
+        out.put("status", "ok");
+        out.put("db", one);
+        return out;
+    }
+
+    /**
+     * Prepared SELECT with one parameter. Same SQL across calls, so
+     * Connector/J's cachePrepStmts cache hits on the second-and-later
+     * call landing on the same physical connection.
+     */
+    @GetMapping("/api/kv/{v}")
+    public Map<String, Object> selectParam(@PathVariable("v") int v) {
+        Integer echoed = jdbc.queryForObject("SELECT ? AS v", Integer.class, v);
+        Map<String, Object> out = new HashMap<>();
+        out.put("echoed", echoed);
+        return out;
+    }
+
+    /**
+     * Prepared INSERT, then prepared SELECT against the same row. Two
+     * distinct prepared statements ("INSERT INTO kv ..." and "SELECT v
+     * FROM kv WHERE id=?") that both go through the Connector/J cache.
+     * Gives the matcher more than one (connID, stmtID) pair to track
+     * concurrently per connection.
+     */
+    @GetMapping("/api/kv/insert-select/{v}")
+    public Map<String, Object> insertThenSelect(@PathVariable("v") int v) {
+        jdbc.update("INSERT INTO kv (v) VALUES (?)", v);
+        Integer last = jdbc.queryForObject(
+                "SELECT v FROM kv ORDER BY id DESC LIMIT 1", Integer.class);
+        Map<String, Object> out = new HashMap<>();
+        out.put("inserted", v);
+        out.put("readback", last);
+        return out;
+    }
+}

tidb-stmt-cache/src/main/java/com/example/tidbstmtcache/SchemaInitializer.java

@@ -0,0 +1,34 @@
+package com.example.tidbstmtcache;
+
+import org.springframework.boot.CommandLineRunner;
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.stereotype.Component;
+
+/**
+ * Creates the {@code kv} table on app boot. Runs as a CommandLineRunner
+ * (not Flyway / Liquibase) deliberately:
+ *
+ *  - Schema lives in one place that a reviewer can read in 30 seconds.
+ *  - The DDL goes through Spring's JdbcTemplate -> Connector/J the same
+ *    way as the rest of the workload, so it benefits from keploy's
+ *    synthetic-OK fallback for unmocked DDL (matchCommand's
+ *    BEGIN/CREATE/DROP/... allowlist in match.go). No mock needs to
+ *    exist for replay to satisfy the CREATE TABLE response.
+ */
+@Component
+public class SchemaInitializer implements CommandLineRunner {
+
+    private final JdbcTemplate jdbc;
+
+    public SchemaInitializer(JdbcTemplate jdbc) {
+        this.jdbc = jdbc;
+    }
+
+    @Override
+    public void run(String... args) {
+        jdbc.execute("CREATE TABLE IF NOT EXISTS kv (" +
+                "id INT PRIMARY KEY AUTO_INCREMENT, " +
+                "v INT NOT NULL" +
+                ")");
+    }
+}

tidb-stmt-cache/src/main/java/com/example/tidbstmtcache/TidbStmtCacheApplication.java

@@ -0,0 +1,11 @@
+package com.example.tidbstmtcache;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication
+public class TidbStmtCacheApplication {
+    public static void main(String[] args) {
+        SpringApplication.run(TidbStmtCacheApplication.class, args);
+    }
+}

tidb-stmt-cache/src/main/resources/application.properties

@@ -0,0 +1,24 @@
+server.port=8080
+
+# --- TiDB DataSource ---
+# TiDB :4000 speaks the MySQL wire protocol, so the MySQL Connector/J 8.x
+# driver works as-is. The three JDBC parameters that matter for the
+# orphan-EXECUTE scenario keploy/keploy@b2e68adb addresses:
+#
+#   useServerPrepStmts=true   server-side prepared statements (stmtIDs
+#                             come from the database, not the client).
+#   cachePrepStmts=true       per-Connection PreparedStatement cache
+#                             keyed by SQL text. This is what causes
+#                             Connector/J to skip COM_STMT_PREPARE on
+#                             cache hits and emit only COM_STMT_EXECUTE.
+#   prepStmtCacheSize=250     ensure the cache actually retains the few
+#                             statements this sample exercises.
+#
+# root / no password matches the default TiDB bootstrap user in the
+# pingcap/tidb:v8.5.x image started by docker-compose.yml -- TiDB does
+# not require an init.sql step the way MySQL does because the SchemaInitializer
+# CommandLineRunner creates the table from inside the app.
+datasource.tidb.jdbc-url=jdbc:mysql://localhost:4000/test?useSSL=false&allowPublicKeyRetrieval=true&useServerPrepStmts=true&cachePrepStmts=true&prepStmtCacheSize=250
+datasource.tidb.username=root
+datasource.tidb.password=
+datasource.tidb.driver-class-name=com.mysql.cj.jdbc.Driver