1use 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#[derive(Debug, Clone)]
95pub struct Config {
96 pub arg_vars: BTreeMap<String, String>,
102 pub seed: Option<String>,
104 pub reset: bool,
107 pub temp_dir: Option<String>,
112 pub source: Option<String>,
114 pub default_timeout: Duration,
116 pub default_max_tries: usize,
118 pub initial_backoff: Duration,
122 pub backoff_factor: f64,
126 pub consistency_checks: consistency::Level,
128 pub consistency_check_timeout: Duration,
130 pub check_statement_logging: bool,
133 pub rewrite_results: bool,
135
136 pub materialize_pgconfig: tokio_postgres::Config,
140 pub materialize_internal_pgconfig: tokio_postgres::Config,
143 pub materialize_use_https: bool,
145 pub materialize_http_port: u16,
148 pub materialize_internal_http_port: u16,
151 pub materialize_password_sql_port: u16,
154 pub materialize_sasl_sql_port: u16,
157 pub materialize_params: Vec<(String, String)>,
159 pub materialize_catalog_config: Option<CatalogConfig>,
161 pub build_info: &'static BuildInfo,
163 pub materialize_cluster_replica_sizes: ClusterReplicaSizeMap,
165
166 pub persist_consensus_url: Option<SensitiveUrl>,
169 pub persist_blob_url: Option<SensitiveUrl>,
171
172 pub kafka_addr: String,
175 pub kafka_default_partitions: usize,
177 pub kafka_opts: Vec<(String, String)>,
180 pub schema_registry_url: Url,
182 pub cert_path: Option<String>,
187 pub cert_password: Option<String>,
189 pub ccsr_username: Option<String>,
192 pub ccsr_password: Option<String>,
195
196 pub aws_config: SdkConfig,
199 pub aws_account: String,
201
202 pub fivetran_destination_url: String,
205 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 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 pub config: Config,
238
239 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: MaterializeState,
261
262 persist_consensus_url: Option<SensitiveUrl>,
264 persist_blob_url: Option<SensitiveUrl>,
265 build_info: &'static BuildInfo,
266 persist_clients: PersistClientCache,
267
268 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 kafka_verify_topics: BTreeSet<String>,
280
281 aws_account: String,
283 aws_config: SdkConfig,
284
285 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 background_tasks: Vec<BackgroundTask>,
293
294 fivetran_destination_url: String,
296 fivetran_destination_files_path: String,
297
298 rewrite_results: bool,
300 pub rewrites: Vec<Rewrite>,
302 pub rewrite_pos_start: usize,
304 pub rewrite_pos_end: usize,
306}
307
308pub struct Rewrite {
309 pub content: String,
310 pub start: usize,
311 pub end: usize,
312}
313
314pub(crate) struct BackgroundTask {
325 desc: String,
326 handle: task::JoinHandle<Result<(), anyhow::Error>>,
327 cancel_token: CancelToken,
328 url: String,
331}
332
333async 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 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 pub fn clear_skip_consistency_checks(&mut self) -> bool {
515 std::mem::replace(&mut self.consistency_checks_adhoc_skip, false)
516 }
517
518 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 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 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 {
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 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 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 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 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 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 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 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 }
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 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 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#[derive(Debug, Clone)]
911pub struct CatalogConfig {
912 pub persist_consensus_url: SensitiveUrl,
914 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 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
1092fn 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 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
1127pub 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 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 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: materialize_state,
1266
1267 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 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_account: config.aws_account.clone(),
1291 aws_config: config.aws_config.clone(),
1292
1293 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_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 #[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 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}