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
18/// URLs for built-in connection names that resolve without a prior
19/// `postgres-connect` registration. An explicit `postgres-connect` with the
20/// same name takes precedence, since `state.postgres_clients` is consulted
21/// first.
22fn builtin_connection_url(state: &State, name: &str) -> Option<String> {
23    match name {
24        "mz_system" => Some(format!(
25            "postgres://mz_system:materialize@{}",
26            state.materialize.internal_sql_addr
27        )),
28        "materialize" => Some(format!(
29            "postgres://materialize:materialize@{}",
30            state.materialize.sql_addr
31        )),
32        _ => None,
33    }
34}
35
36async fn execute_input(cmd: BuiltinCommand, client: &Client) -> Result<(), anyhow::Error> {
37    for query in cmd.input {
38        println!(">> {}", query);
39        // `query` is raw SQL from testdrive input and may contain multiple
40        // statements; this command intentionally forwards it verbatim.
41        #[allow(clippy::disallowed_methods)]
42        client
43            .batch_execute(&query)
44            .await
45            .context("executing postgres query")?;
46    }
47    Ok(())
48}
49
50pub async fn run_execute(
51    mut cmd: BuiltinCommand,
52    state: &mut State,
53) -> Result<ControlFlow, anyhow::Error> {
54    let connection = cmd.args.string("connection")?;
55    let background = cmd.args.opt_bool("background")?.unwrap_or(false);
56    cmd.args.done()?;
57
58    match (connection.starts_with("postgres://"), background) {
59        (true, true) => {
60            let (client_inner, _) = postgres_client(&connection, state.default_timeout).await?;
61            let desc = cmd.input.first().cloned().unwrap_or_default();
62            // Capture a cancel token before moving the client into the task, so
63            // the query can be stopped on the server if the task overruns its
64            // deadline and must be aborted.
65            let cancel_token = client_inner.cancel_token();
66            let handle = task::spawn(|| "postgres-execute", async move {
67                execute_input(cmd, &client_inner).await
68            });
69            // The task is joined at the end of the file so that failures fail
70            // the test, as documented.
71            state.background_tasks.push(BackgroundTask {
72                desc,
73                handle,
74                cancel_token,
75                url: connection,
76            });
77        }
78        (false, true) => bail!("cannot use 'background' arg with referenced connection"),
79        (true, false) => {
80            let (client_inner, _) = postgres_client(&connection, state.default_timeout).await?;
81            execute_input(cmd, &client_inner).await?;
82        }
83        (false, false) => {
84            if !state.postgres_clients.contains_key(&connection)
85                && let Some(url) = builtin_connection_url(state, &connection)
86            {
87                let (client, _) = postgres_client(&url, state.default_timeout).await?;
88                state.postgres_clients.insert(connection.clone(), client);
89            }
90            let client = state
91                .postgres_clients
92                .get(&connection)
93                .ok_or_else(|| anyhow!("connection '{}' not found", connection))?;
94            execute_input(cmd, client).await?;
95        }
96    }
97
98    Ok(ControlFlow::Continue)
99}
mz_testdrive/action/postgres/execute.rs

mz_testdrive/action/postgres/
execute.rs
