mz_testdrive/action/sql_server/
set_from_sql.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};
11use mz_ore::str::StrExt;
12
13use crate::action::{ControlFlow, State};
14use crate::parser::BuiltinCommand;
15
16pub async fn run_set_from_sql(
17    mut cmd: BuiltinCommand,
18    state: &mut State,
19) -> Result<ControlFlow, anyhow::Error> {
20    let name = cmd.args.string("name")?;
21    let var = cmd.args.string("var")?;
22    cmd.args.done()?;
23
24    let client = state
25        .sql_server_clients
26        .get_mut(&name)
27        .ok_or_else(|| anyhow!("connection {} not found", name.quoted()))?;
28
29    let query = cmd.input.join("\n");
30    println!(">> {}", query);
31    // execute uses prepared statements, which will fail for CREATE FUNCTION/PROCEDURE etc, see
32    // https://github.com/prisma/tiberius/issues/236, so using simple_query instead
33    let rows = client
34        .simple_query(query)
35        .await
36        .context("executing SQL Server query")?;
37
38    let row = rows
39        .into_iter()
40        .next()
41        .expect("sql-server-set-from-sql query must return exactly one row");
42
43    let value: &str = row.try_get(0)?.expect("deserializing value as string");
44
45    state.cmd_vars.insert(var, value.to_string());
46
47    Ok(ControlFlow::Continue)
48}