replicators: Add snapshot query comment for MySQL

This commit adds the ability to add a comment to snapshot queries
for MySQL. The comment is added to the snapshot queries by the
caller and is automatically enclosed within the delimiters /* and */.

This is useful to by-pass tools like Percona Toolkit pt-kill which
automatically kills queries that take longer than a threshold.

Fixes REA-5946.
Closes: 1549

Release-Note-Core: Added the ability to add a comment to snapshot
  queries for MySQL.
Change-Id: I9220e91a5c5a42ffe4b98ede5956d8d389478977
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/10242
Tested-by: Buildkite CI
Reviewed-by: Michael Zink <michael.z@readyset.io>

Author: altmannmarcelo

database-utils/src/lib.rs

@@ -165,6 +165,14 @@ pub struct UpstreamConfig {
     )]
     #[serde(default = "default_status_update_interval_secs")]
     pub status_update_interval_secs: u16,
+
+    /// A string to be included as a comment in snapshot queries for MySQL.
+    /// The provided value is automatically enclosed within the delimiters
+    /// /* and */. The caller must supply only the raw comment text, without
+    /// any comment markers.
+    #[arg(long, env = "SNAPSHOT_QUERY_COMMENT", default_value = "")]
+    #[serde(default)]
+    pub snapshot_query_comment: Option<String>,
 }
 
 impl UpstreamConfig {
@@ -321,6 +329,7 @@ impl Default for UpstreamConfig {
             ignore_ulimit_check: false,
             status_update_interval_secs: 10,
             max_parallel_snapshot_tables: default_max_parallel_snapshot_tables(),
+            snapshot_query_comment: Default::default(),
         }
     }
 }

readyset/pkg/common/readyset.conf

@@ -107,3 +107,9 @@ NO_COLOR=true
 ## ALLOWED_USERS should be different than the one you added to the
 ## UPSTREAM_DB_URL - format "user1:pass1,user2:pass2"
 # ALLOWED_USERS=""
+
+## A string to be included as a comment in snapshot queries for MySQL.
+## The provided value is automatically enclosed within the delimiters
+## /* and */. The caller must supply only the raw comment text, without
+## any comment markers.
+# SNAPSHOT_QUERY_COMMENT=

replicators/src/mysql_connector/snapshot.rs

@@ -63,6 +63,8 @@ pub(crate) struct MySqlReplicator<'a> {
     pub(crate) table_filter: &'a mut TableFilter,
     /// Base parsing config to use
     pub(crate) parsing_preset: ParsingPreset,
+    /// Whether to add a comment for snapshot queries
+    pub(crate) snapshot_query_comment: Option<String>,
 }
 
 /// Get the list of tables defined in the database
@@ -423,6 +425,7 @@ impl MySqlReplicator<'_> {
         mut trx: Transaction<'static>,
         mut table_mutator: readyset_client::Table,
         snapshot_report_interval_secs: u16,
+        snapshot_query_comment: Option<String>,
     ) -> ReadySetResult<()> {
         let mut cnt = 0;
         let mut snapshot_type = SnapshotType::new(&table_mutator)?;
@@ -432,7 +435,7 @@ impl MySqlReplicator<'_> {
         // don't have support to lookup a collation from its name. Temporally get the
         // collation ID from querying IS. Later we can avoid the extra query.
         let (count_query, initial_query, bound_base_query, collation_query) =
-            snapshot_type.get_queries(&table_mutator);
+            snapshot_type.get_queries(&table_mutator, snapshot_query_comment);
 
         let collations = trx
             .query(collation_query)
@@ -661,14 +664,20 @@ impl MySqlReplicator<'_> {
         span.in_scope(|| info!("Read lock released"));
 
         let table_mutator = noria.table(table.clone()).instrument(span.clone()).await?;
+        let snapshot_query_comment = self.snapshot_query_comment.clone();
 
         Ok(tokio::spawn(async move {
             (
                 table,
                 repl_offset,
-                Self::snapshot_table(trx, table_mutator, snapshot_report_interval_secs)
-                    .instrument(span)
-                    .await,
+                Self::snapshot_table(
+                    trx,
+                    table_mutator,
+                    snapshot_report_interval_secs,
+                    snapshot_query_comment,
+                )
+                .instrument(span)
+                .await,
             )
         }))
     }

replicators/src/mysql_connector/snapshot_type.rs

@@ -65,7 +65,11 @@ impl SnapshotType {
     /// Returns:
     /// * A tuple containing the count query, the initial query, the bound based query
     ///   and the collation query
-    pub fn get_queries(&self, table: &readyset_client::Table) -> (String, String, String, String) {
+    pub fn get_queries(
+        &self,
+        table: &readyset_client::Table,
+        snapshot_query_comment: Option<String>,
+    ) -> (String, String, String, String) {
         let force_index = match self {
             SnapshotType::KeyBased { name, .. } => {
                 if let Some(name) = name {
@@ -95,6 +99,13 @@ impl SnapshotType {
             table.table_name().name,
             schema
         );
+
+        let snapshot_query_comment = snapshot_query_comment
+            .map(|s| s.replace("/*", "").replace("*/", ""))
+            .filter(|s| !s.is_empty())
+            .map(|s| format!(" /* {s} */"))
+            .unwrap_or_default();
+
         let (initial_query, bound_based_query) = match self {
             SnapshotType::KeyBased { ref keys, .. } => {
                 let keys = keys
@@ -119,14 +130,14 @@ impl SnapshotType {
                 let next_bound = format!("({next_bound})");
 
                 let initial_query = format!(
-                    "SELECT * FROM {} {} ORDER BY {} LIMIT {}",
+                    "SELECT{snapshot_query_comment} * FROM {} {} ORDER BY {} LIMIT {}",
                     table.table_name().display(readyset_sql::Dialect::MySQL),
                     force_index,
                     order_by,
                     MYSQL_BATCH_SIZE
                 );
                 let bound_based_query = format!(
-                    "SELECT * FROM {} {} WHERE {} ORDER BY {} LIMIT {}",
+                    "SELECT{snapshot_query_comment} * FROM {} {} WHERE {} ORDER BY {} LIMIT {}",
                     table.table_name().display(readyset_sql::Dialect::MySQL),
                     force_index,
                     next_bound,
@@ -137,7 +148,7 @@ impl SnapshotType {
             }
             SnapshotType::FullTableScan => {
                 let initial_query = format!(
-                    "SELECT * FROM {}",
+                    "SELECT{snapshot_query_comment} * FROM {}",
                     table.table_name().display(readyset_sql::Dialect::MySQL)
                 );
                 (initial_query.clone(), initial_query)

replicators/src/noria_adapter.rs

@@ -348,6 +348,7 @@ impl<'a> NoriaAdapter<'a> {
                     pool,
                     table_filter,
                     parsing_preset,
+                    snapshot_query_comment: config.snapshot_query_comment.clone(),
                 };
 
                 let snapshot_start = Instant::now();