mz_testdrive/action/sql_server/
connect.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;
13use mz_ore::retry::{Retry, RetryResult};
14use mz_sql_server_util::{Client, Config};
15
16use crate::action::{ControlFlow, State};
17use crate::parser::BuiltinCommand;
18
19/// Whether a connection error is transient and worth retrying.
20///
21/// Azure SQL Database returns error 40613 ("Database ... is not currently
22/// available") while a serverless database is resuming or being scaled, which
23/// can take tens of seconds. The connection succeeds once the database is
24/// online, so we retry rather than fail the test.
25fn is_retryable_connect_error(err: &anyhow::Error) -> bool {
26    // Use alternate Display format `{:#}` to get the full anyhow error chain,
27    // not just the outermost context message.
28    let msg = format!("{:#}", err);
29    msg.contains("40613") || msg.contains("is not currently available")
30}
31
32pub async fn run_connect(
33    mut cmd: BuiltinCommand,
34    state: &mut State,
35) -> Result<ControlFlow, anyhow::Error> {
36    let name = cmd.args.string("name")?;
37    // Retry policy for transient connection errors, all in seconds. Defaults
38    // are generous because an auto-paused Azure SQL Database can take tens of
39    // seconds to resume.
40    let initial_backoff = cmd
41        .args
42        .opt_parse::<f64>("retry-initial-backoff")?
43        .map(Duration::from_secs_f64)
44        .unwrap_or(Duration::from_millis(500));
45    let clamp_backoff = cmd
46        .args
47        .opt_parse::<f64>("retry-clamp-backoff")?
48        .map(Duration::from_secs_f64)
49        .unwrap_or(Duration::from_secs(5));
50    let max_duration = cmd
51        .args
52        .opt_parse::<f64>("retry-max-duration")?
53        .map(Duration::from_secs_f64)
54        .unwrap_or(Duration::from_secs(90));
55    cmd.args.done()?;
56
57    let ado_string = cmd.input.join("\n");
58
59    let config = Config::from_ado_string(&ado_string).context("parsing ADO string")?;
60
61    let client = Retry::default()
62        .initial_backoff(initial_backoff)
63        .clamp_backoff(clamp_backoff)
64        .max_duration(max_duration)
65        .retry_async(|_| async {
66            match Client::connect(config.clone())
67                .await
68                .context("connecting to SQL server")
69            {
70                Ok(client) => RetryResult::Ok(client),
71                Err(err) if is_retryable_connect_error(&err) => {
72                    println!(">> transient connect error, retrying: {:#}", err);
73                    RetryResult::RetryableErr(err)
74                }
75                Err(err) => RetryResult::FatalErr(err),
76            }
77        })
78        .await?;
79    state.sql_server_clients.insert(name.clone(), client);
80
81    Ok(ControlFlow::Continue)
82}
mz_testdrive/action/sql_server/connect.rs

mz_testdrive/action/sql_server/
connect.rs
