Skip to main content

mz_testdrive/
action.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::collections::{BTreeMap, BTreeSet};
11use std::future::Future;
12
13use std::path::PathBuf;
14use std::str::FromStr;
15use std::sync::LazyLock;
16use std::time::Duration;
17use std::{env, fs};
18
19use anyhow::{Context, anyhow, bail};
20use async_trait::async_trait;
21use aws_credential_types::provider::ProvideCredentials;
22use aws_types::SdkConfig;
23use futures::future::FutureExt;
24use itertools::Itertools;
25use mz_adapter::catalog::{Catalog, ConnCatalog, DebugAwsContext};
26use mz_adapter::session::Session;
27use mz_build_info::BuildInfo;
28use mz_catalog::config::ClusterReplicaSizeMap;
29use mz_catalog::durable::BootstrapArgs;
30use mz_ccsr::SubjectVersion;
31use mz_kafka_util::client::{MzClientContext, create_new_client_config_simple};
32use mz_ore::error::ErrorExt;
33use mz_ore::metrics::MetricsRegistry;
34use mz_ore::now::SYSTEM_TIME;
35use mz_ore::retry::Retry;
36use mz_ore::task;
37use mz_ore::url::SensitiveUrl;
38use mz_persist_client::cache::PersistClientCache;
39use mz_persist_client::cfg::PersistConfig;
40use mz_persist_client::rpc::PubSubClientConnection;
41use mz_persist_client::{PersistClient, PersistLocation};
42use mz_postgres_util::{
43    Sql, batch_execute as pg_batch_execute, query as pg_query, query_one as pg_query_one,
44    sql as pg_sql,
45};
46use mz_sql::catalog::EnvironmentId;
47use mz_tls_util::make_tls;
48use rdkafka::ClientConfig;
49use rdkafka::producer::Producer;
50use regex::{Captures, Regex};
51use semver::Version;
52use tokio_postgres::CancelToken;
53use tokio_postgres::error::{DbError, SqlState};
54use tracing::info;
55use url::Url;
56
57use crate::error::PosError;
58use crate::parser::{
59    Command, PosCommand, SqlExpectedError, SqlOutput, VersionConstraint, validate_ident,
60};
61use crate::util;
62use crate::util::postgres::postgres_client;
63
64pub mod consistency;
65
66mod duckdb;
67mod file;
68mod fivetran;
69mod glue;
70mod http;
71mod kafka;
72mod mysql;
73mod nop;
74mod persist;
75mod postgres;
76mod protobuf;
77mod psql;
78mod s3;
79mod schema_registry;
80mod set;
81mod skip_end;
82mod skip_if;
83mod sleep;
84mod sql;
85mod sql_server;
86mod version_check;
87mod webhook;
88
89pub(crate) async fn verify_kafka_topics_exhausted(state: &State) -> Result<(), anyhow::Error> {
90    kafka::verify_topics_exhausted(state).await
91}
92
93/// User-settable configuration parameters.
94#[derive(Debug, Clone)]
95pub struct Config {
96    // === Testdrive options. ===
97    /// Variables to make available to the testdrive script.
98    ///
99    /// The value of each entry will be made available to the script in a
100    /// variable named `arg.KEY`.
101    pub arg_vars: BTreeMap<String, String>,
102    /// A random number to distinguish each run of a testdrive script.
103    pub seed: Option<String>,
104    /// Whether to reset Materialize state before executing each script and
105    /// to clean up AWS state after each script.
106    pub reset: bool,
107    /// Force the use of the specified temporary directory to use.
108    ///
109    /// If unspecified, testdrive creates a temporary directory with a random
110    /// name.
111    pub temp_dir: Option<String>,
112    /// Source string to print out on errors.
113    pub source: Option<String>,
114    /// The default timeout for cancellable operations.
115    pub default_timeout: Duration,
116    /// The default number of tries for retriable operations.
117    pub default_max_tries: usize,
118    /// The initial backoff interval for retry operations.
119    ///
120    /// Set to 0 to retry immediately on failure.
121    pub initial_backoff: Duration,
122    /// Backoff factor to use for retry operations.
123    ///
124    /// Set to 1 to retry at a steady pace.
125    pub backoff_factor: f64,
126    /// Should we skip coordinator and catalog consistency checks.
127    pub consistency_checks: consistency::Level,
128    /// How long a single consistency check may take before the test file fails.
129    pub consistency_check_timeout: Duration,
130    /// Whether to run statement logging consistency checks (adds a few seconds at the end of every
131    /// test file).
132    pub check_statement_logging: bool,
133    /// Whether to automatically rewrite wrong results instead of failing.
134    pub rewrite_results: bool,
135
136    // === Materialize options. ===
137    /// The pgwire connection parameters for the Materialize instance that
138    /// testdrive will connect to.
139    pub materialize_pgconfig: tokio_postgres::Config,
140    /// The internal pgwire connection parameters for the Materialize instance that
141    /// testdrive will connect to.
142    pub materialize_internal_pgconfig: tokio_postgres::Config,
143    /// Whether to use HTTPS instead of plain HTTP for the HTTP(S) connections.
144    pub materialize_use_https: bool,
145    /// The port for the public endpoints of the materialize instance that
146    /// testdrive will connect to via HTTP.
147    pub materialize_http_port: u16,
148    /// The port for the internal endpoints of the materialize instance that
149    /// testdrive will connect to via HTTP.
150    pub materialize_internal_http_port: u16,
151    /// The port for the password endpoints of the materialize instance that
152    /// testdrive will connect to via SQL.
153    pub materialize_password_sql_port: u16,
154    /// The port for the SASL endpoints of the materialize instance that
155    /// testdrive will connect to via SQL.
156    pub materialize_sasl_sql_port: u16,
157    /// Session parameters to set after connecting to materialize.
158    pub materialize_params: Vec<(String, String)>,
159    /// An optional catalog configuration.
160    pub materialize_catalog_config: Option<CatalogConfig>,
161    /// Build information
162    pub build_info: &'static BuildInfo,
163    /// Configured cluster replica sizes
164    pub materialize_cluster_replica_sizes: ClusterReplicaSizeMap,
165
166    // === Persist options. ===
167    /// Handle to the persist consensus system.
168    pub persist_consensus_url: Option<SensitiveUrl>,
169    /// Handle to the persist blob storage.
170    pub persist_blob_url: Option<SensitiveUrl>,
171
172    // === Confluent options. ===
173    /// The address of the Kafka broker that testdrive will interact with.
174    pub kafka_addr: String,
175    /// Default number of partitions to use for topics
176    pub kafka_default_partitions: usize,
177    /// Arbitrary rdkafka options for testdrive to use when connecting to the
178    /// Kafka broker.
179    pub kafka_opts: Vec<(String, String)>,
180    /// The URL of the schema registry that testdrive will connect to.
181    pub schema_registry_url: Url,
182    /// An optional path to a TLS certificate that testdrive will present when
183    /// performing client authentication.
184    ///
185    /// The keystore must be in the PKCS#12 format.
186    pub cert_path: Option<String>,
187    /// An optional password for the TLS certificate.
188    pub cert_password: Option<String>,
189    /// An optional username for basic authentication with the Confluent Schema
190    /// Registry.
191    pub ccsr_username: Option<String>,
192    /// An optional password for basic authentication with the Confluent Schema
193    /// Registry.
194    pub ccsr_password: Option<String>,
195
196    // === AWS options. ===
197    /// The configuration to use when connecting to AWS.
198    pub aws_config: SdkConfig,
199    /// The ID of the AWS account that `aws_config` configures.
200    pub aws_account: String,
201
202    // === Fivetran options. ===
203    /// Address of the Fivetran Destination that is currently running.
204    pub fivetran_destination_url: String,
205    /// Directory that is accessible to the Fivetran Destination.
206    pub fivetran_destination_files_path: String,
207}
208
209pub struct MaterializeState {
210    catalog_config: Option<CatalogConfig>,
211
212    sql_addr: String,
213    use_https: bool,
214    http_addr: String,
215    internal_sql_addr: String,
216    internal_http_addr: String,
217    password_sql_addr: String,
218    sasl_sql_addr: String,
219    user: String,
220    pgclient: tokio_postgres::Client,
221    environment_id: EnvironmentId,
222    bootstrap_args: BootstrapArgs,
223    // The AWS environment context, queried once from the running environmentd at
224    // startup, the same way environment_id is. The builtin AWS connection views
225    // fold these values into their optimized expressions, so the catalog copy
226    // opened for the consistency check must resolve them the same way. Queried
227    // at startup rather than at check time because a check-time query runs after
228    // arbitrary test state and is less reliable, while startup runs against a
229    // clean session on the currently connected version.
230    aws_account_id: Option<String>,
231    aws_external_id_prefix: Option<String>,
232    aws_connection_role_arn: Option<String>,
233}
234
235pub struct State {
236    // The Config that this `State` was originally created from.
237    pub config: Config,
238
239    // === Testdrive state. ===
240    arg_vars: BTreeMap<String, String>,
241    cmd_vars: BTreeMap<String, String>,
242    seed: String,
243    temp_path: PathBuf,
244    _tempfile: Option<tempfile::TempDir>,
245    default_timeout: Duration,
246    timeout: Duration,
247    max_tries: usize,
248    initial_backoff: Duration,
249    backoff_factor: f64,
250    consistency_checks: consistency::Level,
251    consistency_check_timeout: Duration,
252    check_statement_logging: bool,
253    consistency_checks_adhoc_skip: bool,
254    regex: Option<Regex>,
255    regex_replacement: String,
256    error_line_count: usize,
257    error_string: String,
258
259    // === Materialize state. ===
260    materialize: MaterializeState,
261
262    // === Persist state. ===
263    persist_consensus_url: Option<SensitiveUrl>,
264    persist_blob_url: Option<SensitiveUrl>,
265    build_info: &'static BuildInfo,
266    persist_clients: PersistClientCache,
267
268    // === Confluent state. ===
269    schema_registry_url: Url,
270    ccsr_client: mz_ccsr::Client,
271    kafka_addr: String,
272    kafka_admin: rdkafka::admin::AdminClient<MzClientContext>,
273    kafka_admin_opts: rdkafka::admin::AdminOptions,
274    kafka_config: ClientConfig,
275    kafka_default_partitions: usize,
276    kafka_producer: rdkafka::producer::FutureProducer<MzClientContext>,
277    kafka_topics: BTreeMap<String, usize>,
278    /// Topics whose final `kafka-verify-data` must consume the complete topic.
279    kafka_verify_topics: BTreeSet<String>,
280
281    // === AWS state. ===
282    aws_account: String,
283    aws_config: SdkConfig,
284
285    // === Database driver state. ===
286    pub duckdb_clients: BTreeMap<String, std::sync::Arc<std::sync::Mutex<::duckdb::Connection>>>,
287    mysql_clients: BTreeMap<String, mysql_async::Conn>,
288    postgres_clients: BTreeMap<String, tokio_postgres::Client>,
289    sql_server_clients: BTreeMap<String, mz_sql_server_util::Client>,
290    /// Tasks spawned by `postgres-execute background=true`. Joined at the end
291    /// of the file so their failures fail the test.
292    background_tasks: Vec<BackgroundTask>,
293
294    // === Fivetran state. ===
295    fivetran_destination_url: String,
296    fivetran_destination_files_path: String,
297
298    // === Rewrite state. ===
299    rewrite_results: bool,
300    /// Current file, results are replaced inline
301    pub rewrites: Vec<Rewrite>,
302    /// Start position of currently expected result
303    pub rewrite_pos_start: usize,
304    /// End position of currently expected result
305    pub rewrite_pos_end: usize,
306}
307
308pub struct Rewrite {
309    pub content: String,
310    pub start: usize,
311    pub end: usize,
312}
313
314/// A query spawned by `postgres-execute background=true`, retained so it can be
315/// joined at the end of the file.
316///
317/// Aborting `handle` stops the Rust task from awaiting the query's response,
318/// but it does not stop the query on the server. The tokio-postgres connection
319/// driver waits for every pending response before it closes the connection, so
320/// dropping the client leaves the SQL running until Materialize replies. To
321/// actually stop a query that overran its deadline we send an out-of-band
322/// cancel request through `cancel_token`, the same mechanism psql uses for
323/// Ctrl-C, before aborting the task.
324pub(crate) struct BackgroundTask {
325    desc: String,
326    handle: task::JoinHandle<Result<(), anyhow::Error>>,
327    cancel_token: CancelToken,
328    /// Connection URL, used to rebuild the TLS connector that the cancel
329    /// request opens its own connection with.
330    url: String,
331}
332
333/// Best-effort cancellation of the in-progress query on the connection
334/// `cancel_token` was derived from. Opens a fresh connection to send the cancel
335/// request. Cancellation is inherently racy and the caller aborts the task
336/// regardless, so failures are logged rather than propagated.
337async fn cancel_background_query(cancel_token: &CancelToken, url: &str, timeout: Duration) {
338    let tls = match tokio_postgres::Config::from_str(url)
339        .map_err(anyhow::Error::from)
340        .and_then(|config| make_tls(&config).map_err(anyhow::Error::from))
341    {
342        Ok(tls) => tls,
343        Err(e) => {
344            tracing::warn!("could not build TLS connector to cancel background query: {e}");
345            return;
346        }
347    };
348    match tokio::time::timeout(timeout, cancel_token.cancel_query(tls)).await {
349        Ok(Ok(())) => {}
350        Ok(Err(e)) => tracing::warn!("cancel request for background query failed: {e}"),
351        Err(_) => tracing::warn!("cancel request for background query timed out"),
352    }
353}
354
355impl State {
356    pub async fn initialize_cmd_vars(&mut self) -> Result<(), anyhow::Error> {
357        self.cmd_vars
358            .insert("testdrive.kafka-addr".into(), self.kafka_addr.clone());
359        self.cmd_vars.insert(
360            "testdrive.schema-registry-url".into(),
361            self.schema_registry_url.to_string(),
362        );
363        self.cmd_vars
364            .insert("testdrive.seed".into(), self.seed.clone());
365        self.cmd_vars.insert(
366            "testdrive.temp-dir".into(),
367            self.temp_path.display().to_string(),
368        );
369        self.cmd_vars
370            .insert("testdrive.aws-region".into(), self.aws_region().into());
371        self.cmd_vars
372            .insert("testdrive.aws-endpoint".into(), self.aws_endpoint().into());
373        self.cmd_vars
374            .insert("testdrive.aws-account".into(), self.aws_account.clone());
375        {
376            let aws_credentials = self
377                .aws_config
378                .credentials_provider()
379                .ok_or_else(|| anyhow!("no AWS credentials provider configured"))?
380                .provide_credentials()
381                .await
382                .context("fetching AWS credentials")?;
383            self.cmd_vars.insert(
384                "testdrive.aws-access-key-id".into(),
385                aws_credentials.access_key_id().to_owned(),
386            );
387            self.cmd_vars.insert(
388                "testdrive.aws-secret-access-key".into(),
389                aws_credentials.secret_access_key().to_owned(),
390            );
391            self.cmd_vars.insert(
392                "testdrive.aws-token".into(),
393                aws_credentials
394                    .session_token()
395                    .map(|token| token.to_owned())
396                    .unwrap_or_else(String::new),
397            );
398        }
399        self.cmd_vars.insert(
400            "testdrive.materialize-environment-id".into(),
401            self.materialize.environment_id.to_string(),
402        );
403        self.cmd_vars.insert(
404            "testdrive.materialize-sql-addr".into(),
405            self.materialize.sql_addr.clone(),
406        );
407        self.cmd_vars.insert(
408            "testdrive.materialize-internal-sql-addr".into(),
409            self.materialize.internal_sql_addr.clone(),
410        );
411        self.cmd_vars.insert(
412            "testdrive.materialize-password-sql-addr".into(),
413            self.materialize.password_sql_addr.clone(),
414        );
415        self.cmd_vars.insert(
416            "testdrive.materialize-sasl-sql-addr".into(),
417            self.materialize.sasl_sql_addr.clone(),
418        );
419        self.cmd_vars.insert(
420            "testdrive.materialize-user".into(),
421            self.materialize.user.clone(),
422        );
423        self.cmd_vars.insert(
424            "testdrive.fivetran-destination-url".into(),
425            self.fivetran_destination_url.clone(),
426        );
427        self.cmd_vars.insert(
428            "testdrive.fivetran-destination-files-path".into(),
429            self.fivetran_destination_files_path.clone(),
430        );
431
432        for (key, value) in env::vars() {
433            self.cmd_vars.insert(format!("env.{}", key), value);
434        }
435
436        for (key, value) in &self.arg_vars {
437            validate_ident(key)?;
438            self.cmd_vars
439                .insert(format!("arg.{}", key), value.to_string());
440        }
441
442        Ok(())
443    }
444    /// Makes of copy of the durable catalog and runs a function on its
445    /// state. Returns `None` if there's no catalog information in the State.
446    pub async fn with_catalog_copy<F, T>(
447        &self,
448        system_parameter_defaults: BTreeMap<String, String>,
449        build_info: &'static BuildInfo,
450        bootstrap_args: &BootstrapArgs,
451        enable_expression_cache_override: Option<bool>,
452        f: F,
453    ) -> Result<Option<T>, anyhow::Error>
454    where
455        F: FnOnce(ConnCatalog) -> T,
456    {
457        async fn persist_client(
458            persist_consensus_url: SensitiveUrl,
459            persist_blob_url: SensitiveUrl,
460            persist_clients: &PersistClientCache,
461        ) -> Result<PersistClient, anyhow::Error> {
462            let persist_location = PersistLocation {
463                blob_uri: persist_blob_url,
464                consensus_uri: persist_consensus_url,
465            };
466            Ok(persist_clients.open(persist_location).await?)
467        }
468
469        if let Some(CatalogConfig {
470            persist_consensus_url,
471            persist_blob_url,
472        }) = &self.materialize.catalog_config
473        {
474            let persist_client = persist_client(
475                persist_consensus_url.clone(),
476                persist_blob_url.clone(),
477                &self.persist_clients,
478            )
479            .await?;
480            let aws_context = DebugAwsContext {
481                aws_account_id: self.materialize.aws_account_id.clone(),
482                aws_external_id_prefix: self.materialize.aws_external_id_prefix.clone(),
483                aws_connection_role_arn: self.materialize.aws_connection_role_arn.clone(),
484            };
485            let catalog = Catalog::open_debug_read_only_persist_catalog_config(
486                persist_client,
487                SYSTEM_TIME.clone(),
488                self.materialize.environment_id.clone(),
489                system_parameter_defaults,
490                build_info,
491                bootstrap_args,
492                enable_expression_cache_override,
493                Some(aws_context),
494            )
495            .await?;
496            let res = f(catalog.for_session(&Session::dummy()));
497            catalog.expire().await;
498            Ok(Some(res))
499        } else {
500            Ok(None)
501        }
502    }
503
504    pub fn aws_endpoint(&self) -> &str {
505        self.aws_config.endpoint_url().unwrap_or("")
506    }
507
508    pub fn aws_region(&self) -> &str {
509        self.aws_config.region().map(|r| r.as_ref()).unwrap_or("")
510    }
511
512    /// Resets the adhoc skip consistency check that users can toggle per-file, and returns whether
513    /// the consistency checks should be skipped for this current run.
514    pub fn clear_skip_consistency_checks(&mut self) -> bool {
515        std::mem::replace(&mut self.consistency_checks_adhoc_skip, false)
516    }
517
518    /// Joins all tasks spawned by `postgres-execute background=true`,
519    /// returning one error per task that failed or did not complete within the
520    /// default timeout. Must be called before the end of the file so that
521    /// background failures fail the test.
522    pub(crate) async fn join_background_tasks(&mut self) -> Vec<anyhow::Error> {
523        let mut errors = Vec::new();
524        for BackgroundTask {
525            desc,
526            mut handle,
527            cancel_token,
528            url,
529        } in self.background_tasks.drain(..)
530        {
531            // Poll the handle by reference so it survives a timeout. Dropping a
532            // `JoinHandle` only detaches the task, it does not stop it, and a
533            // detached background query would keep running SQL against the same
534            // Materialize instance while later files execute.
535            match tokio::time::timeout(self.default_timeout, &mut handle).await {
536                Ok(Ok(())) => {}
537                Ok(Err(e)) => errors.push(e.context(format!("background query failed: {desc}"))),
538                Err(_) => {
539                    // Aborting `handle` alone leaves the query running on the
540                    // server. Cancel it first so it cannot overlap consistency
541                    // checks or later files, then abort and reap the task.
542                    cancel_background_query(&cancel_token, &url, self.default_timeout).await;
543                    handle.abort_and_wait().await;
544                    errors.push(anyhow!(
545                        "background query did not complete before the end of the file: {desc}"
546                    ));
547                }
548            }
549        }
550        errors
551    }
552
553    pub async fn reset_materialize(&self) -> Result<(), anyhow::Error> {
554        let (inner_client, _) = postgres_client(
555            &format!(
556                "postgres://mz_system:materialize@{}",
557                self.materialize.internal_sql_addr
558            ),
559            self.default_timeout,
560        )
561        .await?;
562
563        let version = pg_query_one(&inner_client, pg_sql!("SELECT mz_version_num()"), &[])
564            .await
565            .context("getting version of materialize")
566            .map(|row| row.get::<_, i32>(0))?;
567
568        let semver = pg_query_one(
569            &inner_client,
570            pg_sql!("SELECT right(split_part(mz_version(), ' ', 1), -1)"),
571            &[],
572        )
573        .await
574        .context("getting semver of materialize")
575        .map(|row| row.get::<_, String>(0))?
576        .parse::<semver::Version>()
577        .context("parsing semver of materialize")?;
578
579        pg_batch_execute(&inner_client, pg_sql!("ALTER SYSTEM RESET ALL"))
580            .await
581            .context("resetting materialize state: ALTER SYSTEM RESET ALL")?;
582
583        // Dangerous functions are useful for tests so we enable it for all tests.
584        {
585            let rename_version = Version::parse("0.128.0-dev.1").expect("known to be valid");
586            let enable_unsafe_functions = if semver >= rename_version {
587                "unsafe_enable_unsafe_functions"
588            } else {
589                "enable_unsafe_functions"
590            };
591            let res = pg_batch_execute(
592                &inner_client,
593                pg_sql!(
594                    "ALTER SYSTEM SET {} = on",
595                    Sql::ident(enable_unsafe_functions)
596                ),
597            )
598            .await
599            .context("enabling dangerous functions");
600            if let Err(e) = res {
601                match e.root_cause().downcast_ref::<DbError>() {
602                    Some(e) if *e.code() == SqlState::CANT_CHANGE_RUNTIME_PARAM => {
603                        info!(
604                            "can't enable unsafe functions because the server is safe mode; \
605                             testdrive scripts will fail if they use unsafe functions",
606                        );
607                    }
608                    _ => return Err(e),
609                }
610            }
611        }
612
613        for row in pg_query(&inner_client, pg_sql!("SHOW DATABASES"), &[])
614            .await
615            .context("resetting materialize state: SHOW DATABASES")?
616        {
617            let db_name: String = row.get(0);
618            if db_name.starts_with("testdrive_no_reset_") {
619                continue;
620            }
621            let drop_database = pg_sql!("DROP DATABASE {}", Sql::ident(&db_name));
622            sql::print_query(drop_database.as_str(), None);
623            pg_batch_execute(&inner_client, drop_database)
624                .await
625                .context(format!(
626                    "resetting materialize state: DROP DATABASE {}",
627                    db_name,
628                ))?;
629        }
630
631        // Get all user clusters not running any objects owned by users
632        let inactive_user_clusters = "
633        WITH
634            active_user_clusters AS
635            (
636                SELECT DISTINCT cluster_id, object_id
637                FROM
638                    (
639                        SELECT cluster_id, id FROM mz_catalog.mz_sources
640                        UNION ALL SELECT cluster_id, id FROM mz_catalog.mz_sinks
641                        UNION ALL
642                            SELECT cluster_id, id
643                            FROM mz_catalog.mz_materialized_views
644                        UNION ALL
645                            SELECT cluster_id, id FROM mz_catalog.mz_indexes
646                        UNION ALL
647                            SELECT cluster_id, id
648                            FROM mz_internal.mz_subscriptions
649                    )
650                    AS t (cluster_id, object_id)
651                WHERE cluster_id IS NOT NULL AND object_id LIKE 'u%'
652            )
653        SELECT name
654        FROM mz_catalog.mz_clusters
655        WHERE
656            id NOT IN ( SELECT cluster_id FROM active_user_clusters ) AND id LIKE 'u%'
657                AND
658            owner_id LIKE 'u%';";
659
660        let inactive_clusters = pg_query(&inner_client, Sql::new(inactive_user_clusters), &[])
661            .await
662            .context("resetting materialize state: inactive_user_clusters")?;
663
664        if !inactive_clusters.is_empty() {
665            println!("cleaning up user clusters from previous tests...")
666        }
667
668        for cluster_name in inactive_clusters {
669            let cluster_name: String = cluster_name.get(0);
670            if cluster_name.starts_with("testdrive_no_reset_") {
671                continue;
672            }
673            let drop_cluster = pg_sql!("DROP CLUSTER {}", Sql::ident(&cluster_name));
674            sql::print_query(drop_cluster.as_str(), None);
675            pg_batch_execute(&inner_client, drop_cluster)
676                .await
677                .context(format!(
678                    "resetting materialize state: DROP CLUSTER {}",
679                    cluster_name,
680                ))?;
681        }
682
683        pg_batch_execute(&inner_client, pg_sql!("CREATE DATABASE materialize"))
684            .await
685            .context("resetting materialize state: CREATE DATABASE materialize")?;
686
687        // Attempt to remove all users but the current user. Old versions of
688        // Materialize did not support roles, so this degrades gracefully if
689        // mz_roles does not exist.
690        if let Ok(rows) = pg_query(&inner_client, pg_sql!("SELECT name FROM mz_roles"), &[]).await {
691            for row in rows {
692                let role_name: String = row.get(0);
693                if role_name == self.materialize.user || role_name.starts_with("mz_") {
694                    continue;
695                }
696                let drop_role = pg_sql!("DROP ROLE {}", Sql::ident(&role_name));
697                sql::print_query(drop_role.as_str(), None);
698                pg_batch_execute(&inner_client, drop_role)
699                    .await
700                    .context(format!(
701                        "resetting materialize state: DROP ROLE {}",
702                        role_name,
703                    ))?;
704            }
705        }
706
707        // Alter materialize user with all system privileges.
708        pg_batch_execute(
709            &inner_client,
710            pg_sql!(
711                "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
712                Sql::ident(&self.materialize.user)
713            ),
714        )
715        .await?;
716
717        // Grant initial privileges.
718        pg_batch_execute(
719            &inner_client,
720            pg_sql!("GRANT USAGE ON DATABASE materialize TO PUBLIC"),
721        )
722        .await?;
723        pg_batch_execute(
724            &inner_client,
725            pg_sql!(
726                "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
727                Sql::ident(&self.materialize.user)
728            ),
729        )
730        .await?;
731        pg_batch_execute(
732            &inner_client,
733            pg_sql!(
734                "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
735                Sql::ident(&self.materialize.user)
736            ),
737        )
738        .await?;
739
740        let cluster = match version {
741            ..=8199 => "default",
742            8200.. => "quickstart",
743        };
744        pg_batch_execute(
745            &inner_client,
746            pg_sql!("GRANT USAGE ON CLUSTER {} TO PUBLIC", Sql::ident(cluster)),
747        )
748        .await?;
749        pg_batch_execute(
750            &inner_client,
751            pg_sql!(
752                "GRANT ALL PRIVILEGES ON CLUSTER {} TO {}",
753                Sql::ident(cluster),
754                Sql::ident(&self.materialize.user)
755            ),
756        )
757        .await?;
758
759        Ok(())
760    }
761
762    /// Delete Kafka topics + CCSR subjects that were created in this run
763    pub async fn reset_kafka(&self) -> Result<(), anyhow::Error> {
764        use rdkafka::types::RDKafkaErrorCode;
765        let mut errors: Vec<anyhow::Error> = Vec::new();
766
767        let metadata = self.kafka_producer.client().fetch_metadata(
768            None,
769            Some(std::cmp::max(Duration::from_secs(1), self.default_timeout)),
770        )?;
771
772        let testdrive_topics: Vec<_> = metadata
773            .topics()
774            .iter()
775            .filter_map(|t| {
776                if t.name().starts_with("testdrive-") {
777                    Some(t.name())
778                } else {
779                    None
780                }
781            })
782            .collect();
783
784        if !testdrive_topics.is_empty() {
785            match self
786                .kafka_admin
787                .delete_topics(&testdrive_topics, &self.kafka_admin_opts)
788                .await
789            {
790                Ok(res) => {
791                    if res.len() != testdrive_topics.len() {
792                        errors.push(anyhow!(
793                            "kafka topic deletion returned {} results, but exactly {} expected",
794                            res.len(),
795                            testdrive_topics.len()
796                        ));
797                    }
798                    for (res, topic) in res.iter().zip_eq(testdrive_topics.iter()) {
799                        match res {
800                            Ok(_) | Err((_, RDKafkaErrorCode::UnknownTopicOrPartition)) => (),
801                            Err((_, err)) => {
802                                errors.push(anyhow!("unable to delete {}: {}", topic, err));
803                            }
804                        }
805                    }
806                }
807                Err(e) => {
808                    errors.push(e.into());
809                }
810            };
811        }
812
813        let schema_registry_errors = self.reset_schema_registry().await;
814
815        errors.extend(schema_registry_errors);
816        if errors.is_empty() {
817            Ok(())
818        } else {
819            bail!(
820                "deleting Kafka topics: {} errors: {}",
821                errors.len(),
822                errors
823                    .into_iter()
824                    .map(|e| e.to_string_with_causes())
825                    .join("\n")
826            );
827        }
828    }
829
830    #[allow(clippy::disallowed_types)]
831    async fn reset_schema_registry(&self) -> Vec<anyhow::Error> {
832        use std::collections::HashMap;
833
834        let mut errors = Vec::new();
835        match self
836            .ccsr_client
837            .list_subjects()
838            .await
839            .context("listing schema registry subjects")
840        {
841            Ok(subjects) => {
842                let testdrive_subjects: Vec<_> = subjects
843                    .into_iter()
844                    .filter(|s| s.starts_with("testdrive-"))
845                    .collect();
846
847                // Build the dependency graphs: subject -> list of subjects it references
848                let mut graphs: HashMap<SubjectVersion, Vec<SubjectVersion>> = HashMap::new();
849
850                for subject in &testdrive_subjects {
851                    match self.ccsr_client.get_subject_with_references(subject).await {
852                        Ok((subj, refs)) => {
853                            // Filter to only testdrive subjects
854                            let refs: Vec<_> = refs
855                                .into_iter()
856                                .filter(|r| r.subject.starts_with("testdrive-"))
857                                .collect();
858                            graphs.insert(
859                                SubjectVersion {
860                                    subject: subj.name,
861                                    version: subj.version,
862                                },
863                                refs,
864                            );
865                        }
866                        Err(mz_ccsr::GetBySubjectError::SubjectNotFound) => {
867                            // Subject was already deleted, skip it
868                        }
869                        Err(e) => {
870                            errors.push(anyhow::anyhow!(
871                                "failed to get references for subject {}: {}",
872                                subject,
873                                e
874                            ));
875                        }
876                    }
877                }
878
879                // Get topological ordering (0 = root/no dependencies, higher = more dependents)
880                // We need to delete in reverse order (highest first = subjects that depend on others)
881                let subjects_to_delete: Vec<_> = match mz_ccsr::topological_sort(&graphs) {
882                    Ok(ordered) => {
883                        let mut subjects: Vec<_> = ordered.into_iter().collect();
884                        subjects.sort_by(|a, b| a.1.cmp(&b.1));
885                        subjects.into_iter().map(|(s, _)| s.clone()).collect()
886                    }
887                    Err(_) => {
888                        tracing::info!("Cycle detected, attempting to delete anyway");
889                        // Cycle detected or other error - fall back to deleting in any order
890                        graphs.into_keys().collect()
891                    }
892                };
893
894                for subject in subjects_to_delete {
895                    match self.ccsr_client.delete_subject(&subject.subject).await {
896                        Ok(()) | Err(mz_ccsr::DeleteError::SubjectNotFound) => (),
897                        Err(e) => errors.push(e.into()),
898                    }
899                }
900            }
901            Err(e) => {
902                errors.push(e);
903            }
904        }
905        errors
906    }
907}
908
909/// Configuration for the Catalog.
910#[derive(Debug, Clone)]
911pub struct CatalogConfig {
912    /// Handle to the persist consensus system.
913    pub persist_consensus_url: SensitiveUrl,
914    /// Handle to the persist blob storage.
915    pub persist_blob_url: SensitiveUrl,
916}
917
918pub enum ControlFlow {
919    Continue,
920    SkipBegin,
921    SkipEnd,
922}
923
924#[async_trait]
925pub(crate) trait Run {
926    async fn run(self, state: &mut State) -> Result<ControlFlow, PosError>;
927}
928
929#[async_trait]
930impl Run for PosCommand {
931    async fn run(self, state: &mut State) -> Result<ControlFlow, PosError> {
932        macro_rules! handle_version {
933            ($version_constraint:expr) => {
934                match $version_constraint {
935                    Some(VersionConstraint { min, max }) => {
936                        match version_check::run_version_check(min, max, state).await {
937                            Ok(true) => return Ok(ControlFlow::Continue),
938                            Ok(false) => {}
939                            Err(err) => return Err(PosError::new(err, self.pos)),
940                        }
941                    }
942                    None => {}
943                }
944            };
945        }
946
947        let wrap_err = |e| PosError::new(e, self.pos);
948        // Substitute variables at startup except for the command-specific ones
949        // Those will be substituted at runtime
950        let ignore_prefix = match &self.command {
951            Command::Builtin(builtin, _) => Some(builtin.name.clone()),
952            _ => None,
953        };
954        let subst = |msg: &str, vars: &BTreeMap<String, String>| {
955            substitute_vars(msg, vars, &ignore_prefix, false).map_err(wrap_err)
956        };
957        let subst_re = |msg: &str, vars: &BTreeMap<String, String>| {
958            substitute_vars(msg, vars, &ignore_prefix, true).map_err(wrap_err)
959        };
960
961        let r = match self.command {
962            Command::Builtin(mut builtin, version_constraint) => {
963                handle_version!(version_constraint);
964                for val in builtin.args.values_mut() {
965                    *val = subst(val, &state.cmd_vars)?;
966                }
967                for line in &mut builtin.input {
968                    *line = subst(line, &state.cmd_vars)?;
969                }
970                match builtin.name.as_ref() {
971                    "check-consistency" => consistency::run_consistency_checks(state).await,
972                    "skip-consistency-checks" => {
973                        consistency::skip_consistency_checks(builtin, state)
974                    }
975                    "check-shard-tombstone" => {
976                        consistency::run_check_shard_tombstone(builtin, state).await
977                    }
978                    "duckdb-execute" => duckdb::run_execute(builtin, state).await,
979                    "duckdb-query" => duckdb::run_query(builtin, state).await,
980                    "fivetran-destination" => {
981                        fivetran::run_destination_command(builtin, state).await
982                    }
983                    "file-append" => file::run_append(builtin, state).await,
984                    "file-delete" => file::run_delete(builtin, state).await,
985                    "glue-create-schema" => glue::run_create_schema(builtin, state).await,
986                    "glue-verify-compatibility" => {
987                        glue::run_verify_compatibility(builtin, state).await
988                    }
989                    "http-request" => http::run_request(builtin, state).await,
990                    "kafka-add-partitions" => kafka::run_add_partitions(builtin, state).await,
991                    "kafka-create-topic" => kafka::run_create_topic(builtin, state).await,
992                    "kafka-wait-topic" => kafka::run_wait_topic(builtin, state).await,
993                    "kafka-delete-records" => kafka::run_delete_records(builtin, state).await,
994                    "kafka-delete-topic-flaky" => kafka::run_delete_topic(builtin, state).await,
995                    "kafka-ingest" => kafka::run_ingest(builtin, state).await,
996                    "kafka-verify-data" => kafka::run_verify_data(builtin, state).await,
997                    "kafka-verify-commit" => kafka::run_verify_commit(builtin, state).await,
998                    "kafka-verify-topic" => kafka::run_verify_topic(builtin, state).await,
999                    "mysql-connect" => mysql::run_connect(builtin, state).await,
1000                    "mysql-execute" => mysql::run_execute(builtin, state).await,
1001                    "nop" => nop::run_nop(),
1002                    "postgres-connect" => postgres::run_connect(builtin, state).await,
1003                    "postgres-execute" => postgres::run_execute(builtin, state).await,
1004                    "postgres-verify-slot" => postgres::run_verify_slot(builtin, state).await,
1005                    "protobuf-compile-descriptors" => {
1006                        protobuf::run_compile_descriptors(builtin, state).await
1007                    }
1008                    "psql-execute" => psql::run_execute(builtin, state).await,
1009                    "s3-verify-data" => s3::run_verify_data(builtin, state).await,
1010                    "s3-verify-keys" => s3::run_verify_keys(builtin, state).await,
1011                    "s3-file-upload" => s3::run_upload(builtin, state).await,
1012                    "s3-set-presigned-url" => s3::run_set_presigned_url(builtin, state).await,
1013                    "s3-upload-parquet-types" => s3::run_upload_parquet_types(builtin, state).await,
1014                    "s3-upload-parquet-unsorted-map" => {
1015                        s3::run_upload_parquet_unsorted_map(builtin, state).await
1016                    }
1017                    "schema-registry-publish" => schema_registry::run_publish(builtin, state).await,
1018                    "schema-registry-verify" => schema_registry::run_verify(builtin, state).await,
1019                    "schema-registry-wait" => schema_registry::run_wait(builtin, state).await,
1020                    "skip-if" => skip_if::run_skip_if(builtin, state).await,
1021                    "skip-end" => skip_end::run_skip_end(),
1022                    "sql-server-connect" => sql_server::run_connect(builtin, state).await,
1023                    "sql-server-execute" => sql_server::run_execute(builtin, state).await,
1024                    "sql-server-set-from-sql" => sql_server::run_set_from_sql(builtin, state).await,
1025                    "persist-force-compaction" => {
1026                        persist::run_force_compaction(builtin, state).await
1027                    }
1028                    "random-sleep" => sleep::run_random_sleep(builtin),
1029                    "set-regex" => set::run_regex_set(builtin, state),
1030                    "unset-regex" => set::run_regex_unset(builtin, state),
1031                    "set-sql-timeout" => set::run_sql_timeout(builtin, state),
1032                    "set-max-tries" => set::run_max_tries(builtin, state),
1033                    "sleep-is-probably-flaky-i-have-justified-my-need-with-a-comment" => {
1034                        sleep::run_sleep(builtin)
1035                    }
1036                    "set" => set::set_vars(builtin, state),
1037                    "set-arg-default" => set::run_set_arg_default(builtin, state),
1038                    "set-from-sql" => set::run_set_from_sql(builtin, state).await,
1039                    "set-from-file" => set::run_set_from_file(builtin, state).await,
1040                    "webhook-append" => webhook::run_append(builtin, state).await,
1041                    _ => {
1042                        return Err(PosError::new(
1043                            anyhow!("unknown built-in command {}", builtin.name),
1044                            self.pos,
1045                        ));
1046                    }
1047                }
1048            }
1049            Command::Sql(mut sql, version_constraint) => {
1050                handle_version!(version_constraint);
1051                sql.query = subst(&sql.query, &state.cmd_vars)?;
1052                if let SqlOutput::Full { expected_rows, .. } = &mut sql.expected_output {
1053                    for row in expected_rows {
1054                        for col in row {
1055                            *col = subst(col, &state.cmd_vars)?;
1056                        }
1057                    }
1058                }
1059                sql::run_sql(sql, state).await
1060            }
1061            Command::FailSql(mut sql, version_constraint) => {
1062                handle_version!(version_constraint);
1063                sql.query = subst(&sql.query, &state.cmd_vars)?;
1064                sql.expected_error = match &sql.expected_error {
1065                    SqlExpectedError::Contains(s) => {
1066                        SqlExpectedError::Contains(subst(s, &state.cmd_vars)?)
1067                    }
1068                    SqlExpectedError::Exact(s) => {
1069                        SqlExpectedError::Exact(subst(s, &state.cmd_vars)?)
1070                    }
1071                    SqlExpectedError::Regex(s) => {
1072                        SqlExpectedError::Regex(subst_re(s, &state.cmd_vars)?)
1073                    }
1074                    SqlExpectedError::Timeout => SqlExpectedError::Timeout,
1075                };
1076                sql.expected_detail = match sql.expected_detail {
1077                    Some(s) => Some(subst(&s, &state.cmd_vars)?),
1078                    None => None,
1079                };
1080                sql.expected_hint = match sql.expected_hint {
1081                    Some(s) => Some(subst(&s, &state.cmd_vars)?),
1082                    None => None,
1083                };
1084                sql::run_fail_sql(sql, state).await
1085            }
1086        };
1087
1088        r.map_err(wrap_err)
1089    }
1090}
1091
1092/// Substituted `${}`-delimited variables from `vars` into `msg`
1093fn substitute_vars(
1094    msg: &str,
1095    vars: &BTreeMap<String, String>,
1096    ignore_prefix: &Option<String>,
1097    regex_escape: bool,
1098) -> Result<String, anyhow::Error> {
1099    static RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\$\{([^}]+)\}").unwrap());
1100    let mut err = None;
1101    let out = RE.replace_all(msg, |caps: &Captures| {
1102        let name = &caps[1];
1103        if let Some(ignore_prefix) = &ignore_prefix {
1104            if name.starts_with(format!("{}.", ignore_prefix).as_str()) {
1105                // Do not substitute, leave original variable name in place
1106                return caps.get(0).unwrap().as_str().to_string();
1107            }
1108        }
1109
1110        if let Some(val) = vars.get(name) {
1111            if regex_escape {
1112                regex::escape(val)
1113            } else {
1114                val.to_string()
1115            }
1116        } else {
1117            err = Some(anyhow!("unknown variable: {}", name));
1118            "#VAR-MISSING#".to_string()
1119        }
1120    });
1121    match err {
1122        Some(err) => Err(err),
1123        None => Ok(out.into_owned()),
1124    }
1125}
1126
1127/// Initializes a [`State`] object by connecting to the various external
1128/// services specified in `config`.
1129///
1130/// Returns the initialized `State` and a cleanup future. The cleanup future
1131/// should be `await`ed only *after* dropping the `State` to check whether any
1132/// errors occured while dropping the `State`. This awkward API is a workaround
1133/// for the lack of `AsyncDrop` support in Rust.
1134pub async fn create_state(
1135    config: &Config,
1136) -> Result<(State, impl Future<Output = Result<(), anyhow::Error>>), anyhow::Error> {
1137    let seed = config
1138        .seed
1139        .clone()
1140        .unwrap_or_else(|| format!("{:010}", rand::random::<u32>()));
1141
1142    let (_tempfile, temp_path) = match &config.temp_dir {
1143        Some(temp_dir) => {
1144            fs::create_dir_all(temp_dir).context("creating temporary directory")?;
1145            (None, PathBuf::from(&temp_dir))
1146        }
1147        _ => {
1148            // Stash the tempfile object so that it does not go out of scope and delete
1149            // the tempdir prematurely
1150            let tempfile_handle = tempfile::tempdir().context("creating temporary directory")?;
1151            let temp_path = tempfile_handle.path().to_path_buf();
1152            (Some(tempfile_handle), temp_path)
1153        }
1154    };
1155
1156    let materialize_catalog_config = config.materialize_catalog_config.clone();
1157
1158    let materialize_url = util::postgres::config_url(&config.materialize_pgconfig)?;
1159    info!("Connecting to {}", materialize_url.as_str());
1160    let (pgclient, pgconn) = Retry::default()
1161        .max_duration(config.default_timeout)
1162        .retry_async_canceling(|_| async move {
1163            let mut pgconfig = config.materialize_pgconfig.clone();
1164            pgconfig.connect_timeout(config.default_timeout);
1165            let tls = make_tls(&pgconfig)?;
1166            pgconfig.connect(tls).await.map_err(|e| anyhow!(e))
1167        })
1168        .await?;
1169
1170    let pgconn_task =
1171        task::spawn(|| "pgconn_task", pgconn).map(|join| join.context("running SQL connection"));
1172
1173    let materialize_state =
1174        create_materialize_state(&config, materialize_catalog_config, pgclient).await?;
1175
1176    let schema_registry_url = config.schema_registry_url.to_owned();
1177
1178    let ccsr_client = {
1179        let mut ccsr_config = mz_ccsr::ClientConfig::new(schema_registry_url.clone());
1180
1181        if let Some(cert_path) = &config.cert_path {
1182            let cert = fs::read(cert_path).context("reading cert")?;
1183            let pass = config.cert_password.as_deref().unwrap_or("").to_owned();
1184            let ident = mz_ccsr::tls::Identity::from_pkcs12_der(cert, pass)
1185                .context("reading keystore file as pkcs12")?;
1186            ccsr_config = ccsr_config.identity(ident);
1187        }
1188
1189        if let Some(ccsr_username) = &config.ccsr_username {
1190            ccsr_config = ccsr_config.auth(ccsr_username.clone(), config.ccsr_password.clone());
1191        }
1192
1193        ccsr_config.build().context("Creating CCSR client")?
1194    };
1195
1196    let (kafka_addr, kafka_admin, kafka_admin_opts, kafka_producer, kafka_topics, kafka_config) = {
1197        use rdkafka::admin::{AdminClient, AdminOptions};
1198        use rdkafka::producer::FutureProducer;
1199
1200        let mut kafka_config = create_new_client_config_simple();
1201        kafka_config.set("bootstrap.servers", &config.kafka_addr);
1202        kafka_config.set("group.id", "materialize-testdrive");
1203        kafka_config.set("auto.offset.reset", "earliest");
1204        kafka_config.set("isolation.level", "read_committed");
1205        if let Some(cert_path) = &config.cert_path {
1206            kafka_config.set("security.protocol", "ssl");
1207            kafka_config.set("ssl.keystore.location", cert_path);
1208            if let Some(cert_password) = &config.cert_password {
1209                kafka_config.set("ssl.keystore.password", cert_password);
1210            }
1211        }
1212        kafka_config.set("message.max.bytes", "15728640");
1213
1214        for (key, value) in &config.kafka_opts {
1215            kafka_config.set(key, value);
1216        }
1217
1218        let admin: AdminClient<_> = kafka_config
1219            .create_with_context(MzClientContext::default())
1220            .with_context(|| format!("opening Kafka connection: {}", config.kafka_addr))?;
1221
1222        let admin_opts = AdminOptions::new().operation_timeout(Some(config.default_timeout));
1223
1224        let producer: FutureProducer<_> = kafka_config
1225            .create_with_context(MzClientContext::default())
1226            .with_context(|| format!("opening Kafka producer connection: {}", config.kafka_addr))?;
1227
1228        let topics = BTreeMap::new();
1229
1230        (
1231            config.kafka_addr.to_owned(),
1232            admin,
1233            admin_opts,
1234            producer,
1235            topics,
1236            kafka_config,
1237        )
1238    };
1239
1240    let mut state = State {
1241        config: config.clone(),
1242
1243        // === Testdrive state. ===
1244        arg_vars: config.arg_vars.clone(),
1245        cmd_vars: BTreeMap::new(),
1246        seed,
1247        temp_path,
1248        _tempfile,
1249        default_timeout: config.default_timeout,
1250        timeout: config.default_timeout,
1251        max_tries: config.default_max_tries,
1252        initial_backoff: config.initial_backoff,
1253        backoff_factor: config.backoff_factor,
1254        consistency_checks: config.consistency_checks,
1255        consistency_check_timeout: config.consistency_check_timeout,
1256        check_statement_logging: config.check_statement_logging,
1257        consistency_checks_adhoc_skip: false,
1258        regex: None,
1259        regex_replacement: set::DEFAULT_REGEX_REPLACEMENT.into(),
1260        rewrite_results: config.rewrite_results,
1261        error_line_count: 0,
1262        error_string: "".to_string(),
1263
1264        // === Materialize state. ===
1265        materialize: materialize_state,
1266
1267        // === Persist state. ===
1268        persist_consensus_url: config.persist_consensus_url.clone(),
1269        persist_blob_url: config.persist_blob_url.clone(),
1270        build_info: config.build_info,
1271        persist_clients: PersistClientCache::new(
1272            PersistConfig::new_default_configs(config.build_info, SYSTEM_TIME.clone()),
1273            &MetricsRegistry::new(),
1274            |_, _| PubSubClientConnection::noop(),
1275        ),
1276
1277        // === Confluent state. ===
1278        schema_registry_url,
1279        ccsr_client,
1280        kafka_addr,
1281        kafka_admin,
1282        kafka_admin_opts,
1283        kafka_config,
1284        kafka_default_partitions: config.kafka_default_partitions,
1285        kafka_producer,
1286        kafka_topics,
1287        kafka_verify_topics: BTreeSet::new(),
1288
1289        // === AWS state. ===
1290        aws_account: config.aws_account.clone(),
1291        aws_config: config.aws_config.clone(),
1292
1293        // === Database driver state. ===
1294        duckdb_clients: BTreeMap::new(),
1295        mysql_clients: BTreeMap::new(),
1296        postgres_clients: BTreeMap::new(),
1297        sql_server_clients: BTreeMap::new(),
1298        background_tasks: Vec::new(),
1299
1300        // === Fivetran state. ===
1301        fivetran_destination_url: config.fivetran_destination_url.clone(),
1302        fivetran_destination_files_path: config.fivetran_destination_files_path.clone(),
1303
1304        rewrites: Vec::new(),
1305        rewrite_pos_start: 0,
1306        rewrite_pos_end: 0,
1307    };
1308    state.initialize_cmd_vars().await?;
1309    Ok((state, pgconn_task))
1310}
1311
1312async fn create_materialize_state(
1313    config: &&Config,
1314    materialize_catalog_config: Option<CatalogConfig>,
1315    pgclient: tokio_postgres::Client,
1316) -> Result<MaterializeState, anyhow::Error> {
1317    let materialize_url = util::postgres::config_url(&config.materialize_pgconfig)?;
1318    let materialize_internal_url =
1319        util::postgres::config_url(&config.materialize_internal_pgconfig)?;
1320
1321    for (key, value) in &config.materialize_params {
1322        // Session parameter values are raw SQL fragments from testdrive config.
1323        #[allow(clippy::disallowed_methods)]
1324        pgclient
1325            .batch_execute(&format!("SET {key} = {value}"))
1326            .await
1327            .context("setting session parameter")?;
1328    }
1329
1330    let materialize_user = config
1331        .materialize_pgconfig
1332        .get_user()
1333        .expect("testdrive URL must contain user")
1334        .to_string();
1335
1336    let materialize_sql_addr = format!(
1337        "{}:{}",
1338        materialize_url.host_str().unwrap(),
1339        materialize_url.port().unwrap()
1340    );
1341    let materialize_http_addr = format!(
1342        "{}:{}",
1343        materialize_url.host_str().unwrap(),
1344        config.materialize_http_port
1345    );
1346    let materialize_internal_sql_addr = format!(
1347        "{}:{}",
1348        materialize_internal_url.host_str().unwrap(),
1349        materialize_internal_url.port().unwrap()
1350    );
1351    let materialize_password_sql_addr = format!(
1352        "{}:{}",
1353        materialize_url.host_str().unwrap(),
1354        config.materialize_password_sql_port
1355    );
1356    let materialize_sasl_sql_addr = format!(
1357        "{}:{}",
1358        materialize_url.host_str().unwrap(),
1359        config.materialize_sasl_sql_port
1360    );
1361    let materialize_internal_http_addr = format!(
1362        "{}:{}",
1363        materialize_internal_url.host_str().unwrap(),
1364        config.materialize_internal_http_port
1365    );
1366    let environment_id = pg_query_one(&pgclient, pg_sql!("SELECT mz_environment_id()"), &[])
1367        .await?
1368        .get::<_, String>(0)
1369        .parse()
1370        .context("parsing environment ID")?;
1371
1372    // Tolerate the functions not existing: upgrade tests run testdrive against an
1373    // older environmentd whose catalog predates them. A version without them has
1374    // no context-folding views either, so no context is the correct fold.
1375    let (aws_account_id, aws_external_id_prefix, aws_connection_role_arn) = match pg_query_one(
1376        &pgclient,
1377        pg_sql!(
1378            "SELECT mz_aws_account_id(), mz_aws_external_id_prefix(), \
1379                 mz_aws_connection_role_arn()"
1380        ),
1381        &[],
1382    )
1383    .await
1384    {
1385        Ok(row) => (
1386            row.get::<_, Option<String>>(0),
1387            row.get::<_, Option<String>>(1),
1388            row.get::<_, Option<String>>(2),
1389        ),
1390        Err(_) => (None, None, None),
1391    };
1392
1393    let bootstrap_args = BootstrapArgs {
1394        cluster_replica_size_map: config.materialize_cluster_replica_sizes.clone(),
1395        default_cluster_replica_size: "ABC".to_string(),
1396        default_cluster_replication_factor: 1,
1397        bootstrap_role: None,
1398    };
1399
1400    let materialize_state = MaterializeState {
1401        catalog_config: materialize_catalog_config,
1402        sql_addr: materialize_sql_addr,
1403        use_https: config.materialize_use_https,
1404        http_addr: materialize_http_addr,
1405        internal_sql_addr: materialize_internal_sql_addr,
1406        internal_http_addr: materialize_internal_http_addr,
1407        password_sql_addr: materialize_password_sql_addr,
1408        sasl_sql_addr: materialize_sasl_sql_addr,
1409        user: materialize_user,
1410        pgclient,
1411        environment_id,
1412        bootstrap_args,
1413        aws_account_id,
1414        aws_external_id_prefix,
1415        aws_connection_role_arn,
1416    };
1417
1418    Ok(materialize_state)
1419}