mz_testdrive/action/postgres/
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 anyhow::{Context, anyhow, bail};
11use mz_ore::task;
12use tokio_postgres::Client;
13
14use crate::action::{BackgroundTask, ControlFlow, State};
15use crate::parser::BuiltinCommand;
16use crate::util::postgres::postgres_client;
17
18async fn execute_input(cmd: BuiltinCommand, client: &Client) -> Result<(), anyhow::Error> {
19    for query in cmd.input {
20        println!(">> {}", query);
21        // `query` is raw SQL from testdrive input and may contain multiple
22        // statements; this command intentionally forwards it verbatim.
23        #[allow(clippy::disallowed_methods)]
24        client
25            .batch_execute(&query)
26            .await
27            .context("executing postgres query")?;
28    }
29    Ok(())
30}
31
32pub async fn run_execute(
33    mut cmd: BuiltinCommand,
34    state: &mut State,
35) -> Result<ControlFlow, anyhow::Error> {
36    let connection = cmd.args.string("connection")?;
37    let background = cmd.args.opt_bool("background")?.unwrap_or(false);
38    cmd.args.done()?;
39
40    match (connection.starts_with("postgres://"), background) {
41        (true, true) => {
42            let (client_inner, _) = postgres_client(&connection, state.default_timeout).await?;
43            let desc = cmd.input.first().cloned().unwrap_or_default();
44            // Capture a cancel token before moving the client into the task, so
45            // the query can be stopped on the server if the task overruns its
46            // deadline and must be aborted.
47            let cancel_token = client_inner.cancel_token();
48            let handle = task::spawn(|| "postgres-execute", async move {
49                execute_input(cmd, &client_inner).await
50            });
51            // The task is joined at the end of the file so that failures fail
52            // the test, as documented.
53            state.background_tasks.push(BackgroundTask {
54                desc,
55                handle,
56                cancel_token,
57                url: connection,
58            });
59        }
60        (false, true) => bail!("cannot use 'background' arg with referenced connection"),
61        (true, false) => {
62            let (client_inner, _) = postgres_client(&connection, state.default_timeout).await?;
63            execute_input(cmd, &client_inner).await?;
64        }
65        (false, false) => {
66            let client = state
67                .postgres_clients
68                .get(&connection)
69                .ok_or_else(|| anyhow!("connection '{}' not found", &connection))?;
70            execute_input(cmd, client).await?;
71        }
72    }
73
74    Ok(ControlFlow::Continue)
75}
mz_testdrive/action/postgres/execute.rs

mz_testdrive/action/postgres/
execute.rs
