mz_testdrive/action/sql_server/
execute.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::time::Duration;
11
12use anyhow::{Context, anyhow};
13use mz_ore::str::StrExt;
14
15use crate::action::{ControlFlow, State};
16use crate::parser::BuiltinCommand;
17
18/// Check if an error is a transient SQL Server error that should be retried.
19///
20/// Covers:
21/// - Deadlock victim (error 1205)
22/// - SQL Server Agent still starting (error 14258 inside 22836/22832) — the
23///   Agent is needed for CDC job creation and may not be ready even though the
24///   healthcheck (`SELECT 1`) already passes.
25/// - Database not yet available during startup (error 904)
26fn is_retryable_error(err: &anyhow::Error) -> bool {
27    // Use alternate Display format `{:#}` to get the full anyhow error chain,
28    // not just the outermost context message.
29    let msg = format!("{:#}", err);
30    (msg.contains("1205") && msg.contains("deadlock"))
31        || msg.contains("SQLServerAgent is starting")
32        || msg.contains("cannot be autostarted during server shutdown or startup")
33}
34
35/// Maximum number of retries for transient errors.
36const MAX_RETRIES: usize = 20;
37
38/// Fixed backoff duration between retries.
39const RETRY_BACKOFF: Duration = Duration::from_millis(100);
40
41/// Executes `query`, retrying transient errors.
42///
43/// NOTE: a retry re-executes `query` in its entirety. If `query` contains
44/// multiple statements (e.g. with `split-lines=false`), statements that
45/// committed before the transient error are executed again, duplicating their
46/// side effects. Callers must pass a single statement, or statements that are
47/// safe to re-execute.
48async fn execute_with_retry(
49    client: &mut mz_sql_server_util::Client,
50    query: &str,
51) -> Result<(), anyhow::Error> {
52    for attempt in 0..=MAX_RETRIES {
53        match client
54            .simple_query(query.to_string())
55            .await
56            .context("executing SQL Server query")
57        {
58            Ok(_) => return Ok(()),
59            Err(err) if is_retryable_error(&err) && attempt < MAX_RETRIES => {
60                println!(
61                    ">> transient error (attempt {}/{}), retrying after {:?}: {:#}",
62                    attempt + 1,
63                    MAX_RETRIES,
64                    RETRY_BACKOFF,
65                    err,
66                );
67                tokio::time::sleep(RETRY_BACKOFF).await;
68            }
69            Err(err) => return Err(err),
70        }
71    }
72    unreachable!()
73}
74
75pub async fn run_execute(
76    mut cmd: BuiltinCommand,
77    state: &mut State,
78) -> Result<ControlFlow, anyhow::Error> {
79    let name = cmd.args.string("name")?;
80    let split_lines = cmd.args.opt_bool("split-lines")?.unwrap_or(true);
81    // When set, wraps the SQL in a Transaction and then drops it without
82    // calling commit or rollback. Used to test that Transaction::drop sends
83    // ROLLBACK correctly.
84    let abandon_txn = cmd.args.opt_bool("abandon-txn")?.unwrap_or(false);
85    cmd.args.done()?;
86
87    let client = state
88        .sql_server_clients
89        .get_mut(&name)
90        .ok_or_else(|| anyhow!("connection {} not found", name.quoted()))?;
91
92    if abandon_txn {
93        let mut txn = client
94            .transaction()
95            .await
96            .context("begin transaction for abandon-txn")?;
97        if split_lines {
98            for query in &cmd.input {
99                println!(">> (abandon-txn) {}", query);
100                txn.simple_query(query.to_string())
101                    .await
102                    .context("executing SQL Server query in transaction")?;
103            }
104        } else {
105            let query = cmd.input.join("\n");
106            println!(">> (abandon-txn) {}", query);
107            txn.simple_query(query)
108                .await
109                .context("executing SQL Server query in transaction")?;
110        }
111        // Transaction dropped here without commit or rollback.
112        // If Drop is correct, a ROLLBACK is sent via the channel.
113    } else {
114        if split_lines {
115            for query in &cmd.input {
116                println!(">> {}", query);
117                execute_with_retry(client, query).await?;
118            }
119        } else {
120            let query = cmd.input.join("\n");
121            println!(">> {}", query);
122            // execute uses prepared statements, which will fail for CREATE FUNCTION/PROCEDURE etc, see
123            // https://github.com/prisma/tiberius/issues/236, so using simple_query instead
124            execute_with_retry(client, &query).await?;
125        }
126    }
127
128    Ok(ControlFlow::Continue)
129}
mz_testdrive/action/sql_server/execute.rs

mz_testdrive/action/sql_server/
execute.rs
