Skip to main content

mz_sqllogictest/
runner.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
10//! The Materialize-specific runner for sqllogictest.
11//!
12//! slt tests expect a serialized execution of sql statements and queries.
13//! To get the same results in materialize we track current_timestamp and increment it whenever we execute a statement.
14//!
15//! The high-level workflow is:
16//!   for each record in the test file:
17//!     if record is a sql statement:
18//!       run sql in postgres, observe changes and copy them to materialize using LocalInput::Updates(..)
19//!       advance current_timestamp
20//!       promise to never send updates for times < current_timestamp using LocalInput::Watermark(..)
21//!       compare to expected results
22//!       if wrong, bail out and stop processing this file
23//!     if record is a sql query:
24//!       peek query at current_timestamp
25//!       compare to expected results
26//!       if wrong, record the error
27
28use std::collections::BTreeMap;
29use std::error::Error;
30use std::fs::{File, OpenOptions};
31use std::io::{Read, Seek, SeekFrom, Write};
32use std::net::{IpAddr, Ipv4Addr, SocketAddr};
33use std::path::Path;
34use std::sync::Arc;
35use std::sync::LazyLock;
36use std::time::Duration;
37use std::{env, fmt, ops, str, thread};
38
39use anyhow::{anyhow, bail};
40use bytes::BytesMut;
41use chrono::{DateTime, NaiveDateTime, NaiveTime, Utc};
42use futures::sink::SinkExt;
43use itertools::Itertools;
44use maplit::btreemap;
45use md5::{Digest, Md5};
46use mz_adapter_types::bootstrap_builtin_cluster_config::{
47    ANALYTICS_CLUSTER_DEFAULT_REPLICATION_FACTOR, BootstrapBuiltinClusterConfig,
48    CATALOG_SERVER_CLUSTER_DEFAULT_REPLICATION_FACTOR, PROBE_CLUSTER_DEFAULT_REPLICATION_FACTOR,
49    SUPPORT_CLUSTER_DEFAULT_REPLICATION_FACTOR, SYSTEM_CLUSTER_DEFAULT_REPLICATION_FACTOR,
50};
51use mz_catalog::config::ClusterReplicaSizeMap;
52use mz_controller::{ControllerConfig, ReplicaHttpLocator};
53use mz_environmentd::CatalogConfig;
54use mz_license_keys::ValidatedLicenseKey;
55use mz_orchestrator_process::{ProcessOrchestrator, ProcessOrchestratorConfig};
56use mz_orchestrator_tracing::{TracingCliArgs, TracingOrchestrator};
57use mz_ore::cast::{CastFrom, ReinterpretCast};
58use mz_ore::channel::trigger;
59use mz_ore::error::ErrorExt;
60use mz_ore::metrics::MetricsRegistry;
61use mz_ore::now::SYSTEM_TIME;
62use mz_ore::retry::Retry;
63use mz_ore::sql;
64use mz_ore::sql::Sql;
65use mz_ore::task;
66use mz_ore::thread::{JoinHandleExt, JoinOnDropHandle};
67use mz_ore::tracing::TracingHandle;
68use mz_ore::url::SensitiveUrl;
69use mz_persist_client::PersistLocation;
70use mz_persist_client::cache::PersistClientCache;
71use mz_persist_client::cfg::PersistConfig;
72use mz_persist_client::rpc::{
73    MetricsSameProcessPubSubSender, PersistGrpcPubSubServer, PubSubClientConnection, PubSubSender,
74};
75use mz_pgrepr::{Interval, Jsonb, Numeric, UInt2, UInt4, UInt8, Value, oid};
76use mz_repr::ColumnName;
77use mz_repr::adt::date::Date;
78use mz_repr::adt::mz_acl_item::{AclItem, MzAclItem};
79use mz_repr::adt::numeric;
80use mz_secrets::SecretsController;
81use mz_server_core::listeners::v26_32_0::ListenersConfig;
82use mz_server_core::listeners::{
83    AllowedRoles, AuthenticatorKind, HttpListenerConfig, HttpRoutesEnabled, RouteGroup,
84    SqlListenerConfig,
85};
86use mz_sql::ast::{Expr, Raw, Statement};
87use mz_sql::catalog::EnvironmentId;
88use mz_sql_parser::ast::display::AstDisplay;
89use mz_sql_parser::ast::{
90    CreateIndexStatement, CreateViewStatement, CteBlock, Distinct, DropObjectsStatement, Ident,
91    IfExistsBehavior, ObjectType, OrderByExpr, Query, RawItemName, Select, SelectItem,
92    SelectStatement, SetExpr, Statement as AstStatement, TableFactor, TableWithJoins,
93    UnresolvedItemName, UnresolvedObjectName, ViewDefinition,
94};
95use mz_sql_parser::parser;
96use mz_storage_types::connections::ConnectionContext;
97use postgres_protocol::types;
98use regex::Regex;
99use tempfile::TempDir;
100use tokio::net::TcpListener;
101use tokio::runtime::Runtime;
102use tokio::sync::oneshot;
103use tokio_postgres::fallible_iterator::FallibleIterator;
104use tokio_postgres::types::{FromSql, Kind as PgKind, Type as PgType};
105use tokio_postgres::{NoTls, Row, SimpleQueryMessage};
106use tokio_stream::wrappers::TcpListenerStream;
107use tower_http::cors::AllowOrigin;
108use tracing::{error, info};
109use uuid::Uuid;
110use uuid::fmt::Simple;
111
112use crate::ast::{Location, Mode, Output, QueryOutput, Record, Sort, Type};
113use crate::util;
114
115#[derive(Debug)]
116pub enum Outcome<'a> {
117    Unsupported {
118        error: anyhow::Error,
119        location: Location,
120    },
121    ParseFailure {
122        error: anyhow::Error,
123        location: Location,
124    },
125    PlanFailure {
126        error: anyhow::Error,
127        expected_error: Option<String>,
128        location: Location,
129    },
130    UnexpectedPlanSuccess {
131        expected_error: &'a str,
132        location: Location,
133    },
134    WrongNumberOfRowsInserted {
135        expected_count: u64,
136        actual_count: u64,
137        location: Location,
138    },
139    WrongColumnCount {
140        expected_count: usize,
141        actual_count: usize,
142        location: Location,
143    },
144    WrongColumnNames {
145        expected_column_names: &'a Vec<ColumnName>,
146        actual_column_names: Vec<ColumnName>,
147        actual_output: Output,
148        location: Location,
149    },
150    OutputFailure {
151        expected_output: &'a Output,
152        actual_raw_output: Vec<Row>,
153        actual_output: Output,
154        location: Location,
155    },
156    InconsistentViewOutcome {
157        query_outcome: Box<Outcome<'a>>,
158        view_outcome: Box<Outcome<'a>>,
159        location: Location,
160    },
161    Bail {
162        cause: Box<Outcome<'a>>,
163        location: Location,
164    },
165    Warning {
166        cause: Box<Outcome<'a>>,
167        location: Location,
168    },
169    Success,
170}
171
172const NUM_OUTCOMES: usize = 12;
173const WARNING_OUTCOME: usize = NUM_OUTCOMES - 2;
174const SUCCESS_OUTCOME: usize = NUM_OUTCOMES - 1;
175
176impl<'a> Outcome<'a> {
177    fn code(&self) -> usize {
178        match self {
179            Outcome::Unsupported { .. } => 0,
180            Outcome::ParseFailure { .. } => 1,
181            Outcome::PlanFailure { .. } => 2,
182            Outcome::UnexpectedPlanSuccess { .. } => 3,
183            Outcome::WrongNumberOfRowsInserted { .. } => 4,
184            Outcome::WrongColumnCount { .. } => 5,
185            Outcome::WrongColumnNames { .. } => 6,
186            Outcome::OutputFailure { .. } => 7,
187            Outcome::InconsistentViewOutcome { .. } => 8,
188            Outcome::Bail { .. } => 9,
189            Outcome::Warning { .. } => 10,
190            Outcome::Success => 11,
191        }
192    }
193
194    fn success(&self) -> bool {
195        matches!(self, Outcome::Success)
196    }
197
198    fn failure(&self) -> bool {
199        !matches!(self, Outcome::Success) && !matches!(self, Outcome::Warning { .. })
200    }
201
202    /// Returns an error message that will match self. Appropriate for
203    /// rewriting error messages (i.e. not inserting error messages where we
204    /// currently expect success).
205    fn err_msg(&self) -> Option<String> {
206        match self {
207            Outcome::Unsupported { error, .. }
208            | Outcome::ParseFailure { error, .. }
209            | Outcome::PlanFailure { error, .. } => {
210                // Take only the first line, which should be sufficient for
211                // meaningfully matching the error.
212                let err_str = error.to_string_with_causes();
213                let err_str = err_str.split('\n').next().unwrap();
214                // Strip the "db error: ERROR: " prefix added by the postgres
215                // client library, as it's noisy and not useful for matching.
216                let err_str = err_str.strip_prefix("db error: ERROR: ").unwrap_or(err_str);
217                // This value gets fed back into regex to check that it matches
218                // `self`, so escape its meta characters. We need to undo the
219                // escaping of #. `regex::escape` escapes this because it
220                // expects that we use the `x` flag when building a regex, but
221                // this is not the case, so \# would end up being an invalid
222                // escape sequence, which would choke the parsing of the slt
223                // file the next time around.
224                Some(regex::escape(err_str).replace(r"\#", "#"))
225            }
226            _ => None,
227        }
228    }
229}
230
231impl fmt::Display for Outcome<'_> {
232    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
233        use Outcome::*;
234        const INDENT: &str = "\n        ";
235        match self {
236            Unsupported { error, location } => write!(
237                f,
238                "Unsupported:{}:\n{}",
239                location,
240                error.display_with_causes()
241            ),
242            ParseFailure { error, location } => {
243                write!(
244                    f,
245                    "ParseFailure:{}:\n{}",
246                    location,
247                    error.display_with_causes()
248                )
249            }
250            PlanFailure {
251                error,
252                expected_error,
253                location,
254            } => {
255                if let Some(expected_error) = expected_error {
256                    write!(
257                        f,
258                        "PlanFailure:{}:\nerror does not match expected pattern:\n  expected: /{}/\n  actual:    {}",
259                        location,
260                        expected_error,
261                        error.display_with_causes()
262                    )
263                } else {
264                    write!(f, "PlanFailure:{}:\n{:#}", location, error)
265                }
266            }
267            UnexpectedPlanSuccess {
268                expected_error,
269                location,
270            } => write!(
271                f,
272                "UnexpectedPlanSuccess:{} expected error: {}",
273                location, expected_error
274            ),
275            WrongNumberOfRowsInserted {
276                expected_count,
277                actual_count,
278                location,
279            } => write!(
280                f,
281                "WrongNumberOfRowsInserted:{}{}expected: {}{}actually: {}",
282                location, INDENT, expected_count, INDENT, actual_count
283            ),
284            WrongColumnCount {
285                expected_count,
286                actual_count,
287                location,
288            } => write!(
289                f,
290                "WrongColumnCount:{}{}expected: {}{}actually: {}",
291                location, INDENT, expected_count, INDENT, actual_count
292            ),
293            WrongColumnNames {
294                expected_column_names,
295                actual_column_names,
296                actual_output: _,
297                location,
298            } => write!(
299                f,
300                "Wrong Column Names:{}:{}expected column names: {}{}inferred column names: {}",
301                location,
302                INDENT,
303                expected_column_names
304                    .iter()
305                    .map(|n| n.to_string())
306                    .collect::<Vec<_>>()
307                    .join(" "),
308                INDENT,
309                actual_column_names
310                    .iter()
311                    .map(|n| n.to_string())
312                    .collect::<Vec<_>>()
313                    .join(" ")
314            ),
315            OutputFailure {
316                expected_output,
317                actual_raw_output,
318                actual_output,
319                location,
320            } => write!(
321                f,
322                "OutputFailure:{}{}expected: {:?}{}actually: {:?}{}actual raw: {:?}",
323                location, INDENT, expected_output, INDENT, actual_output, INDENT, actual_raw_output
324            ),
325            InconsistentViewOutcome {
326                query_outcome,
327                view_outcome,
328                location,
329            } => write!(
330                f,
331                "InconsistentViewOutcome:{}{}expected from query: {}{}actually from indexed view: {}",
332                location, INDENT, query_outcome, INDENT, view_outcome
333            ),
334            Bail { cause, location } => write!(f, "Bail:{} {}", location, cause),
335            Warning { cause, location } => write!(f, "Warning:{} {}", location, cause),
336            Success => f.write_str("Success"),
337        }
338    }
339}
340
341#[derive(Default, Debug)]
342pub struct Outcomes {
343    stats: [usize; NUM_OUTCOMES],
344    details: Vec<String>,
345}
346
347impl ops::AddAssign<Outcomes> for Outcomes {
348    fn add_assign(&mut self, rhs: Outcomes) {
349        for (lhs, rhs) in self.stats.iter_mut().zip_eq(rhs.stats.iter()) {
350            *lhs += rhs
351        }
352    }
353}
354impl Outcomes {
355    pub fn any_failed(&self) -> bool {
356        self.stats[SUCCESS_OUTCOME] + self.stats[WARNING_OUTCOME] < self.stats.iter().sum::<usize>()
357    }
358
359    pub fn as_json(&self) -> serde_json::Value {
360        serde_json::json!({
361            "unsupported": self.stats[0],
362            "parse_failure": self.stats[1],
363            "plan_failure": self.stats[2],
364            "unexpected_plan_success": self.stats[3],
365            "wrong_number_of_rows_affected": self.stats[4],
366            "wrong_column_count": self.stats[5],
367            "wrong_column_names": self.stats[6],
368            "output_failure": self.stats[7],
369            "inconsistent_view_outcome": self.stats[8],
370            "bail": self.stats[9],
371            "warning": self.stats[10],
372            "success": self.stats[11],
373        })
374    }
375
376    pub fn display(&self, no_fail: bool, failure_details: bool) -> OutcomesDisplay<'_> {
377        OutcomesDisplay {
378            inner: self,
379            no_fail,
380            failure_details,
381        }
382    }
383}
384
385pub struct OutcomesDisplay<'a> {
386    inner: &'a Outcomes,
387    no_fail: bool,
388    failure_details: bool,
389}
390
391impl<'a> fmt::Display for OutcomesDisplay<'a> {
392    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
393        let total: usize = self.inner.stats.iter().sum();
394        if self.failure_details
395            && (self.inner.stats[SUCCESS_OUTCOME] + self.inner.stats[WARNING_OUTCOME] != total
396                || self.no_fail)
397        {
398            for outcome in &self.inner.details {
399                writeln!(f, "{}", outcome)?;
400            }
401            Ok(())
402        } else {
403            write!(
404                f,
405                "{}:",
406                if self.inner.stats[SUCCESS_OUTCOME] + self.inner.stats[WARNING_OUTCOME] == total {
407                    "PASS"
408                } else if self.no_fail {
409                    "FAIL-IGNORE"
410                } else {
411                    "FAIL"
412                }
413            )?;
414            static NAMES: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
415                vec![
416                    "unsupported",
417                    "parse-failure",
418                    "plan-failure",
419                    "unexpected-plan-success",
420                    "wrong-number-of-rows-inserted",
421                    "wrong-column-count",
422                    "wrong-column-names",
423                    "output-failure",
424                    "inconsistent-view-outcome",
425                    "bail",
426                    "warning",
427                    "success",
428                    "total",
429                ]
430            });
431            for (i, n) in self.inner.stats.iter().enumerate() {
432                if *n > 0 {
433                    write!(f, " {}={}", NAMES[i], n)?;
434                }
435            }
436            write!(f, " total={}", total)
437        }
438    }
439}
440
441struct QueryInfo {
442    is_select: bool,
443    num_attributes: Option<usize>,
444    /// Whether the `SELECT` carries an `AS OF` clause. Such queries are
445    /// excluded from the `--auto-index-selects` consistency check: re-running
446    /// them against a freshly created indexed view is inherently racy, because
447    /// the view's `since` advances with time and can move past a historical
448    /// `AS OF`, yielding a spurious "could not find a valid timestamp" failure.
449    has_as_of: bool,
450}
451
452enum PrepareQueryOutcome<'a> {
453    QueryPrepared(QueryInfo),
454    Outcome(Outcome<'a>),
455}
456
457pub struct Runner<'a> {
458    config: &'a RunConfig<'a>,
459    inner: Option<RunnerInner<'a>>,
460    /// Active `replace` substitutions, applied to actual query output before it
461    /// is compared/rewritten. Held on the outer `Runner` so they survive a
462    /// `reset-server`. See `Record::Replace`.
463    replacements: Vec<(Regex, String)>,
464}
465
466pub struct RunnerInner<'a> {
467    server_addr: SocketAddr,
468    internal_server_addr: SocketAddr,
469    password_server_addr: SocketAddr,
470    internal_http_server_addr: SocketAddr,
471    // Drop order matters for these fields.
472    client: tokio_postgres::Client,
473    system_client: tokio_postgres::Client,
474    clients: BTreeMap<String, tokio_postgres::Client>,
475    /// The user set by the most recent `user` directive, if any. Records run
476    /// on that user's connection (cached in `clients`) instead of `client`.
477    active_user: Option<String>,
478    auto_index_tables: bool,
479    auto_index_selects: bool,
480    auto_transactions: bool,
481    enable_table_keys: bool,
482    verbose: bool,
483    stdout: &'a dyn WriteFmt,
484    _shutdown_trigger: trigger::Trigger,
485    _server_thread: JoinOnDropHandle<()>,
486    _temp_dir: TempDir,
487}
488
489#[derive(Debug)]
490pub struct Slt(Value);
491
492impl<'a> FromSql<'a> for Slt {
493    fn from_sql(
494        ty: &PgType,
495        mut raw: &'a [u8],
496    ) -> Result<Self, Box<dyn Error + 'static + Send + Sync>> {
497        Ok(match *ty {
498            PgType::ACLITEM => Self(Value::AclItem(AclItem::decode_binary(
499                types::bytea_from_sql(raw),
500            )?)),
501            PgType::BOOL => Self(Value::Bool(types::bool_from_sql(raw)?)),
502            PgType::BYTEA => Self(Value::Bytea(types::bytea_from_sql(raw).to_vec())),
503            PgType::CHAR => Self(Value::Char(u8::from_be_bytes(
504                types::char_from_sql(raw)?.to_be_bytes(),
505            ))),
506            PgType::FLOAT4 => Self(Value::Float4(types::float4_from_sql(raw)?)),
507            PgType::FLOAT8 => Self(Value::Float8(types::float8_from_sql(raw)?)),
508            PgType::DATE => Self(Value::Date(Date::from_pg_epoch(types::int4_from_sql(
509                raw,
510            )?)?)),
511            PgType::INT2 => Self(Value::Int2(types::int2_from_sql(raw)?)),
512            PgType::INT4 => Self(Value::Int4(types::int4_from_sql(raw)?)),
513            PgType::INT8 => Self(Value::Int8(types::int8_from_sql(raw)?)),
514            PgType::INTERVAL => Self(Value::Interval(Interval::from_sql(ty, raw)?)),
515            PgType::JSONB => Self(Value::Jsonb(Jsonb::from_sql(ty, raw)?)),
516            PgType::NAME => Self(Value::Name(types::text_from_sql(raw)?.to_string())),
517            PgType::NUMERIC => Self(Value::Numeric(Numeric::from_sql(ty, raw)?)),
518            PgType::OID => Self(Value::Oid(types::oid_from_sql(raw)?)),
519            PgType::REGCLASS => Self(Value::Oid(types::oid_from_sql(raw)?)),
520            PgType::REGPROC => Self(Value::Oid(types::oid_from_sql(raw)?)),
521            PgType::REGTYPE => Self(Value::Oid(types::oid_from_sql(raw)?)),
522            PgType::TEXT | PgType::BPCHAR | PgType::VARCHAR => {
523                Self(Value::Text(types::text_from_sql(raw)?.to_string()))
524            }
525            PgType::TIME => Self(Value::Time(NaiveTime::from_sql(ty, raw)?)),
526            PgType::TIMESTAMP => Self(Value::Timestamp(
527                NaiveDateTime::from_sql(ty, raw)?.try_into()?,
528            )),
529            PgType::TIMESTAMPTZ => Self(Value::TimestampTz(
530                DateTime::<Utc>::from_sql(ty, raw)?.try_into()?,
531            )),
532            PgType::UUID => Self(Value::Uuid(Uuid::from_sql(ty, raw)?)),
533            PgType::RECORD => {
534                let num_fields = read_be_i32(&mut raw)?;
535                let mut tuple = vec![];
536                for _ in 0..num_fields {
537                    let oid = u32::reinterpret_cast(read_be_i32(&mut raw)?);
538                    let typ = match PgType::from_oid(oid) {
539                        Some(typ) => typ,
540                        None => return Err("unknown oid".into()),
541                    };
542                    let v = read_value::<Option<Slt>>(&typ, &mut raw)?;
543                    tuple.push(v.map(|v| v.0));
544                }
545                Self(Value::Record(tuple))
546            }
547            PgType::INT4_RANGE
548            | PgType::INT8_RANGE
549            | PgType::DATE_RANGE
550            | PgType::NUM_RANGE
551            | PgType::TS_RANGE
552            | PgType::TSTZ_RANGE => {
553                use mz_repr::adt::range::Range;
554                let range: Range<Slt> = Range::from_sql(ty, raw)?;
555                Self(Value::Range(range.into_bounds(|b| Box::new(b.0))))
556            }
557
558            _ => match ty.kind() {
559                PgKind::Array(arr_type) => {
560                    let arr = types::array_from_sql(raw)?;
561                    let elements: Vec<Option<Value>> = arr
562                        .values()
563                        .map(|v| match v {
564                            Some(v) => Ok(Some(Slt::from_sql(arr_type, v)?)),
565                            None => Ok(None),
566                        })
567                        .collect::<Vec<Option<Slt>>>()?
568                        .into_iter()
569                        // Map a Vec<Option<Slt>> to Vec<Option<Value>>.
570                        .map(|v| v.map(|v| v.0))
571                        .collect();
572
573                    Self(Value::Array {
574                        dims: arr
575                            .dimensions()
576                            .map(|d| {
577                                Ok(mz_repr::adt::array::ArrayDimension {
578                                    lower_bound: isize::cast_from(d.lower_bound),
579                                    length: usize::try_from(d.len)
580                                        .expect("cannot have negative length"),
581                                })
582                            })
583                            .collect()?,
584                        elements,
585                    })
586                }
587                _ => match ty.oid() {
588                    oid::TYPE_UINT2_OID => Self(Value::UInt2(UInt2::from_sql(ty, raw)?)),
589                    oid::TYPE_UINT4_OID => Self(Value::UInt4(UInt4::from_sql(ty, raw)?)),
590                    oid::TYPE_UINT8_OID => Self(Value::UInt8(UInt8::from_sql(ty, raw)?)),
591                    oid::TYPE_MZ_TIMESTAMP_OID => {
592                        let s = types::text_from_sql(raw)?;
593                        let t: mz_repr::Timestamp = s.parse()?;
594                        Self(Value::MzTimestamp(t))
595                    }
596                    oid::TYPE_MZ_ACL_ITEM_OID => Self(Value::MzAclItem(MzAclItem::decode_binary(
597                        types::bytea_from_sql(raw),
598                    )?)),
599                    _ => unreachable!(),
600                },
601            },
602        })
603    }
604    fn accepts(ty: &PgType) -> bool {
605        match ty.kind() {
606            PgKind::Array(_) | PgKind::Composite(_) => return true,
607            _ => {}
608        }
609        match ty.oid() {
610            oid::TYPE_UINT2_OID
611            | oid::TYPE_UINT4_OID
612            | oid::TYPE_UINT8_OID
613            | oid::TYPE_MZ_TIMESTAMP_OID
614            | oid::TYPE_MZ_ACL_ITEM_OID => return true,
615            _ => {}
616        }
617        matches!(
618            *ty,
619            PgType::ACLITEM
620                | PgType::BOOL
621                | PgType::BYTEA
622                | PgType::CHAR
623                | PgType::DATE
624                | PgType::FLOAT4
625                | PgType::FLOAT8
626                | PgType::INT2
627                | PgType::INT4
628                | PgType::INT8
629                | PgType::INTERVAL
630                | PgType::JSONB
631                | PgType::NAME
632                | PgType::NUMERIC
633                | PgType::OID
634                | PgType::REGCLASS
635                | PgType::REGPROC
636                | PgType::REGTYPE
637                | PgType::RECORD
638                | PgType::TEXT
639                | PgType::BPCHAR
640                | PgType::VARCHAR
641                | PgType::TIME
642                | PgType::TIMESTAMP
643                | PgType::TIMESTAMPTZ
644                | PgType::UUID
645                | PgType::INT4_RANGE
646                | PgType::INT4_RANGE_ARRAY
647                | PgType::INT8_RANGE
648                | PgType::INT8_RANGE_ARRAY
649                | PgType::DATE_RANGE
650                | PgType::DATE_RANGE_ARRAY
651                | PgType::NUM_RANGE
652                | PgType::NUM_RANGE_ARRAY
653                | PgType::TS_RANGE
654                | PgType::TS_RANGE_ARRAY
655                | PgType::TSTZ_RANGE
656                | PgType::TSTZ_RANGE_ARRAY
657        )
658    }
659}
660
661// From postgres-types/src/private.rs.
662fn read_be_i32(buf: &mut &[u8]) -> Result<i32, Box<dyn Error + Sync + Send>> {
663    if buf.len() < 4 {
664        return Err("invalid buffer size".into());
665    }
666    let mut bytes = [0; 4];
667    bytes.copy_from_slice(&buf[..4]);
668    *buf = &buf[4..];
669    Ok(i32::from_be_bytes(bytes))
670}
671
672// From postgres-types/src/private.rs.
673fn read_value<'a, T>(type_: &PgType, buf: &mut &'a [u8]) -> Result<T, Box<dyn Error + Sync + Send>>
674where
675    T: FromSql<'a>,
676{
677    let value = match usize::try_from(read_be_i32(buf)?) {
678        Err(_) => None,
679        Ok(len) => {
680            if len > buf.len() {
681                return Err("invalid buffer size".into());
682            }
683            let (head, tail) = buf.split_at(len);
684            *buf = tail;
685            Some(head)
686        }
687    };
688    T::from_sql_nullable(type_, value)
689}
690
691fn format_datum(d: Slt, typ: &Type, mode: Mode) -> String {
692    match (typ, d.0) {
693        (Type::Bool, Value::Bool(b)) => b.to_string(),
694
695        (Type::Integer, Value::Int2(i)) => i.to_string(),
696        (Type::Integer, Value::Int4(i)) => i.to_string(),
697        (Type::Integer, Value::Int8(i)) => i.to_string(),
698        (Type::Integer, Value::UInt2(u)) => u.0.to_string(),
699        (Type::Integer, Value::UInt4(u)) => u.0.to_string(),
700        (Type::Integer, Value::UInt8(u)) => u.0.to_string(),
701        (Type::Integer, Value::Oid(i)) => i.to_string(),
702        // TODO(benesch): rewrite to avoid `as`.
703        #[allow(clippy::as_conversions)]
704        (Type::Integer, Value::Float4(f)) => format!("{}", f as i64),
705        // TODO(benesch): rewrite to avoid `as`.
706        #[allow(clippy::as_conversions)]
707        (Type::Integer, Value::Float8(f)) => format!("{}", f as i64),
708        // This is so wrong, but sqlite needs it.
709        (Type::Integer, Value::Text(_)) => "0".to_string(),
710        (Type::Integer, Value::Bool(b)) => i8::from(b).to_string(),
711        (Type::Integer, Value::Numeric(d)) => {
712            let mut d = d.0.0.clone();
713            let mut cx = numeric::cx_datum();
714            // Truncate the decimal to match sqlite.
715            if mode == Mode::Standard {
716                cx.set_rounding(dec::Rounding::Down);
717            }
718            cx.round(&mut d);
719            numeric::munge_numeric(&mut d).unwrap();
720            d.to_standard_notation_string()
721        }
722
723        (Type::Real, Value::Int2(i)) => format!("{:.3}", i),
724        (Type::Real, Value::Int4(i)) => format!("{:.3}", i),
725        (Type::Real, Value::Int8(i)) => format!("{:.3}", i),
726        (Type::Real, Value::Float4(f)) => match mode {
727            Mode::Standard => format!("{:.3}", f),
728            Mode::Cockroach => format!("{}", f),
729        },
730        (Type::Real, Value::Float8(f)) => match mode {
731            Mode::Standard => format!("{:.3}", f),
732            Mode::Cockroach => format!("{}", f),
733        },
734        (Type::Real, Value::Numeric(d)) => match mode {
735            Mode::Standard => {
736                let mut d = d.0.0.clone();
737                if d.exponent() < -3 {
738                    numeric::rescale(&mut d, 3).unwrap();
739                }
740                numeric::munge_numeric(&mut d).unwrap();
741                d.to_standard_notation_string()
742            }
743            Mode::Cockroach => d.0.0.to_standard_notation_string(),
744        },
745
746        (Type::Text, Value::Text(s)) => {
747            if s.is_empty() {
748                "(empty)".to_string()
749            } else {
750                s
751            }
752        }
753        (Type::Text, Value::Bool(b)) => b.to_string(),
754        (Type::Text, Value::Float4(f)) => format!("{:.3}", f),
755        (Type::Text, Value::Float8(f)) => format!("{:.3}", f),
756        // Bytes are printed as text iff they are valid UTF-8. This
757        // seems guaranteed to confuse everyone, but it is required for
758        // compliance with the CockroachDB sqllogictest runner. [0]
759        //
760        // [0]: https://github.com/cockroachdb/cockroach/blob/970782487/pkg/sql/logictest/logic.go#L2038-L2043
761        (Type::Text, Value::Bytea(b)) => match str::from_utf8(&b) {
762            Ok(s) => s.to_string(),
763            Err(_) => format!("{:?}", b),
764        },
765        (Type::Text, Value::Numeric(d)) => d.0.0.to_standard_notation_string(),
766        // Everything else gets normal text encoding. This correctly handles things
767        // like arrays, tuples, and strings that need to be quoted.
768        (Type::Text, d) => {
769            let mut buf = BytesMut::new();
770            d.encode_text(&mut buf, mz_pgrepr::TextEncodeSettings::STABLE);
771            String::from_utf8_lossy(&buf).into_owned()
772        }
773
774        (Type::Oid, Value::Oid(o)) => o.to_string(),
775
776        // A mismatch between the declared column type and the value the query
777        // actually returned. Fall back to normal text encoding so the
778        // mismatch surfaces as an output diff rather than a panic that aborts
779        // the entire run.
780        (_, d) => {
781            let mut buf = BytesMut::new();
782            d.encode_text(&mut buf, mz_pgrepr::TextEncodeSettings::STABLE);
783            String::from_utf8_lossy(&buf).into_owned()
784        }
785    }
786}
787
788fn format_row(row: &Row, types: &[Type], mode: Mode) -> Vec<String> {
789    let mut formatted: Vec<String> = vec![];
790    for i in 0..row.len() {
791        let t: Option<Slt> = row.get::<usize, Option<Slt>>(i);
792        let t: Option<String> = t.map(|d| format_datum(d, &types[i], mode));
793        formatted.push(match t {
794            Some(t) => t,
795            None => "NULL".into(),
796        });
797    }
798
799    formatted
800}
801
802/// Reports whether `err` matches the `expected_error` regex of a record.
803///
804/// Expected errors in files imported from CockroachDB are not always valid
805/// regexes. An invalid regex cannot match, so the record fails with the usual
806/// mismatch outcome rather than aborting the entire run.
807fn error_matches(expected_error: &str, err: &str) -> bool {
808    match Regex::new(expected_error) {
809        Ok(re) => re.is_match(err),
810        Err(_) => false,
811    }
812}
813
814/// Removes CockroachDB-specific physical-layout items (`INDEX`,
815/// `UNIQUE INDEX`, `INVERTED INDEX`, and `FAMILY` definitions) from the
816/// column list of a `CREATE TABLE` statement. Constraints are kept, since
817/// dropping them would change the statement's semantics. Returns `None` if
818/// `sql` does not look like a `CREATE TABLE` statement.
819fn strip_crdb_table_items(sql: &str) -> Option<String> {
820    if !sql.trim_start().to_uppercase().starts_with("CREATE TABLE") {
821        return None;
822    }
823    let open = sql.find('(')?;
824    let mut depth = 1;
825    let mut in_string = false;
826    let mut in_ident = false;
827    let mut items: Vec<&str> = vec![];
828    let mut item_start = open + 1;
829    let mut close = None;
830    for (i, c) in sql[open + 1..].char_indices() {
831        let i = open + 1 + i;
832        match c {
833            '\'' if !in_ident => in_string = !in_string,
834            '"' if !in_string => in_ident = !in_ident,
835            _ if in_string || in_ident => (),
836            '(' => depth += 1,
837            ')' => {
838                depth -= 1;
839                if depth == 0 {
840                    close = Some(i);
841                    break;
842                }
843            }
844            ',' if depth == 1 => {
845                items.push(&sql[item_start..i]);
846                item_start = i + 1;
847            }
848            _ => (),
849        }
850    }
851    let close = close?;
852    items.push(&sql[item_start..close]);
853    fn first_word(s: &str) -> String {
854        s.trim_start()
855            .chars()
856            .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
857            .collect::<String>()
858            .to_uppercase()
859    }
860    let kept: Vec<&str> = items
861        .into_iter()
862        .filter(|item| {
863            let first = first_word(item);
864            let is_physical = match first.as_str() {
865                "INDEX" | "FAMILY" => true,
866                "UNIQUE" | "INVERTED" => {
867                    let rest = &item.trim_start()[first.len()..];
868                    first_word(rest) == "INDEX"
869                }
870                _ => false,
871            };
872            !is_physical
873        })
874        .collect();
875    Some(format!(
876        "{}{}{}",
877        &sql[..open + 1],
878        kept.join(","),
879        &sql[close..]
880    ))
881}
882
883impl<'a> Runner<'a> {
884    pub async fn start(config: &'a RunConfig<'a>) -> Result<Runner<'a>, anyhow::Error> {
885        let mut runner = Self {
886            config,
887            inner: None,
888            replacements: Vec::new(),
889        };
890        runner.reset().await?;
891        Ok(runner)
892    }
893
894    pub async fn reset(&mut self) -> Result<(), anyhow::Error> {
895        // Explicitly drop the old runner here to ensure that we wait for threads to terminate
896        // before starting a new runner
897        drop(self.inner.take());
898        self.inner = Some(RunnerInner::start(self.config).await?);
899
900        Ok(())
901    }
902
903    async fn run_record<'r>(
904        &mut self,
905        record: &'r Record<'r>,
906        in_transaction: &mut bool,
907    ) -> Result<Outcome<'r>, anyhow::Error> {
908        if let Record::ResetServer = record {
909            self.reset().await?;
910            Ok(Outcome::Success)
911        } else if let Record::Replace {
912            pattern,
913            replacement,
914        } = record
915        {
916            // Validated at parse time, so this compile cannot fail.
917            let regex = Regex::new(pattern).expect("replace regex validated by parser");
918            self.replacements.push((regex, replacement.clone()));
919            Ok(Outcome::Success)
920        } else {
921            self.inner
922                .as_mut()
923                .expect("RunnerInner missing")
924                .run_record(record, in_transaction, &self.replacements)
925                .await
926        }
927    }
928
929    async fn check_catalog(&self) -> Result<(), anyhow::Error> {
930        self.inner
931            .as_ref()
932            .expect("RunnerInner missing")
933            .check_catalog()
934            .await
935    }
936
937    #[allow(clippy::disallowed_methods)]
938    async fn reset_database(&mut self) -> Result<(), anyhow::Error> {
939        let inner = self.inner.as_mut().expect("RunnerInner missing");
940
941        inner.client.batch_execute("ROLLBACK;").await?;
942
943        inner
944            .system_client
945            .batch_execute(
946                "ROLLBACK;
947                 SET cluster = mz_catalog_server;
948                 RESET cluster_replica;",
949            )
950            .await?;
951
952        inner
953            .system_client
954            .batch_execute("ALTER SYSTEM RESET ALL")
955            .await?;
956
957        // Drop all databases, then recreate the `materialize` database.
958        for row in inner
959            .system_client
960            .query("SELECT name FROM mz_databases", &[])
961            .await?
962        {
963            let name: &str = row.get("name");
964            inner
965                .system_client
966                .batch_execute(sql!("DROP DATABASE {}", Sql::ident(name)).as_str())
967                .await?;
968        }
969
970        inner
971            .system_client
972            .batch_execute("CREATE DATABASE materialize")
973            .await?;
974
975        // Ensure quickstart cluster exists with one replica of size `self.config.replica_size`.
976        // We don't destroy the existing quickstart cluster replica if it exists, as turning
977        // on a cluster replica is exceptionally slow.
978        let mut needs_default_cluster = true;
979        for row in inner
980            .system_client
981            .query("SELECT name FROM mz_clusters WHERE id LIKE 'u%'", &[])
982            .await?
983        {
984            match row.get("name") {
985                "quickstart" => needs_default_cluster = false,
986                name => {
987                    inner
988                        .system_client
989                        .batch_execute(sql!("DROP CLUSTER {}", Sql::ident(name)).as_str())
990                        .await?
991                }
992            }
993        }
994        if needs_default_cluster {
995            inner
996                .system_client
997                .batch_execute("CREATE CLUSTER quickstart REPLICAS ()")
998                .await?;
999        }
1000        let mut needs_default_replica = false;
1001        let rows = inner
1002            .system_client
1003            .query(
1004                "SELECT name, size FROM mz_cluster_replicas
1005                 WHERE cluster_id = (SELECT id FROM mz_clusters WHERE name = 'quickstart')
1006                 ORDER BY name",
1007                &[],
1008            )
1009            .await?;
1010        if rows.len() != self.config.replicas {
1011            needs_default_replica = true;
1012        } else {
1013            for (i, row) in rows.iter().enumerate() {
1014                let name: &str = row.get("name");
1015                let size: &str = row.get("size");
1016                if name != format!("r{}", i + 1) || size != self.config.replica_size {
1017                    needs_default_replica = true;
1018                    break;
1019                }
1020            }
1021        }
1022
1023        if needs_default_replica {
1024            inner
1025                .system_client
1026                .batch_execute("ALTER CLUSTER quickstart SET (MANAGED = false)")
1027                .await?;
1028            for row in inner
1029                .system_client
1030                .query(
1031                    "SELECT name FROM mz_cluster_replicas
1032                     WHERE cluster_id = (SELECT id FROM mz_clusters WHERE name = 'quickstart')",
1033                    &[],
1034                )
1035                .await?
1036            {
1037                let name: &str = row.get("name");
1038                inner
1039                    .system_client
1040                    .batch_execute(
1041                        sql!("DROP CLUSTER REPLICA quickstart.{}", Sql::ident(name)).as_str(),
1042                    )
1043                    .await?;
1044            }
1045            for i in 1..=self.config.replicas {
1046                inner
1047                    .system_client
1048                    .batch_execute(
1049                        sql!(
1050                            "CREATE CLUSTER REPLICA quickstart.r{} SIZE {}",
1051                            i,
1052                            Sql::literal(&self.config.replica_size)
1053                        )
1054                        .as_str(),
1055                    )
1056                    .await?;
1057            }
1058            inner
1059                .system_client
1060                .batch_execute("ALTER CLUSTER quickstart SET (MANAGED = true)")
1061                .await?;
1062        }
1063
1064        // Drop all user roles except the default one, so that roles created
1065        // via the `user` directive or by the tests themselves do not leak
1066        // into later files. Best-effort: a role that still owns objects (for
1067        // example through default privileges) stays behind rather than
1068        // failing the reset.
1069        for row in inner
1070            .system_client
1071            .query(
1072                "SELECT name FROM mz_roles WHERE id LIKE 'u%' AND name != 'materialize'",
1073                &[],
1074            )
1075            .await?
1076        {
1077            let name: &str = row.get("name");
1078            let _ = inner
1079                .system_client
1080                .batch_execute(sql!("DROP ROLE {}", Sql::ident(name)).as_str())
1081                .await;
1082        }
1083
1084        // Grant initial privileges.
1085        inner
1086            .system_client
1087            .batch_execute("GRANT USAGE ON DATABASE materialize TO PUBLIC")
1088            .await?;
1089        inner
1090            .system_client
1091            .batch_execute("GRANT CREATE ON DATABASE materialize TO materialize")
1092            .await?;
1093        inner
1094            .system_client
1095            .batch_execute("GRANT CREATE ON SCHEMA materialize.public TO materialize")
1096            .await?;
1097        inner
1098            .system_client
1099            .batch_execute("GRANT USAGE ON CLUSTER quickstart TO PUBLIC")
1100            .await?;
1101        inner
1102            .system_client
1103            .batch_execute("GRANT CREATE ON CLUSTER quickstart TO materialize")
1104            .await?;
1105
1106        // Some sqllogictests require more than the default amount of tables, so we increase the
1107        // limit for all tests.
1108        inner
1109            .system_client
1110            .simple_query("ALTER SYSTEM SET max_tables = 100")
1111            .await?;
1112
1113        if inner.enable_table_keys {
1114            inner
1115                .system_client
1116                .simple_query("ALTER SYSTEM SET unsafe_enable_table_keys = true")
1117                .await?;
1118        }
1119
1120        inner.ensure_fixed_features().await?;
1121
1122        inner.client = connect(inner.server_addr, None, None).await.unwrap();
1123        inner.system_client = connect(inner.internal_server_addr, Some("mz_system"), None)
1124            .await
1125            .unwrap();
1126        inner.clients = BTreeMap::new();
1127        inner.active_user = None;
1128
1129        Ok(())
1130    }
1131}
1132
1133impl<'a> RunnerInner<'a> {
1134    pub async fn start(config: &RunConfig<'a>) -> Result<RunnerInner<'a>, anyhow::Error> {
1135        let temp_dir = tempfile::tempdir()?;
1136        let scratch_dir = tempfile::tempdir()?;
1137        let environment_id = EnvironmentId::for_tests();
1138        let (consensus_uri, timestamp_oracle_url): (SensitiveUrl, SensitiveUrl) = {
1139            let postgres_url = &config.postgres_url;
1140            let prefix = &config.prefix;
1141            info!(%postgres_url, "starting server");
1142            let (client, conn) = Retry::default()
1143                .max_tries(5)
1144                .retry_async(|_| async {
1145                    match tokio_postgres::connect(postgres_url, NoTls).await {
1146                        Ok(c) => Ok(c),
1147                        Err(e) => {
1148                            error!(%e, "failed to connect to postgres");
1149                            Err(e)
1150                        }
1151                    }
1152                })
1153                .await?;
1154            task::spawn(|| "sqllogictest_connect", async move {
1155                if let Err(e) = conn.await {
1156                    panic!("connection error: {}", e);
1157                }
1158            });
1159            // `prefix` is generated by sqllogictest harness configuration and
1160            // cannot be represented with composable `Sql`.
1161            #[allow(clippy::disallowed_methods)]
1162            client
1163                .batch_execute(&format!(
1164                    "DROP SCHEMA IF EXISTS {prefix}_tsoracle CASCADE;
1165                     CREATE SCHEMA IF NOT EXISTS {prefix}_consensus;
1166                     CREATE SCHEMA {prefix}_tsoracle;"
1167                ))
1168                .await?;
1169            (
1170                format!("{postgres_url}?options=--search_path={prefix}_consensus")
1171                    .parse()
1172                    .expect("invalid consensus URI"),
1173                format!("{postgres_url}?options=--search_path={prefix}_tsoracle")
1174                    .parse()
1175                    .expect("invalid timestamp oracle URI"),
1176            )
1177        };
1178
1179        let secrets_dir = temp_dir.path().join("secrets");
1180        let orchestrator = Arc::new(
1181            ProcessOrchestrator::new(ProcessOrchestratorConfig {
1182                image_dir: env::current_exe()?.parent().unwrap().to_path_buf(),
1183                suppress_output: false,
1184                environment_id: environment_id.to_string(),
1185                secrets_dir: secrets_dir.clone(),
1186                command_wrapper: config
1187                    .orchestrator_process_wrapper
1188                    .as_ref()
1189                    .map_or(Ok(vec![]), |s| shell_words::split(s))?,
1190                propagate_crashes: true,
1191                tcp_proxy: None,
1192                scratch_directory: scratch_dir.path().to_path_buf(),
1193            })
1194            .await?,
1195        );
1196        let now = SYSTEM_TIME.clone();
1197        let metrics_registry = MetricsRegistry::new();
1198
1199        let persist_config = PersistConfig::new(
1200            &mz_environmentd::BUILD_INFO,
1201            now.clone(),
1202            mz_dyncfgs::all_dyncfgs(),
1203        );
1204        let persist_pubsub_server =
1205            PersistGrpcPubSubServer::new(&persist_config, &metrics_registry);
1206        let persist_pubsub_client = persist_pubsub_server.new_same_process_connection();
1207        let persist_pubsub_tcp_listener =
1208            TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0))
1209                .await
1210                .expect("pubsub addr binding");
1211        let persist_pubsub_server_port = persist_pubsub_tcp_listener
1212            .local_addr()
1213            .expect("pubsub addr has local addr")
1214            .port();
1215        info!("listening for persist pubsub connections on localhost:{persist_pubsub_server_port}");
1216        mz_ore::task::spawn(|| "persist_pubsub_server", async move {
1217            persist_pubsub_server
1218                .serve_with_stream(TcpListenerStream::new(persist_pubsub_tcp_listener))
1219                .await
1220                .expect("success")
1221        });
1222        let persist_clients =
1223            PersistClientCache::new(persist_config, &metrics_registry, |cfg, metrics| {
1224                let sender: Arc<dyn PubSubSender> = Arc::new(MetricsSameProcessPubSubSender::new(
1225                    cfg,
1226                    persist_pubsub_client.sender,
1227                    metrics,
1228                ));
1229                PubSubClientConnection::new(sender, persist_pubsub_client.receiver)
1230            });
1231        let persist_clients = Arc::new(persist_clients);
1232
1233        let secrets_controller = Arc::clone(&orchestrator);
1234        let connection_context = ConnectionContext::for_tests(orchestrator.reader());
1235        let orchestrator = Arc::new(TracingOrchestrator::new(
1236            orchestrator,
1237            config.tracing.clone(),
1238        ));
1239        let listeners_config = ListenersConfig {
1240            sql: btreemap! {
1241                "external".to_owned() => SqlListenerConfig {
1242                    addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
1243                    authenticator_kind: AuthenticatorKind::None,
1244                    allowed_roles: AllowedRoles::Normal,
1245                    enable_tls: false,
1246                },
1247                "internal".to_owned() => SqlListenerConfig {
1248                    addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
1249                    authenticator_kind: AuthenticatorKind::None,
1250                    allowed_roles: AllowedRoles::Internal,
1251                    enable_tls: false,
1252                },
1253                "password".to_owned() => SqlListenerConfig {
1254                    addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
1255                    authenticator_kind: AuthenticatorKind::Password,
1256                    allowed_roles: AllowedRoles::Normal,
1257                    enable_tls: false,
1258                },
1259            },
1260            http: btreemap![
1261                "external".to_owned() => HttpListenerConfig {
1262                    addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
1263                    authenticator_kind: AuthenticatorKind::None,
1264                    enable_tls: false,
1265                    routes: HttpRoutesEnabled {
1266                        base: RouteGroup::Enabled(AllowedRoles::Normal),
1267                        webhook: RouteGroup::Enabled(AllowedRoles::Normal),
1268                        internal: RouteGroup::Disabled,
1269                        metrics: RouteGroup::Disabled,
1270                        profiling: RouteGroup::Disabled,
1271                        mcp_agent: RouteGroup::Disabled,
1272                        mcp_developer: RouteGroup::Disabled,
1273                        mcp: RouteGroup::Disabled,
1274                        console_config: RouteGroup::Enabled(AllowedRoles::Normal),
1275                    },
1276                },
1277                "internal".to_owned() => HttpListenerConfig {
1278                    addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
1279                    authenticator_kind: AuthenticatorKind::None,
1280                    enable_tls: false,
1281                    routes: HttpRoutesEnabled {
1282                        base: RouteGroup::Enabled(AllowedRoles::NormalAndInternal),
1283                        webhook: RouteGroup::Enabled(AllowedRoles::NormalAndInternal),
1284                        internal: RouteGroup::Enabled(AllowedRoles::NormalAndInternal),
1285                        metrics: RouteGroup::Enabled(AllowedRoles::NormalAndInternal),
1286                        profiling: RouteGroup::Enabled(AllowedRoles::NormalAndInternal),
1287                        mcp_agent: RouteGroup::Disabled,
1288                        mcp_developer: RouteGroup::Disabled,
1289                        mcp: RouteGroup::Disabled,
1290                        console_config: RouteGroup::Disabled,
1291                    },
1292                },
1293            ],
1294        };
1295        let listeners = mz_environmentd::Listeners::bind(listeners_config)?;
1296        let host_name = format!(
1297            "localhost:{}",
1298            listeners.http["external"].handle.local_addr.port()
1299        );
1300        let catalog_config = CatalogConfig {
1301            persist_clients: Arc::clone(&persist_clients),
1302            metrics: Arc::new(mz_catalog::durable::Metrics::new(&MetricsRegistry::new())),
1303        };
1304        let system_dyncfgs = Arc::clone(&persist_clients.cfg().configs);
1305        let server_config = mz_environmentd::Config {
1306            catalog_config,
1307            timestamp_oracle_url: Some(timestamp_oracle_url),
1308            controller: ControllerConfig {
1309                build_info: &mz_environmentd::BUILD_INFO,
1310                orchestrator,
1311                clusterd_image: "clusterd".into(),
1312                init_container_image: None,
1313                deploy_generation: 0,
1314                persist_location: PersistLocation {
1315                    blob_uri: format!(
1316                        "file://{}/persist/blob",
1317                        config.persist_dir.path().display()
1318                    )
1319                    .parse()
1320                    .expect("invalid blob URI"),
1321                    consensus_uri,
1322                },
1323                persist_clients,
1324                now: SYSTEM_TIME.clone(),
1325                metrics_registry: metrics_registry.clone(),
1326                persist_pubsub_url: format!("http://localhost:{}", persist_pubsub_server_port),
1327                secrets_args: mz_secrets_cli::SecretsReaderCliArgs {
1328                    secrets_reader: mz_secrets_cli::SecretsControllerKind::LocalFile,
1329                    secrets_reader_local_file_dir: Some(secrets_dir),
1330                    secrets_reader_kubernetes_context: None,
1331                    secrets_reader_aws_prefix: None,
1332                    secrets_reader_name_prefix: None,
1333                },
1334                connection_context,
1335                replica_http_locator: Arc::new(ReplicaHttpLocator::default()),
1336            },
1337            secrets_controller,
1338            cloud_resource_controller: None,
1339            system_dyncfgs,
1340            tls: None,
1341            frontegg: None,
1342            frontegg_oauth_issuer_url: None,
1343            cors_allowed_origin: AllowOrigin::list([]),
1344            cors_allowed_origin_list: Vec::new(),
1345            unsafe_mode: true,
1346            all_features: false,
1347            restarted_in_place: false,
1348            metrics_registry,
1349            now,
1350            environment_id,
1351            cluster_replica_sizes: ClusterReplicaSizeMap::for_tests(),
1352            bootstrap_default_cluster_replica_size: config.replica_size.clone(),
1353            bootstrap_default_cluster_replication_factor: config
1354                .replicas
1355                .try_into()
1356                .expect("replicas must fit"),
1357            bootstrap_builtin_system_cluster_config: BootstrapBuiltinClusterConfig {
1358                replication_factor: SYSTEM_CLUSTER_DEFAULT_REPLICATION_FACTOR,
1359                size: config.replica_size.clone(),
1360            },
1361            bootstrap_builtin_catalog_server_cluster_config: BootstrapBuiltinClusterConfig {
1362                replication_factor: CATALOG_SERVER_CLUSTER_DEFAULT_REPLICATION_FACTOR,
1363                size: config.replica_size.clone(),
1364            },
1365            bootstrap_builtin_probe_cluster_config: BootstrapBuiltinClusterConfig {
1366                replication_factor: PROBE_CLUSTER_DEFAULT_REPLICATION_FACTOR,
1367                size: config.replica_size.clone(),
1368            },
1369            bootstrap_builtin_support_cluster_config: BootstrapBuiltinClusterConfig {
1370                replication_factor: SUPPORT_CLUSTER_DEFAULT_REPLICATION_FACTOR,
1371                size: config.replica_size.clone(),
1372            },
1373            bootstrap_builtin_analytics_cluster_config: BootstrapBuiltinClusterConfig {
1374                replication_factor: ANALYTICS_CLUSTER_DEFAULT_REPLICATION_FACTOR,
1375                size: config.replica_size.clone(),
1376            },
1377            system_parameter_defaults: {
1378                let mut params = BTreeMap::new();
1379                params.insert(
1380                    "log_filter".to_string(),
1381                    config.tracing.startup_log_filter.to_string(),
1382                );
1383                params.extend(config.system_parameter_defaults.clone());
1384                params
1385            },
1386            availability_zones: Default::default(),
1387            tracing_handle: config.tracing_handle.clone(),
1388            storage_usage_collection_interval: Duration::from_secs(3600),
1389            storage_usage_retention_period: None,
1390            segment_api_key: None,
1391            segment_client_side: false,
1392            // SLT doesn't like eternally running tasks since it waits for them to finish inbetween SLT files
1393            test_only_dummy_segment_client: false,
1394            egress_addresses: vec![],
1395            aws_account_id: None,
1396            aws_privatelink_availability_zones: None,
1397            launchdarkly_sdk_key: None,
1398            launchdarkly_base_uri: None,
1399            launchdarkly_key_map: Default::default(),
1400            config_sync_file_path: None,
1401            config_sync_timeout: Duration::from_secs(30),
1402            config_sync_loop_interval: None,
1403            bootstrap_role: Some("materialize".into()),
1404            http_host_name: Some(host_name),
1405            internal_console_redirect_url: None,
1406            tls_reload_certs: mz_server_core::cert_reload_never_reload(),
1407            helm_chart_version: None,
1408            license_key: ValidatedLicenseKey::for_tests(),
1409            external_login_password_mz_system: None,
1410            force_builtin_schema_migration: None,
1411        };
1412        // We need to run the server on its own Tokio runtime, which in turn
1413        // requires its own thread, so that we can wait for any tasks spawned
1414        // by the server to be shutdown at the end of each file. If we were to
1415        // share a Tokio runtime, tasks from the last file's server would still
1416        // be live at the start of the next file's server.
1417        let (server_addr_tx, server_addr_rx): (oneshot::Sender<Result<_, anyhow::Error>>, _) =
1418            oneshot::channel();
1419        let (internal_server_addr_tx, internal_server_addr_rx) = oneshot::channel();
1420        let (password_server_addr_tx, password_server_addr_rx) = oneshot::channel();
1421        let (internal_http_server_addr_tx, internal_http_server_addr_rx) = oneshot::channel();
1422        let (shutdown_trigger, shutdown_trigger_rx) = trigger::channel();
1423        let server_thread = thread::spawn(|| {
1424            let runtime = match Runtime::new() {
1425                Ok(runtime) => runtime,
1426                Err(e) => {
1427                    server_addr_tx
1428                        .send(Err(e.into()))
1429                        .expect("receiver should not drop first");
1430                    return;
1431                }
1432            };
1433            let server = match runtime.block_on(listeners.serve(server_config)) {
1434                Ok(runtime) => runtime,
1435                Err(e) => {
1436                    server_addr_tx
1437                        .send(Err(e.into()))
1438                        .expect("receiver should not drop first");
1439                    return;
1440                }
1441            };
1442            server_addr_tx
1443                .send(Ok(server.sql_listener_handles["external"].local_addr))
1444                .expect("receiver should not drop first");
1445            internal_server_addr_tx
1446                .send(server.sql_listener_handles["internal"].local_addr)
1447                .expect("receiver should not drop first");
1448            password_server_addr_tx
1449                .send(server.sql_listener_handles["password"].local_addr)
1450                .expect("receiver should not drop first");
1451            internal_http_server_addr_tx
1452                .send(server.http_listener_handles["internal"].local_addr)
1453                .expect("receiver should not drop first");
1454            runtime.block_on(shutdown_trigger_rx);
1455        });
1456        let server_addr = server_addr_rx.await??;
1457        let internal_server_addr = internal_server_addr_rx.await?;
1458        let password_server_addr = password_server_addr_rx.await?;
1459        let internal_http_server_addr = internal_http_server_addr_rx.await?;
1460
1461        let system_client = connect(internal_server_addr, Some("mz_system"), None)
1462            .await
1463            .unwrap();
1464        let client = connect(server_addr, None, None).await.unwrap();
1465
1466        let inner = RunnerInner {
1467            server_addr,
1468            internal_server_addr,
1469            password_server_addr,
1470            internal_http_server_addr,
1471            _shutdown_trigger: shutdown_trigger,
1472            _server_thread: server_thread.join_on_drop(),
1473            _temp_dir: temp_dir,
1474            client,
1475            system_client,
1476            clients: BTreeMap::new(),
1477            active_user: None,
1478            auto_index_tables: config.auto_index_tables,
1479            auto_index_selects: config.auto_index_selects,
1480            auto_transactions: config.auto_transactions,
1481            enable_table_keys: config.enable_table_keys,
1482            verbose: config.verbose,
1483            stdout: config.stdout,
1484        };
1485        inner.ensure_fixed_features().await?;
1486
1487        Ok(inner)
1488    }
1489
1490    /// Set features that should be enabled regardless of whether reset-server was
1491    /// called. These features may be set conditionally depending on the run configuration.
1492    #[allow(clippy::disallowed_methods)]
1493    async fn ensure_fixed_features(&self) -> Result<(), anyhow::Error> {
1494        // We turn on enable_reduce_mfp_fusion, as we wish
1495        // to get as much coverage of these features as we can.
1496        // TODO(vmarcos): Remove this code when we retire this feature flag.
1497        self.system_client
1498            .execute("ALTER SYSTEM SET enable_reduce_mfp_fusion = on", &[])
1499            .await?;
1500
1501        // Dangerous functions are useful for tests so we enable it for all tests.
1502        self.system_client
1503            .execute("ALTER SYSTEM SET unsafe_enable_unsafe_functions = on", &[])
1504            .await?;
1505        Ok(())
1506    }
1507
1508    /// The connection records run on: the `user` directive's connection when
1509    /// one is active, the default connection otherwise.
1510    fn active_client(&self) -> &tokio_postgres::Client {
1511        match &self.active_user {
1512            Some(user) => self
1513                .clients
1514                .get(user)
1515                .expect("connection for the active user exists"),
1516            None => &self.client,
1517        }
1518    }
1519
1520    /// Handles a `user` directive: ensures the role exists, connects as it,
1521    /// and makes it the active user. CockroachDB's `root` maps to the
1522    /// default `materialize` user.
1523    #[allow(clippy::disallowed_methods)]
1524    async fn run_user<'r>(
1525        &mut self,
1526        user: &'r str,
1527        location: &Location,
1528        in_transaction: &mut bool,
1529    ) -> Result<Outcome<'r>, anyhow::Error> {
1530        if self.auto_transactions && *in_transaction {
1531            self.active_client().execute("COMMIT", &[]).await?;
1532            *in_transaction = false;
1533        }
1534        if user == "root" || user == "materialize" {
1535            self.active_user = None;
1536            return Ok(Outcome::Success);
1537        }
1538        if !self.clients.contains_key(user) {
1539            let create = sql!("CREATE ROLE {}", Sql::ident(user));
1540            if let Err(error) = self.system_client.batch_execute(create.as_str()).await {
1541                if !error.to_string_with_causes().contains("already exists") {
1542                    return Ok(Outcome::Bail {
1543                        cause: Box::new(Outcome::PlanFailure {
1544                            error: anyhow!(error),
1545                            expected_error: None,
1546                            location: location.clone(),
1547                        }),
1548                        location: location.clone(),
1549                    });
1550                }
1551            }
1552            match connect(self.server_addr, Some(user), None).await {
1553                Ok(client) => {
1554                    self.clients.insert(user.to_string(), client);
1555                }
1556                Err(error) => {
1557                    return Ok(Outcome::Bail {
1558                        cause: Box::new(Outcome::PlanFailure {
1559                            error: anyhow!(error),
1560                            expected_error: None,
1561                            location: location.clone(),
1562                        }),
1563                        location: location.clone(),
1564                    });
1565                }
1566            }
1567        }
1568        self.active_user = Some(user.to_string());
1569        Ok(Outcome::Success)
1570    }
1571
1572    #[allow(clippy::disallowed_methods)]
1573    async fn run_record<'r>(
1574        &mut self,
1575        record: &'r Record<'r>,
1576        in_transaction: &mut bool,
1577        replacements: &[(Regex, String)],
1578    ) -> Result<Outcome<'r>, anyhow::Error> {
1579        match &record {
1580            Record::User { user, location } => self.run_user(user, location, in_transaction).await,
1581            Record::Statement {
1582                expected_error,
1583                rows_affected,
1584                sql,
1585                location,
1586            } => {
1587                if self.auto_transactions && *in_transaction {
1588                    self.active_client().execute("COMMIT", &[]).await?;
1589                    *in_transaction = false;
1590                }
1591                match self
1592                    .run_statement(*expected_error, *rows_affected, sql, location.clone())
1593                    .await?
1594                {
1595                    Outcome::Success => {
1596                        if self.auto_index_tables {
1597                            let additional = mutate(sql);
1598                            for stmt in additional {
1599                                self.active_client().execute(&stmt, &[]).await?;
1600                            }
1601                        }
1602                        Ok(Outcome::Success)
1603                    }
1604                    other => {
1605                        if expected_error.is_some() {
1606                            Ok(other)
1607                        } else {
1608                            // If we failed to execute a statement that was supposed to succeed,
1609                            // running the rest of the tests in this file will probably cause
1610                            // false positives, so just give up on the file entirely.
1611                            Ok(Outcome::Bail {
1612                                cause: Box::new(other),
1613                                location: location.clone(),
1614                            })
1615                        }
1616                    }
1617                }
1618            }
1619            Record::Query {
1620                sql,
1621                output,
1622                location,
1623            } => {
1624                self.run_query(sql, output, location.clone(), in_transaction, replacements)
1625                    .await
1626            }
1627            Record::Simple {
1628                conn,
1629                user,
1630                password,
1631                sql,
1632                sort,
1633                output,
1634                location,
1635                ..
1636            } => {
1637                self.run_simple(
1638                    *conn,
1639                    *user,
1640                    *password,
1641                    sql,
1642                    sort.clone(),
1643                    output,
1644                    location.clone(),
1645                )
1646                .await
1647            }
1648            Record::Copy {
1649                table_name,
1650                tsv_path,
1651            } => {
1652                let tsv = tokio::fs::read(tsv_path).await?;
1653                let copy = self
1654                    .client
1655                    .copy_in(sql!("COPY {} FROM STDIN", Sql::ident(*table_name)).as_str())
1656                    .await?;
1657                tokio::pin!(copy);
1658                copy.send(bytes::Bytes::from(tsv)).await?;
1659                copy.finish().await?;
1660                Ok(Outcome::Success)
1661            }
1662            _ => Ok(Outcome::Success),
1663        }
1664    }
1665
1666    #[allow(clippy::disallowed_methods)]
1667    async fn run_statement<'r>(
1668        &self,
1669        expected_error: Option<&'r str>,
1670        expected_rows_affected: Option<u64>,
1671        sql: &'r str,
1672        location: Location,
1673    ) -> Result<Outcome<'r>, anyhow::Error> {
1674        static UNSUPPORTED_INDEX_STATEMENT_REGEX: LazyLock<Regex> =
1675            LazyLock::new(|| Regex::new("^(CREATE UNIQUE INDEX|REINDEX)").unwrap());
1676        if UNSUPPORTED_INDEX_STATEMENT_REGEX.is_match(sql) {
1677            // sure, we totally made you an index
1678            return Ok(Outcome::Success);
1679        }
1680
1681        // CockroachDB's logic tests frequently pack multiple commands into
1682        // one statement record. The extended protocol used by `execute`
1683        // rejects those, and Materialize restricts DDL inside the implicit
1684        // transaction the simple protocol would wrap them in, so run each
1685        // command separately and report the first error. This path cannot
1686        // report an affected-row count, so records that assert one keep the
1687        // single-command path (and fail if they contain multiple commands).
1688        let commands = if expected_rows_affected.is_none() {
1689            match mz_sql::parse::parse(sql) {
1690                Ok(stmts) if stmts.len() > 1 => {
1691                    Some(stmts.iter().map(|s| s.sql.to_owned()).collect::<Vec<_>>())
1692                }
1693                Ok(_) => None,
1694                // A statement Materialize cannot parse may be a CockroachDB
1695                // `CREATE TABLE` with inline physical-layout items. Retry
1696                // with those stripped, so the semantic records that follow
1697                // can still run against the table.
1698                Err(_) => strip_crdb_table_items(sql)
1699                    .filter(|stripped| mz_sql::parse::parse(stripped).is_ok())
1700                    .map(|stripped| vec![stripped]),
1701            }
1702        } else {
1703            None
1704        };
1705        let result = match commands {
1706            Some(commands) => {
1707                let mut result = Ok(0);
1708                for command in &commands {
1709                    if let Err(error) = self.active_client().execute(command, &[]).await {
1710                        result = Err(error);
1711                        break;
1712                    }
1713                }
1714                result
1715            }
1716            None => self.active_client().execute(sql, &[]).await,
1717        };
1718        match result {
1719            Ok(actual) => {
1720                if let Some(expected_error) = expected_error {
1721                    return Ok(Outcome::UnexpectedPlanSuccess {
1722                        expected_error,
1723                        location,
1724                    });
1725                }
1726                match expected_rows_affected {
1727                    None => Ok(Outcome::Success),
1728                    Some(expected) => {
1729                        if expected != actual {
1730                            Ok(Outcome::WrongNumberOfRowsInserted {
1731                                expected_count: expected,
1732                                actual_count: actual,
1733                                location,
1734                            })
1735                        } else {
1736                            Ok(Outcome::Success)
1737                        }
1738                    }
1739                }
1740            }
1741            Err(error) => {
1742                if let Some(expected_error) = expected_error {
1743                    if error_matches(expected_error, &error.to_string_with_causes()) {
1744                        return Ok(Outcome::Success);
1745                    }
1746                    return Ok(Outcome::PlanFailure {
1747                        error: anyhow!(error),
1748                        expected_error: Some(expected_error.to_string()),
1749                        location,
1750                    });
1751                }
1752                // CockroachDB tests configure many CockroachDB-specific
1753                // session settings. Treat setting an unknown parameter as a
1754                // no-op success so the records that follow still run.
1755                static SET_STATEMENT_REGEX: LazyLock<Regex> =
1756                    LazyLock::new(|| Regex::new("(?i)^(SET|RESET) ").unwrap());
1757                if SET_STATEMENT_REGEX.is_match(sql.trim_start())
1758                    && error
1759                        .to_string_with_causes()
1760                        .contains("unrecognized configuration parameter")
1761                {
1762                    return Ok(Outcome::Success);
1763                }
1764                Ok(Outcome::PlanFailure {
1765                    error: anyhow!(error),
1766                    expected_error: None,
1767                    location,
1768                })
1769            }
1770        }
1771    }
1772
1773    #[allow(clippy::disallowed_methods)]
1774    async fn prepare_query<'r>(
1775        &self,
1776        sql: &str,
1777        output: &'r Result<QueryOutput<'_>, &'r str>,
1778        location: Location,
1779        in_transaction: &mut bool,
1780    ) -> Result<PrepareQueryOutcome<'r>, anyhow::Error> {
1781        // get statement
1782        let statements = match mz_sql::parse::parse(sql) {
1783            Ok(statements) => statements,
1784            Err(e) => match output {
1785                Ok(_) => {
1786                    return Ok(PrepareQueryOutcome::Outcome(Outcome::ParseFailure {
1787                        error: e.into(),
1788                        location,
1789                    }));
1790                }
1791                Err(expected_error) => {
1792                    if error_matches(expected_error, &e.to_string_with_causes()) {
1793                        return Ok(PrepareQueryOutcome::Outcome(Outcome::Success));
1794                    } else {
1795                        return Ok(PrepareQueryOutcome::Outcome(Outcome::ParseFailure {
1796                            error: e.into(),
1797                            location,
1798                        }));
1799                    }
1800                }
1801            },
1802        };
1803        // Query records in files imported from CockroachDB occasionally
1804        // contain zero or multiple statements. Fail the record rather than
1805        // the entire run.
1806        let statement = match &*statements {
1807            [statement] => &statement.ast,
1808            _ => {
1809                return Ok(PrepareQueryOutcome::Outcome(Outcome::ParseFailure {
1810                    error: anyhow!("expected one statement, but got {}", statements.len()),
1811                    location,
1812                }));
1813            }
1814        };
1815        let (is_select, num_attributes, has_as_of) = match statement {
1816            Statement::Select(stmt) => (
1817                true,
1818                derive_num_attributes(&stmt.query.body),
1819                stmt.as_of.is_some(),
1820            ),
1821            _ => (false, None, false),
1822        };
1823
1824        match output {
1825            Ok(_) => {
1826                if self.auto_transactions && !*in_transaction {
1827                    // No ISOLATION LEVEL SERIALIZABLE because of database-issues#5323
1828                    self.active_client().execute("BEGIN", &[]).await?;
1829                    *in_transaction = true;
1830                }
1831            }
1832            Err(_) => {
1833                if self.auto_transactions && *in_transaction {
1834                    self.active_client().execute("COMMIT", &[]).await?;
1835                    *in_transaction = false;
1836                }
1837            }
1838        }
1839
1840        // `SHOW` commands reference catalog schema, thus are not in the same timedomain and not
1841        // allowed in the same transaction, see:
1842        // https://materialize.com/docs/sql/begin/#same-timedomain-error
1843        match statement {
1844            Statement::Show(..) => {
1845                if self.auto_transactions && *in_transaction {
1846                    self.active_client().execute("COMMIT", &[]).await?;
1847                    *in_transaction = false;
1848                }
1849            }
1850            _ => (),
1851        }
1852        Ok(PrepareQueryOutcome::QueryPrepared(QueryInfo {
1853            is_select,
1854            num_attributes,
1855            has_as_of,
1856        }))
1857    }
1858
1859    #[allow(clippy::disallowed_methods)]
1860    async fn execute_query<'r>(
1861        &self,
1862        sql: &str,
1863        output: &'r Result<QueryOutput<'_>, &'r str>,
1864        location: Location,
1865        replacements: &[(Regex, String)],
1866    ) -> Result<Outcome<'r>, anyhow::Error> {
1867        let rows = match self.active_client().query(sql, &[]).await {
1868            Ok(rows) => rows,
1869            Err(error) => {
1870                let error_string = error.to_string_with_causes();
1871                return match output {
1872                    Ok(_) => {
1873                        if error_string.contains("supported") || error_string.contains("overload") {
1874                            // this is a failure, but it's caused by lack of support rather than by bugs
1875                            Ok(Outcome::Unsupported {
1876                                error: anyhow!(error),
1877                                location,
1878                            })
1879                        } else {
1880                            Ok(Outcome::PlanFailure {
1881                                error: anyhow!(error),
1882                                expected_error: None,
1883                                location,
1884                            })
1885                        }
1886                    }
1887                    Err(expected_error) => {
1888                        if error_matches(expected_error, &error_string) {
1889                            Ok(Outcome::Success)
1890                        } else {
1891                            Ok(Outcome::PlanFailure {
1892                                error: anyhow!(error),
1893                                expected_error: Some(expected_error.to_string()),
1894                                location,
1895                            })
1896                        }
1897                    }
1898                };
1899            }
1900        };
1901
1902        // unpack expected output
1903        let QueryOutput {
1904            sort,
1905            types: expected_types,
1906            column_names: expected_column_names,
1907            output: expected_output,
1908            mode,
1909            ..
1910        } = match output {
1911            Err(expected_error) => {
1912                return Ok(Outcome::UnexpectedPlanSuccess {
1913                    expected_error,
1914                    location,
1915                });
1916            }
1917            Ok(query_output) => query_output,
1918        };
1919
1920        // format output
1921        let mut formatted_rows = vec![];
1922        for row in &rows {
1923            if row.len() != expected_types.len() {
1924                return Ok(Outcome::WrongColumnCount {
1925                    expected_count: expected_types.len(),
1926                    actual_count: row.len(),
1927                    location,
1928                });
1929            }
1930            let row = format_row(row, expected_types, *mode);
1931            formatted_rows.push(row);
1932        }
1933
1934        // sort formatted output
1935        if let Sort::Row = sort {
1936            formatted_rows.sort();
1937        }
1938        let mut values = formatted_rows.into_iter().flatten().collect::<Vec<_>>();
1939        if let Sort::Value = sort {
1940            values.sort();
1941        }
1942
1943        // Apply any `replace` substitutions to the actual output, so that
1944        // non-deterministic tokens (e.g. a folded `mz_now()` timestamp) are
1945        // masked before the output is compared against, or rewritten into, the
1946        // expected output.
1947        if !replacements.is_empty() {
1948            for value in &mut values {
1949                for (regex, replacement) in replacements {
1950                    *value = regex.replace_all(value, replacement.as_str()).into_owned();
1951                }
1952            }
1953        }
1954
1955        // Various checks as long as there are returned rows.
1956        if let Some(row) = rows.get(0) {
1957            // check column names
1958            if let Some(expected_column_names) = expected_column_names {
1959                let actual_column_names = row
1960                    .columns()
1961                    .iter()
1962                    .map(|t| ColumnName::from(t.name()))
1963                    .collect::<Vec<_>>();
1964                if expected_column_names != &actual_column_names {
1965                    return Ok(Outcome::WrongColumnNames {
1966                        expected_column_names,
1967                        actual_column_names,
1968                        actual_output: Output::Values(values),
1969                        location,
1970                    });
1971                }
1972            }
1973        }
1974
1975        // check output
1976        match expected_output {
1977            Output::Values(expected_values) => {
1978                if values != *expected_values {
1979                    return Ok(Outcome::OutputFailure {
1980                        expected_output,
1981                        actual_raw_output: rows,
1982                        actual_output: Output::Values(values),
1983                        location,
1984                    });
1985                }
1986            }
1987            Output::Hashed {
1988                num_values,
1989                md5: expected_md5,
1990            } => {
1991                let mut hasher = Md5::new();
1992                for value in &values {
1993                    hasher.update(value);
1994                    hasher.update("\n");
1995                }
1996                let md5 = hex::encode(hasher.finalize());
1997                if values.len() != *num_values || md5 != *expected_md5 {
1998                    return Ok(Outcome::OutputFailure {
1999                        expected_output,
2000                        actual_raw_output: rows,
2001                        actual_output: Output::Hashed {
2002                            num_values: values.len(),
2003                            md5,
2004                        },
2005                        location,
2006                    });
2007                }
2008            }
2009        }
2010
2011        Ok(Outcome::Success)
2012    }
2013
2014    #[allow(clippy::disallowed_methods)]
2015    async fn execute_view_inner<'r>(
2016        &self,
2017        sql: &str,
2018        output: &'r Result<QueryOutput<'_>, &'r str>,
2019        location: Location,
2020    ) -> Result<Option<Outcome<'r>>, anyhow::Error> {
2021        print_sql_if(self.stdout, sql, self.verbose);
2022        let sql_result = self.active_client().execute(sql, &[]).await;
2023
2024        // Evaluate if we already reached an outcome or not.
2025        let tentative_outcome = if let Err(view_error) = sql_result {
2026            if let Err(expected_error) = output {
2027                if error_matches(expected_error, &view_error.to_string_with_causes()) {
2028                    Some(Outcome::Success)
2029                } else {
2030                    Some(Outcome::PlanFailure {
2031                        error: view_error.into(),
2032                        expected_error: Some(expected_error.to_string()),
2033                        location: location.clone(),
2034                    })
2035                }
2036            } else {
2037                Some(Outcome::PlanFailure {
2038                    error: view_error.into(),
2039                    expected_error: None,
2040                    location: location.clone(),
2041                })
2042            }
2043        } else {
2044            None
2045        };
2046        Ok(tentative_outcome)
2047    }
2048
2049    #[allow(clippy::disallowed_methods)]
2050    async fn execute_view<'r>(
2051        &self,
2052        sql: &str,
2053        num_attributes: Option<usize>,
2054        output: &'r Result<QueryOutput<'_>, &'r str>,
2055        location: Location,
2056        replacements: &[(Regex, String)],
2057    ) -> Result<Outcome<'r>, anyhow::Error> {
2058        // Create indexed view SQL commands and execute `CREATE VIEW`.
2059        let expected_column_names = if let Ok(QueryOutput { column_names, .. }) = output {
2060            column_names.clone()
2061        } else {
2062            None
2063        };
2064        let (create_view, create_index, view_sql, drop_view) = generate_view_sql(
2065            sql,
2066            Uuid::new_v4().as_simple(),
2067            num_attributes,
2068            expected_column_names,
2069        );
2070        let tentative_outcome = self
2071            .execute_view_inner(create_view.as_str(), output, location.clone())
2072            .await?;
2073
2074        // Either we already have an outcome or alternatively,
2075        // we proceed to index and query the view.
2076        if let Some(view_outcome) = tentative_outcome {
2077            return Ok(view_outcome);
2078        }
2079
2080        let tentative_outcome = self
2081            .execute_view_inner(create_index.as_str(), output, location.clone())
2082            .await?;
2083
2084        let view_outcome;
2085        if let Some(outcome) = tentative_outcome {
2086            view_outcome = outcome;
2087        } else {
2088            print_sql_if(self.stdout, view_sql.as_str(), self.verbose);
2089            view_outcome = self
2090                .execute_query(view_sql.as_str(), output, location.clone(), replacements)
2091                .await?;
2092        }
2093
2094        // Remember to clean up after ourselves by dropping the view.
2095        print_sql_if(self.stdout, drop_view.as_str(), self.verbose);
2096        self.active_client()
2097            .execute(drop_view.as_str(), &[])
2098            .await?;
2099
2100        Ok(view_outcome)
2101    }
2102
2103    async fn run_query<'r>(
2104        &self,
2105        sql: &'r str,
2106        output: &'r Result<QueryOutput<'_>, &'r str>,
2107        location: Location,
2108        in_transaction: &mut bool,
2109        replacements: &[(Regex, String)],
2110    ) -> Result<Outcome<'r>, anyhow::Error> {
2111        let prepare_outcome = self
2112            .prepare_query(sql, output, location.clone(), in_transaction)
2113            .await?;
2114        match prepare_outcome {
2115            PrepareQueryOutcome::QueryPrepared(QueryInfo {
2116                is_select,
2117                num_attributes,
2118                has_as_of,
2119            }) => {
2120                let query_outcome = self
2121                    .execute_query(sql, output, location.clone(), replacements)
2122                    .await?;
2123                // `AS OF` queries are excluded from the indexed-view consistency
2124                // check: re-running them against a freshly created indexed view
2125                // is racy, since the view's `since` can advance past a historical
2126                // `AS OF` and spuriously fail with "could not find a valid
2127                // timestamp" even though the one-shot query succeeded.
2128                if is_select && !has_as_of && self.auto_index_selects && query_outcome.success() {
2129                    let view_outcome = self
2130                        .execute_view(sql, None, output, location.clone(), replacements)
2131                        .await?;
2132
2133                    if !view_outcome.success() {
2134                        // Before producing a failure outcome, we try to obtain a new
2135                        // outcome for view-based execution exploiting analysis of the
2136                        // number of attributes. This two-level strategy can avoid errors
2137                        // produced by column ambiguity in the `SELECT`.
2138                        let view_outcome = if num_attributes.is_some() {
2139                            self.execute_view(
2140                                sql,
2141                                num_attributes,
2142                                output,
2143                                location.clone(),
2144                                replacements,
2145                            )
2146                            .await?
2147                        } else {
2148                            view_outcome
2149                        };
2150
2151                        if !view_outcome.success() {
2152                            let inconsistent_view_outcome = Outcome::InconsistentViewOutcome {
2153                                query_outcome: Box::new(query_outcome),
2154                                view_outcome: Box::new(view_outcome),
2155                                location: location.clone(),
2156                            };
2157                            // Determine if this inconsistent view outcome should be reported
2158                            // as an error or only as a warning.
2159                            let outcome = if should_warn(&inconsistent_view_outcome) {
2160                                Outcome::Warning {
2161                                    cause: Box::new(inconsistent_view_outcome),
2162                                    location: location.clone(),
2163                                }
2164                            } else {
2165                                inconsistent_view_outcome
2166                            };
2167                            return Ok(outcome);
2168                        }
2169                    }
2170                }
2171                Ok(query_outcome)
2172            }
2173            PrepareQueryOutcome::Outcome(outcome) => Ok(outcome),
2174        }
2175    }
2176
2177    async fn get_conn(
2178        &mut self,
2179        name: Option<&str>,
2180        user: Option<&str>,
2181        password: Option<&str>,
2182    ) -> Result<&tokio_postgres::Client, tokio_postgres::Error> {
2183        match name {
2184            None => Ok(&self.client),
2185            Some(name) => {
2186                if !self.clients.contains_key(name) {
2187                    let addr = if matches!(user, Some("mz_system") | Some("mz_support")) {
2188                        self.internal_server_addr
2189                    } else if password.is_some() {
2190                        // Use password server for password authentication
2191                        self.password_server_addr
2192                    } else {
2193                        self.server_addr
2194                    };
2195                    let client = connect(addr, user, password).await?;
2196                    self.clients.insert(name.into(), client);
2197                }
2198                Ok(self.clients.get(name).unwrap())
2199            }
2200        }
2201    }
2202
2203    #[allow(clippy::disallowed_methods)]
2204    async fn run_simple<'r>(
2205        &mut self,
2206        conn: Option<&'r str>,
2207        user: Option<&'r str>,
2208        password: Option<&'r str>,
2209        sql: &'r str,
2210        sort: Sort,
2211        output: &'r Output,
2212        location: Location,
2213    ) -> Result<Outcome<'r>, anyhow::Error> {
2214        let actual = match self.get_conn(conn, user, password).await {
2215            Ok(client) => match client.simple_query(sql).await {
2216                Ok(result) => {
2217                    let mut rows = Vec::new();
2218
2219                    for m in result.into_iter() {
2220                        match m {
2221                            SimpleQueryMessage::Row(row) => {
2222                                let mut s = vec![];
2223                                for i in 0..row.len() {
2224                                    s.push(row.get(i).unwrap_or("NULL"));
2225                                }
2226                                rows.push(s.join(","));
2227                            }
2228                            SimpleQueryMessage::CommandComplete(count) => {
2229                                // This applies any sort on the COMPLETE line as
2230                                // well, but we do the same for the expected output.
2231                                rows.push(format!("COMPLETE {}", count));
2232                            }
2233                            SimpleQueryMessage::RowDescription(_) => {}
2234                            _ => panic!("unexpected"),
2235                        }
2236                    }
2237
2238                    if let Sort::Row = sort {
2239                        rows.sort();
2240                    }
2241
2242                    Output::Values(rows)
2243                }
2244                // Errors can contain multiple lines (say if there are details), and rewrite
2245                // sticks them each on their own line, so we need to split up the lines here to
2246                // each be its own String in the Vec.
2247                Err(error) => Output::Values(
2248                    error
2249                        .to_string_with_causes()
2250                        .lines()
2251                        .map(|s| s.to_string())
2252                        .collect(),
2253                ),
2254            },
2255            Err(error) => Output::Values(
2256                error
2257                    .to_string_with_causes()
2258                    .lines()
2259                    .map(|s| s.to_string())
2260                    .collect(),
2261            ),
2262        };
2263        if *output != actual {
2264            Ok(Outcome::OutputFailure {
2265                expected_output: output,
2266                actual_raw_output: vec![],
2267                actual_output: actual,
2268                location,
2269            })
2270        } else {
2271            Ok(Outcome::Success)
2272        }
2273    }
2274
2275    async fn check_catalog(&self) -> Result<(), anyhow::Error> {
2276        let url = format!(
2277            "http://{}/api/catalog/check",
2278            self.internal_http_server_addr
2279        );
2280        let response: serde_json::Value = reqwest::get(&url).await?.json().await?;
2281
2282        if let Some(inconsistencies) = response.get("err") {
2283            let inconsistencies = serde_json::to_string_pretty(&inconsistencies)
2284                .expect("serializing Value cannot fail");
2285            Err(anyhow::anyhow!("Catalog inconsistency\n{inconsistencies}"))
2286        } else {
2287            Ok(())
2288        }
2289    }
2290}
2291
2292async fn connect(
2293    addr: SocketAddr,
2294    user: Option<&str>,
2295    password: Option<&str>,
2296) -> Result<tokio_postgres::Client, tokio_postgres::Error> {
2297    let mut config = tokio_postgres::Config::new();
2298    config.host(addr.ip().to_string());
2299    config.port(addr.port());
2300    config.user(user.unwrap_or("materialize"));
2301    if let Some(password) = password {
2302        config.password(password);
2303    }
2304    let (client, connection) = config.connect(NoTls).await?;
2305
2306    task::spawn(|| "sqllogictest_connect", async move {
2307        if let Err(e) = connection.await {
2308            eprintln!("connection error: {}", e);
2309        }
2310    });
2311    Ok(client)
2312}
2313
2314pub trait WriteFmt {
2315    fn write_fmt(&self, fmt: fmt::Arguments<'_>);
2316}
2317
2318pub struct RunConfig<'a> {
2319    pub stdout: &'a dyn WriteFmt,
2320    pub stderr: &'a dyn WriteFmt,
2321    pub verbose: bool,
2322    pub quiet: bool,
2323    pub postgres_url: String,
2324    pub prefix: String,
2325    pub no_fail: bool,
2326    pub fail_fast: bool,
2327    pub auto_index_tables: bool,
2328    pub auto_index_selects: bool,
2329    pub auto_transactions: bool,
2330    pub enable_table_keys: bool,
2331    pub orchestrator_process_wrapper: Option<String>,
2332    pub tracing: TracingCliArgs,
2333    pub tracing_handle: TracingHandle,
2334    pub system_parameter_defaults: BTreeMap<String, String>,
2335    /// Persist state is handled specially because:
2336    /// - Persist background workers do not necessarily shut down immediately once the server is
2337    ///   shut down, and may panic if their storage is deleted out from under them.
2338    /// - It's safe for different databases to reference the same state: all data is scoped by UUID.
2339    pub persist_dir: TempDir,
2340    pub replicas: usize,
2341    pub replica_size: String,
2342}
2343
2344/// Indentation used for verbose output of SQL statements.
2345const PRINT_INDENT: usize = 4;
2346
2347fn print_record(config: &RunConfig<'_>, record: &Record) {
2348    match record {
2349        Record::Statement { sql, .. } | Record::Query { sql, .. } => {
2350            print_sql(config.stdout, sql, None)
2351        }
2352        Record::Simple { conn, sql, .. } => print_sql(config.stdout, sql, *conn),
2353        Record::Copy {
2354            table_name,
2355            tsv_path,
2356        } => {
2357            writeln!(
2358                config.stdout,
2359                "{}slt copy {} from {}",
2360                " ".repeat(PRINT_INDENT),
2361                table_name,
2362                tsv_path
2363            )
2364        }
2365        Record::User { user, .. } => {
2366            writeln!(config.stdout, "{}user {}", " ".repeat(PRINT_INDENT), user)
2367        }
2368        Record::ResetServer => {
2369            writeln!(config.stdout, "{}reset-server", " ".repeat(PRINT_INDENT))
2370        }
2371        Record::Halt => {
2372            writeln!(config.stdout, "{}halt", " ".repeat(PRINT_INDENT))
2373        }
2374        Record::HashThreshold { threshold } => {
2375            writeln!(
2376                config.stdout,
2377                "{}hash-threshold {}",
2378                " ".repeat(PRINT_INDENT),
2379                threshold
2380            )
2381        }
2382        Record::Replace {
2383            pattern,
2384            replacement,
2385        } => {
2386            writeln!(
2387                config.stdout,
2388                "{}replace {}  {}",
2389                " ".repeat(PRINT_INDENT),
2390                pattern,
2391                replacement
2392            )
2393        }
2394    }
2395}
2396
2397fn print_sql_if<'a>(stdout: &'a dyn WriteFmt, sql: &str, cond: bool) {
2398    if cond {
2399        print_sql(stdout, sql, None)
2400    }
2401}
2402
2403fn print_sql<'a>(stdout: &'a dyn WriteFmt, sql: &str, conn: Option<&str>) {
2404    let text = if let Some(conn) = conn {
2405        format!("[conn={}] {}", conn, sql)
2406    } else {
2407        sql.to_string()
2408    };
2409    writeln!(stdout, "{}", util::indent(&text, PRINT_INDENT))
2410}
2411
2412/// Regular expressions for matching error messages that should force a plan failure
2413/// in an inconsistent view outcome into a warning if the corresponding query succeeds.
2414const INCONSISTENT_VIEW_OUTCOME_WARNING_REGEXPS: [&str; 9] = [
2415    // The following are unfixable errors in indexed views given our
2416    // current constraints.
2417    "cannot materialize call to",
2418    "SHOW commands are not allowed in views",
2419    "cannot create view with unstable dependencies",
2420    "cannot use wildcard expansions or NATURAL JOINs in a view that depends on system objects",
2421    "no valid schema selected",
2422    r#"system schema '\w+' cannot be modified"#,
2423    r#"permission denied for (SCHEMA|CLUSTER) "(\w+\.)?\w+""#,
2424    // NOTE(vmarcos): Column ambiguity that could not be eliminated by our
2425    // currently implemented syntactic rewrites is considered unfixable.
2426    // In addition, if some column cannot be dealt with, e.g., in `ORDER BY`
2427    // references, we treat this condition as unfixable as well.
2428    r#"column "[\w\?]+" specified more than once"#,
2429    r#"column "(\w+\.)?\w+" does not exist"#,
2430];
2431
2432/// Evaluates if the given outcome should be returned directly or if it should
2433/// be wrapped as a warning. Note that this function should be used for outcomes
2434/// that can be judged in a context-independent manner, i.e., the outcome itself
2435/// provides enough information as to whether a warning should be emitted or not.
2436fn should_warn(outcome: &Outcome) -> bool {
2437    match outcome {
2438        Outcome::InconsistentViewOutcome { view_outcome, .. } => match view_outcome.as_ref() {
2439            Outcome::PlanFailure { error, .. } => {
2440                INCONSISTENT_VIEW_OUTCOME_WARNING_REGEXPS.iter().any(|s| {
2441                    Regex::new(s)
2442                        .expect("unexpected error in regular expression parsing")
2443                        .is_match(&error.to_string_with_causes())
2444                })
2445            }
2446            _ => false,
2447        },
2448        _ => false,
2449    }
2450}
2451
2452pub async fn run_string(
2453    runner: &mut Runner<'_>,
2454    source: &str,
2455    input: &str,
2456) -> Result<Outcomes, anyhow::Error> {
2457    runner.reset_database().await?;
2458    // `replace` directives are scoped to a single file; clear any registered by
2459    // a previous file so they do not leak into this one.
2460    runner.replacements.clear();
2461
2462    let mut outcomes = Outcomes::default();
2463    let mut parser = crate::parser::Parser::new(source, input);
2464    // Transactions are currently relatively slow. Since sqllogictest runs in a single connection
2465    // there should be no difference in having longer running transactions.
2466    let mut in_transaction = false;
2467    writeln!(runner.config.stdout, "--- {}", source);
2468
2469    for record in parser.parse_records()? {
2470        // In maximal-verbose mode, print the query before attempting to run
2471        // it. Running the query might panic, so it is important to print out
2472        // what query we are trying to run *before* we panic.
2473        if runner.config.verbose {
2474            print_record(runner.config, &record);
2475        }
2476
2477        let outcome = runner
2478            .run_record(&record, &mut in_transaction)
2479            .await
2480            .map_err(|err| format!("In {}:\n{}", source, err))
2481            .unwrap();
2482
2483        // Print warnings and failures in verbose mode.
2484        if !runner.config.quiet && !outcome.success() {
2485            if !runner.config.verbose {
2486                // If `verbose` is enabled, we'll already have printed the record,
2487                // so don't print it again. Yes, this is an ugly bit of logic.
2488                // Please don't try to consolidate it with the `print_record`
2489                // call above, as it's important to have a mode in which records
2490                // are printed before they are run, so that if running the
2491                // record panics, you can tell which record caused it.
2492                if !outcome.failure() {
2493                    writeln!(
2494                        runner.config.stdout,
2495                        "{}",
2496                        util::indent("Warning detected for: ", 4)
2497                    );
2498                }
2499                print_record(runner.config, &record);
2500            }
2501            if runner.config.verbose || outcome.failure() {
2502                writeln!(
2503                    runner.config.stdout,
2504                    "{}",
2505                    util::indent(&outcome.to_string(), 4)
2506                );
2507                writeln!(runner.config.stdout, "{}", util::indent("----", 4));
2508            }
2509        }
2510
2511        outcomes.stats[outcome.code()] += 1;
2512        if outcome.failure() {
2513            outcomes.details.push(format!("{}", outcome));
2514        }
2515
2516        if let Outcome::Bail { .. } = outcome {
2517            break;
2518        }
2519
2520        if runner.config.fail_fast && outcome.failure() {
2521            break;
2522        }
2523    }
2524    Ok(outcomes)
2525}
2526
2527pub async fn run_file(runner: &mut Runner<'_>, filename: &Path) -> Result<Outcomes, anyhow::Error> {
2528    let mut input = String::new();
2529    File::open(filename)?.read_to_string(&mut input)?;
2530    let outcomes = run_string(runner, &format!("{}", filename.display()), &input).await?;
2531    runner.check_catalog().await?;
2532
2533    Ok(outcomes)
2534}
2535
2536pub async fn rewrite_file(runner: &mut Runner<'_>, filename: &Path) -> Result<(), anyhow::Error> {
2537    runner.reset_database().await?;
2538    // `replace` directives are scoped to a single file; clear any registered by
2539    // a previous file so they do not leak into this one.
2540    runner.replacements.clear();
2541
2542    let mut file = OpenOptions::new().read(true).write(true).open(filename)?;
2543
2544    let mut input = String::new();
2545    file.read_to_string(&mut input)?;
2546
2547    let mut buf = RewriteBuffer::new(&input);
2548
2549    let mut parser = crate::parser::Parser::new(filename.to_str().unwrap_or(""), &input);
2550    writeln!(runner.config.stdout, "--- {}", filename.display());
2551    let mut in_transaction = false;
2552
2553    fn append_values_output(
2554        buf: &mut RewriteBuffer,
2555        input: &String,
2556        expected_output: &str,
2557        mode: &Mode,
2558        types: &Vec<Type>,
2559        column_names: Option<&Vec<ColumnName>>,
2560        actual_output: &Vec<String>,
2561        multiline: bool,
2562    ) {
2563        buf.append_header(input, expected_output, column_names);
2564
2565        for (i, row) in actual_output.chunks(types.len()).enumerate() {
2566            match mode {
2567                // In Cockroach mode, output each row on its own line, with
2568                // two spaces between each column.
2569                Mode::Cockroach => {
2570                    if i != 0 {
2571                        buf.append("\n");
2572                    }
2573
2574                    if row.len() == 0 {
2575                        // nothing to do
2576                    } else if row.len() == 1 {
2577                        // If there is only one column, then there is no need for space
2578                        // substitution, so we only do newline substitution.
2579                        if multiline {
2580                            buf.append(&row[0]);
2581                        } else {
2582                            buf.append(&row[0].replace('\n', "⏎"))
2583                        }
2584                    } else {
2585                        // Substitute spaces with ␠ to avoid mistaking the spaces in the result
2586                        // values with spaces that separate columns.
2587                        buf.append(
2588                            &row.iter()
2589                                .map(|col| {
2590                                    let mut col = col.replace(' ', "␠");
2591                                    if !multiline {
2592                                        col = col.replace('\n', "⏎");
2593                                    }
2594                                    col
2595                                })
2596                                .join("  "),
2597                        );
2598                    }
2599                }
2600                // In standard mode, output each value on its own line,
2601                // and ignore row boundaries.
2602                // No need to substitute spaces, because every value (not row) is on a separate
2603                // line. But we do need to substitute newlines.
2604                Mode::Standard => {
2605                    for (j, col) in row.iter().enumerate() {
2606                        if i != 0 || j != 0 {
2607                            buf.append("\n");
2608                        }
2609                        buf.append(&if multiline {
2610                            col.clone()
2611                        } else {
2612                            col.replace('\n', "⏎")
2613                        });
2614                    }
2615                }
2616            }
2617        }
2618    }
2619
2620    for record in parser.parse_records()? {
2621        let outcome = runner.run_record(&record, &mut in_transaction).await?;
2622
2623        match (&record, &outcome) {
2624            // If we see an output failure for a query, rewrite the expected output
2625            // to match the observed output.
2626            (
2627                Record::Query {
2628                    output:
2629                        Ok(QueryOutput {
2630                            mode,
2631                            output: Output::Values(_),
2632                            output_str: expected_output,
2633                            types,
2634                            column_names,
2635                            multiline,
2636                            ..
2637                        }),
2638                    ..
2639                },
2640                Outcome::OutputFailure {
2641                    actual_output: Output::Values(actual_output),
2642                    ..
2643                },
2644            ) => {
2645                append_values_output(
2646                    &mut buf,
2647                    &input,
2648                    expected_output,
2649                    mode,
2650                    types,
2651                    column_names.as_ref(),
2652                    actual_output,
2653                    *multiline,
2654                );
2655            }
2656            (
2657                Record::Query {
2658                    output:
2659                        Ok(QueryOutput {
2660                            mode,
2661                            output: Output::Values(_),
2662                            output_str: expected_output,
2663                            types,
2664                            multiline,
2665                            ..
2666                        }),
2667                    ..
2668                },
2669                Outcome::WrongColumnNames {
2670                    actual_column_names,
2671                    actual_output: Output::Values(actual_output),
2672                    ..
2673                },
2674            ) => {
2675                append_values_output(
2676                    &mut buf,
2677                    &input,
2678                    expected_output,
2679                    mode,
2680                    types,
2681                    Some(actual_column_names),
2682                    actual_output,
2683                    *multiline,
2684                );
2685            }
2686            (
2687                Record::Query {
2688                    output:
2689                        Ok(QueryOutput {
2690                            output: Output::Hashed { .. },
2691                            output_str: expected_output,
2692                            column_names,
2693                            ..
2694                        }),
2695                    ..
2696                },
2697                Outcome::OutputFailure {
2698                    actual_output: Output::Hashed { num_values, md5 },
2699                    ..
2700                },
2701            ) => {
2702                buf.append_header(&input, expected_output, column_names.as_ref());
2703
2704                buf.append(format!("{} values hashing to {}\n", num_values, md5).as_str())
2705            }
2706            (
2707                Record::Simple {
2708                    output_str: expected_output,
2709                    ..
2710                },
2711                Outcome::OutputFailure {
2712                    actual_output: Output::Values(actual_output),
2713                    ..
2714                },
2715            ) => {
2716                buf.append_header(&input, expected_output, None);
2717
2718                for (i, row) in actual_output.iter().enumerate() {
2719                    if i != 0 {
2720                        buf.append("\n");
2721                    }
2722                    buf.append(row);
2723                }
2724            }
2725            (
2726                Record::Query {
2727                    sql,
2728                    output: Err(err),
2729                    ..
2730                },
2731                outcome,
2732            )
2733            | (
2734                Record::Statement {
2735                    expected_error: Some(err),
2736                    sql,
2737                    ..
2738                },
2739                outcome,
2740            ) if outcome.err_msg().is_some() => {
2741                buf.rewrite_expected_error(&input, err, &outcome.err_msg().unwrap(), sql)
2742            }
2743            (_, Outcome::Success) => {}
2744            _ => bail!("unexpected: {:?} {:?}", record, outcome),
2745        }
2746    }
2747
2748    file.set_len(0)?;
2749    file.seek(SeekFrom::Start(0))?;
2750    file.write_all(buf.finish().as_bytes())?;
2751    file.sync_all()?;
2752    Ok(())
2753}
2754
2755/// Provides a means to rewrite the `.slt` file while iterating over it.
2756///
2757/// This struct takes the slt file as its `input`, tracks a cursor into it
2758/// (`input_offset`), and provides a buffer (`output`) to store the rewritten
2759/// results.
2760///
2761/// Functions that modify the file will lazily move `input` into `output` using
2762/// `flush_to`. However, those calls should all be interior to other functions.
2763#[derive(Debug)]
2764struct RewriteBuffer<'a> {
2765    input: &'a str,
2766    input_offset: usize,
2767    output: String,
2768}
2769
2770impl<'a> RewriteBuffer<'a> {
2771    fn new(input: &'a str) -> RewriteBuffer<'a> {
2772        RewriteBuffer {
2773            input,
2774            input_offset: 0,
2775            output: String::new(),
2776        }
2777    }
2778
2779    fn flush_to(&mut self, offset: usize) {
2780        assert!(offset >= self.input_offset);
2781        let chunk = &self.input[self.input_offset..offset];
2782        self.output.push_str(chunk);
2783        self.input_offset = offset;
2784    }
2785
2786    fn skip_to(&mut self, offset: usize) {
2787        assert!(offset >= self.input_offset);
2788        self.input_offset = offset;
2789    }
2790
2791    fn append(&mut self, s: &str) {
2792        self.output.push_str(s);
2793    }
2794
2795    fn append_header(
2796        &mut self,
2797        input: &String,
2798        expected_output: &str,
2799        column_names: Option<&Vec<ColumnName>>,
2800    ) {
2801        // Output everything before this record.
2802        // TODO(benesch): is it possible to rewrite this to avoid `as`?
2803        #[allow(clippy::as_conversions)]
2804        let offset = expected_output.as_ptr() as usize - input.as_ptr() as usize;
2805        self.flush_to(offset);
2806        self.skip_to(offset + expected_output.len());
2807
2808        // Attempt to install the result separator (----), if it does
2809        // not already exist.
2810        if self.peek_last(5) == "\n----" {
2811            self.append("\n");
2812        } else if self.peek_last(6) != "\n----\n" {
2813            self.append("\n----\n");
2814        }
2815
2816        let Some(names) = column_names else {
2817            return;
2818        };
2819        self.append(
2820            &names
2821                .iter()
2822                .map(|name| name.replace(' ', "␠"))
2823                .collect::<Vec<_>>()
2824                .join(" "),
2825        );
2826        self.append("\n");
2827    }
2828
2829    fn rewrite_expected_error(
2830        &mut self,
2831        input: &String,
2832        old_err: &str,
2833        new_err: &str,
2834        query: &str,
2835    ) {
2836        // Output everything before this error message.
2837        // TODO(benesch): is it possible to rewrite this to avoid `as`?
2838        #[allow(clippy::as_conversions)]
2839        let err_offset = old_err.as_ptr() as usize - input.as_ptr() as usize;
2840        self.flush_to(err_offset);
2841        self.append(new_err);
2842        self.append("\n");
2843        self.append(query);
2844        // TODO(benesch): is it possible to rewrite this to avoid `as`?
2845        #[allow(clippy::as_conversions)]
2846        self.skip_to(query.as_ptr() as usize - input.as_ptr() as usize + query.len())
2847    }
2848
2849    fn peek_last(&self, n: usize) -> &str {
2850        &self.output[self.output.len() - n..]
2851    }
2852
2853    fn finish(mut self) -> String {
2854        self.flush_to(self.input.len());
2855        self.output
2856    }
2857}
2858
2859/// Generates view creation, view indexing, view querying, and view
2860/// dropping SQL commands for a given `SELECT` query. If the number
2861/// of attributes produced by the query is known, the view commands
2862/// are specialized to avoid issues with column ambiguity. This
2863/// function is a helper for `--auto_index_selects` and assumes that
2864/// the provided input SQL has already been run through the parser,
2865/// resulting in a valid `SELECT` statement.
2866fn generate_view_sql(
2867    sql: &str,
2868    view_uuid: &Simple,
2869    num_attributes: Option<usize>,
2870    expected_column_names: Option<Vec<ColumnName>>,
2871) -> (String, String, String, String) {
2872    // To create the view, re-parse the sql; note that we must find exactly
2873    // one statement and it must be a `SELECT`.
2874    // NOTE(vmarcos): Direct string manipulation was attempted while
2875    // prototyping the code below, which avoids the extra parsing and
2876    // data structure cloning. However, running DDL is so slow that
2877    // it did not matter in terms of runtime. We can revisit this if
2878    // DDL cost drops dramatically in the future.
2879    let stmts = parser::parse_statements(sql).unwrap_or_default();
2880    assert!(stmts.len() == 1);
2881    let (query, query_as_of) = match &stmts[0].ast {
2882        Statement::Select(stmt) => (&stmt.query, &stmt.as_of),
2883        _ => unreachable!("This function should only be called for SELECTs"),
2884    };
2885
2886    // Prior to creating the view, process the `ORDER BY` clause of
2887    // the `SELECT` query, if any. Ordering is not preserved when a
2888    // view includes an `ORDER BY` clause and must be re-enforced by
2889    // an external `ORDER BY` clause when querying the view.
2890    let (view_order_by, extra_columns, distinct) = if num_attributes.is_none() {
2891        (query.order_by.clone(), vec![], None)
2892    } else {
2893        derive_order_by(&query.body, &query.order_by)
2894    };
2895
2896    // Since one-shot SELECT statements may contain ambiguous column names,
2897    // we either use the expected column names, if that option was
2898    // provided, or else just rename the output schema of the view
2899    // using numerically increasing attribute names, whenever possible.
2900    // This strategy makes it possible to use `CREATE INDEX`, thus
2901    // matching the behavior of the option `auto_index_tables`. However,
2902    // we may be presented with a `SELECT *` query, in which case the parser
2903    // does not produce sufficient information to allow us to compute
2904    // the number of output columns. In the latter case, we are supplied
2905    // with `None` for `num_attributes` and just employ the command
2906    // `CREATE DEFAULT INDEX` instead. Additionally, the view is created
2907    // without schema renaming. This strategy is insufficient to dodge
2908    // column name ambiguity in all cases, but we assume here that we
2909    // can adjust the (hopefully) small number of tests that eventually
2910    // challenge us in this particular way.
2911    let name = UnresolvedItemName(vec![Ident::new_unchecked(format!("v{}", view_uuid))]);
2912    let projection = expected_column_names.map_or_else(
2913        || {
2914            num_attributes.map_or(vec![], |n| {
2915                (1..=n)
2916                    .map(|i| Ident::new_unchecked(format!("a{i}")))
2917                    .collect()
2918            })
2919        },
2920        |cols| {
2921            cols.iter()
2922                .map(|c| Ident::new_unchecked(c.as_str()))
2923                .collect()
2924        },
2925    );
2926    let columns: Vec<Ident> = projection
2927        .iter()
2928        .cloned()
2929        .chain(extra_columns.iter().map(|item| {
2930            if let SelectItem::Expr {
2931                expr: _,
2932                alias: Some(ident),
2933            } = item
2934            {
2935                ident.clone()
2936            } else {
2937                unreachable!("alias must be given for extra column")
2938            }
2939        }))
2940        .collect();
2941
2942    // Build a `CREATE VIEW` with the columns computed above.
2943    let mut query = query.clone();
2944    if extra_columns.len() > 0 {
2945        match &mut query.body {
2946            SetExpr::Select(stmt) => stmt.projection.extend(extra_columns.iter().cloned()),
2947            _ => unimplemented!("cannot yet rewrite projections of nested queries"),
2948        }
2949    }
2950    let create_view = AstStatement::<Raw>::CreateView(CreateViewStatement {
2951        if_exists: IfExistsBehavior::Error,
2952        temporary: false,
2953        definition: ViewDefinition {
2954            name: name.clone(),
2955            columns: columns.clone(),
2956            query,
2957        },
2958    })
2959    .to_ast_string_stable();
2960
2961    // We then create either a `CREATE INDEX` or a `CREATE DEFAULT INDEX`
2962    // statement, depending on whether we could obtain the number of
2963    // attributes from the original `SELECT`.
2964    let create_index = AstStatement::<Raw>::CreateIndex(CreateIndexStatement {
2965        name: None,
2966        in_cluster: None,
2967        on_name: RawItemName::Name(name.clone()),
2968        key_parts: if columns.len() == 0 {
2969            None
2970        } else {
2971            Some(
2972                columns
2973                    .iter()
2974                    .map(|ident| Expr::Identifier(vec![ident.clone()]))
2975                    .collect(),
2976            )
2977        },
2978        with_options: Vec::new(),
2979        if_not_exists: false,
2980    })
2981    .to_ast_string_stable();
2982
2983    // Assert if DISTINCT semantics are unchanged from view
2984    let distinct_unneeded = extra_columns.len() == 0
2985        || match distinct {
2986            None | Some(Distinct::On(_)) => true,
2987            Some(Distinct::EntireRow) => false,
2988        };
2989    let distinct = if distinct_unneeded { None } else { distinct };
2990
2991    // `SELECT [* | {projection}] FROM {name} [ORDER BY {view_order_by}]`
2992    let view_sql = AstStatement::<Raw>::Select(SelectStatement {
2993        query: Query {
2994            ctes: CteBlock::Simple(vec![]),
2995            body: SetExpr::Select(Box::new(Select {
2996                distinct,
2997                projection: if projection.len() == 0 {
2998                    vec![SelectItem::Wildcard]
2999                } else {
3000                    projection
3001                        .iter()
3002                        .map(|ident| SelectItem::Expr {
3003                            expr: Expr::Identifier(vec![ident.clone()]),
3004                            alias: None,
3005                        })
3006                        .collect()
3007                },
3008                from: vec![TableWithJoins {
3009                    relation: TableFactor::Table {
3010                        name: RawItemName::Name(name.clone()),
3011                        alias: None,
3012                    },
3013                    joins: vec![],
3014                }],
3015                selection: None,
3016                group_by: vec![],
3017                having: None,
3018                qualify: None,
3019                options: vec![],
3020            })),
3021            order_by: view_order_by,
3022            limit: None,
3023            offset: None,
3024        },
3025        as_of: query_as_of.clone(),
3026    })
3027    .to_ast_string_stable();
3028
3029    // `DROP VIEW {name}`
3030    let drop_view = AstStatement::<Raw>::DropObjects(DropObjectsStatement {
3031        object_type: ObjectType::View,
3032        if_exists: false,
3033        names: vec![UnresolvedObjectName::Item(name)],
3034        cascade: false,
3035    })
3036    .to_ast_string_stable();
3037
3038    (create_view, create_index, view_sql, drop_view)
3039}
3040
3041/// Analyzes the provided query `body` to derive the number of
3042/// attributes in the query. We only consider syntactic cues,
3043/// so we may end up deriving `None` for the number of attributes
3044/// as a conservative approximation.
3045fn derive_num_attributes(body: &SetExpr<Raw>) -> Option<usize> {
3046    let Some((projection, _)) = find_projection(body) else {
3047        return None;
3048    };
3049    derive_num_attributes_from_projection(projection)
3050}
3051
3052/// Analyzes a query's `ORDER BY` clause to derive an `ORDER BY`
3053/// clause that makes numeric references to any expressions in
3054/// the projection and generated-attribute references to expressions
3055/// that need to be added as extra columns to the projection list.
3056/// The rewritten `ORDER BY` clause is then usable when querying a
3057/// view that contains the same `SELECT` as the given query.
3058/// This function returns both the rewritten `ORDER BY` clause
3059/// as well as a list of extra columns that need to be added
3060/// to the query's projection for the `ORDER BY` clause to
3061/// succeed.
3062fn derive_order_by(
3063    body: &SetExpr<Raw>,
3064    order_by: &Vec<OrderByExpr<Raw>>,
3065) -> (
3066    Vec<OrderByExpr<Raw>>,
3067    Vec<SelectItem<Raw>>,
3068    Option<Distinct<Raw>>,
3069) {
3070    let Some((projection, distinct)) = find_projection(body) else {
3071        return (vec![], vec![], None);
3072    };
3073    let (view_order_by, extra_columns) = derive_order_by_from_projection(projection, order_by);
3074    (view_order_by, extra_columns, distinct.clone())
3075}
3076
3077/// Finds the projection list in a `SELECT` query body.
3078fn find_projection(body: &SetExpr<Raw>) -> Option<(&Vec<SelectItem<Raw>>, &Option<Distinct<Raw>>)> {
3079    // Iterate to peel off the query body until the query's
3080    // projection list is found.
3081    let mut set_expr = body;
3082    loop {
3083        match set_expr {
3084            SetExpr::Select(select) => {
3085                return Some((&select.projection, &select.distinct));
3086            }
3087            SetExpr::SetOperation { left, .. } => set_expr = left.as_ref(),
3088            SetExpr::Query(query) => set_expr = &query.body,
3089            _ => return None,
3090        }
3091    }
3092}
3093
3094/// Computes the number of attributes that are obtained by the
3095/// projection of a `SELECT` query. The projection may include
3096/// wildcards, in which case the analysis just returns `None`.
3097fn derive_num_attributes_from_projection(projection: &Vec<SelectItem<Raw>>) -> Option<usize> {
3098    let mut num_attributes = 0usize;
3099    for item in projection.iter() {
3100        let SelectItem::Expr { expr, .. } = item else {
3101            return None;
3102        };
3103        match expr {
3104            Expr::QualifiedWildcard(..) | Expr::WildcardAccess(..) => {
3105                return None;
3106            }
3107            _ => {
3108                num_attributes += 1;
3109            }
3110        }
3111    }
3112    Some(num_attributes)
3113}
3114
3115/// Computes an `ORDER BY` clause with only numeric references
3116/// from given projection and `ORDER BY` of a `SELECT` query.
3117/// If the derivation fails to match a given expression, the
3118/// matched prefix is returned. Note that this could be empty.
3119fn derive_order_by_from_projection(
3120    projection: &Vec<SelectItem<Raw>>,
3121    order_by: &Vec<OrderByExpr<Raw>>,
3122) -> (Vec<OrderByExpr<Raw>>, Vec<SelectItem<Raw>>) {
3123    let mut view_order_by: Vec<OrderByExpr<Raw>> = vec![];
3124    let mut extra_columns: Vec<SelectItem<Raw>> = vec![];
3125    for order_by_expr in order_by.iter() {
3126        let query_expr = &order_by_expr.expr;
3127        let view_expr = match query_expr {
3128            Expr::Value(mz_sql_parser::ast::Value::Number(_)) => query_expr.clone(),
3129            _ => {
3130                // Find expression in query projection, if we can.
3131                if let Some(i) = projection.iter().position(|item| match item {
3132                    SelectItem::Expr { expr, alias } => {
3133                        expr == query_expr
3134                            || match query_expr {
3135                                Expr::Identifier(ident) => {
3136                                    ident.len() == 1 && Some(&ident[0]) == alias.as_ref()
3137                                }
3138                                _ => false,
3139                            }
3140                    }
3141                    SelectItem::Wildcard => false,
3142                }) {
3143                    Expr::Value(mz_sql_parser::ast::Value::Number((i + 1).to_string()))
3144                } else {
3145                    // If the expression is not found in the
3146                    // projection, add extra column.
3147                    let ident = Ident::new_unchecked(format!(
3148                        "a{}",
3149                        (projection.len() + extra_columns.len() + 1)
3150                    ));
3151                    extra_columns.push(SelectItem::Expr {
3152                        expr: query_expr.clone(),
3153                        alias: Some(ident.clone()),
3154                    });
3155                    Expr::Identifier(vec![ident])
3156                }
3157            }
3158        };
3159        view_order_by.push(OrderByExpr {
3160            expr: view_expr,
3161            asc: order_by_expr.asc,
3162            nulls_last: order_by_expr.nulls_last,
3163        });
3164    }
3165    (view_order_by, extra_columns)
3166}
3167
3168/// Returns extra statements to execute after `stmt` is executed.
3169fn mutate(sql: &str) -> Vec<String> {
3170    let stmts = parser::parse_statements(sql).unwrap_or_default();
3171    let mut additional = Vec::new();
3172    for stmt in stmts {
3173        match stmt.ast {
3174            AstStatement::CreateTable(stmt) => additional.push(
3175                // CREATE TABLE -> CREATE INDEX. Specify all columns manually in case CREATE
3176                // DEFAULT INDEX ever goes away.
3177                AstStatement::<Raw>::CreateIndex(CreateIndexStatement {
3178                    name: None,
3179                    in_cluster: None,
3180                    on_name: RawItemName::Name(stmt.name.clone()),
3181                    key_parts: Some(
3182                        stmt.columns
3183                            .iter()
3184                            .map(|def| Expr::Identifier(vec![def.name.clone()]))
3185                            .collect(),
3186                    ),
3187                    with_options: Vec::new(),
3188                    if_not_exists: false,
3189                })
3190                .to_ast_string_stable(),
3191            ),
3192            _ => {}
3193        }
3194    }
3195    additional
3196}
3197
3198#[mz_ore::test]
3199#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
3200fn test_generate_view_sql() {
3201    let uuid = Uuid::parse_str("67e5504410b1426f9247bb680e5fe0c8").unwrap();
3202    let cases = vec![
3203        (("SELECT * FROM t", None, None),
3204        (
3205            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" AS SELECT * FROM "t""#.to_string(),
3206            r#"CREATE DEFAULT INDEX ON "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3207            r#"SELECT * FROM "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3208            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3209        )),
3210        (("SELECT a, b, c FROM t1, t2", Some(3), Some(vec![ColumnName::from("a"), ColumnName::from("b"), ColumnName::from("c")])),
3211        (
3212            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a", "b", "c") AS SELECT "a", "b", "c" FROM "t1", "t2""#.to_string(),
3213            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a", "b", "c")"#.to_string(),
3214            r#"SELECT "a", "b", "c" FROM "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3215            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3216        )),
3217        (("SELECT a, b, c FROM t1, t2", Some(3), None),
3218        (
3219            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3") AS SELECT "a", "b", "c" FROM "t1", "t2""#.to_string(),
3220            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3")"#.to_string(),
3221            r#"SELECT "a1", "a2", "a3" FROM "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3222            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3223        )),
3224        // A case with ambiguity that is accepted by the function, illustrating that
3225        // our measures to dodge this issue are imperfect.
3226        (("SELECT * FROM (SELECT a, sum(b) AS a FROM t GROUP BY a)", None, None),
3227        (
3228            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" AS SELECT * FROM (SELECT "a", "sum"("b") AS "a" FROM "t" GROUP BY "a")"#.to_string(),
3229            r#"CREATE DEFAULT INDEX ON "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3230            r#"SELECT * FROM "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3231            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3232        )),
3233        (("SELECT a, b, b + d AS c, a + b AS d FROM t1, t2 ORDER BY a, c, a + b", Some(4), Some(vec![ColumnName::from("a"), ColumnName::from("b"), ColumnName::from("c"), ColumnName::from("d")])),
3234        (
3235            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a", "b", "c", "d") AS SELECT "a", "b", "b" + "d" AS "c", "a" + "b" AS "d" FROM "t1", "t2" ORDER BY "a", "c", "a" + "b""#.to_string(),
3236            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a", "b", "c", "d")"#.to_string(),
3237            r#"SELECT "a", "b", "c", "d" FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY 1, 3, 4"#.to_string(),
3238            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3239        )),
3240        (("((SELECT 1 AS a UNION SELECT 2 AS b) UNION SELECT 3 AS c) ORDER BY a", Some(1), None),
3241        (
3242            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a1") AS (SELECT 1 AS "a" UNION SELECT 2 AS "b") UNION SELECT 3 AS "c" ORDER BY "a""#.to_string(),
3243            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a1")"#.to_string(),
3244            r#"SELECT "a1" FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY 1"#.to_string(),
3245            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3246        )),
3247        (("SELECT * FROM (SELECT a, sum(b) AS a FROM t GROUP BY a) ORDER BY 1", None, None),
3248        (
3249            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" AS SELECT * FROM (SELECT "a", "sum"("b") AS "a" FROM "t" GROUP BY "a") ORDER BY 1"#.to_string(),
3250            r#"CREATE DEFAULT INDEX ON "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3251            r#"SELECT * FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY 1"#.to_string(),
3252            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3253        )),
3254        (("SELECT * FROM (SELECT a, sum(b) AS a FROM t GROUP BY a) ORDER BY a", None, None),
3255        (
3256            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" AS SELECT * FROM (SELECT "a", "sum"("b") AS "a" FROM "t" GROUP BY "a") ORDER BY "a""#.to_string(),
3257            r#"CREATE DEFAULT INDEX ON "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3258            r#"SELECT * FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY "a""#.to_string(),
3259            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3260        )),
3261        (("SELECT a, sum(b) AS a FROM t GROUP BY a, c ORDER BY a, c", Some(2), None),
3262        (
3263            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3") AS SELECT "a", "sum"("b") AS "a", "c" AS "a3" FROM "t" GROUP BY "a", "c" ORDER BY "a", "c""#.to_string(),
3264            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3")"#.to_string(),
3265            r#"SELECT "a1", "a2" FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY 1, "a3""#.to_string(),
3266            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3267        )),
3268        (("SELECT a, sum(b) AS a FROM t GROUP BY a, c ORDER BY c, a", Some(2), None),
3269        (
3270            r#"CREATE VIEW "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3") AS SELECT "a", "sum"("b") AS "a", "c" AS "a3" FROM "t" GROUP BY "a", "c" ORDER BY "c", "a""#.to_string(),
3271            r#"CREATE INDEX ON "v67e5504410b1426f9247bb680e5fe0c8" ("a1", "a2", "a3")"#.to_string(),
3272            r#"SELECT "a1", "a2" FROM "v67e5504410b1426f9247bb680e5fe0c8" ORDER BY "a3", 1"#.to_string(),
3273            r#"DROP VIEW "v67e5504410b1426f9247bb680e5fe0c8""#.to_string(),
3274        )),
3275    ];
3276    for ((sql, num_attributes, expected_column_names), expected) in cases {
3277        let view_sql =
3278            generate_view_sql(sql, uuid.as_simple(), num_attributes, expected_column_names);
3279        assert_eq!(expected, view_sql);
3280    }
3281}
3282
3283#[mz_ore::test]
3284fn test_mutate() {
3285    let cases = vec![
3286        ("CREATE TABLE t ()", vec![r#"CREATE INDEX ON "t" ()"#]),
3287        (
3288            "CREATE TABLE t (a INT)",
3289            vec![r#"CREATE INDEX ON "t" ("a")"#],
3290        ),
3291        (
3292            "CREATE TABLE t (a INT, b TEXT)",
3293            vec![r#"CREATE INDEX ON "t" ("a", "b")"#],
3294        ),
3295        // Invalid syntax works, just returns nothing.
3296        ("BAD SYNTAX", Vec::new()),
3297    ];
3298    for (sql, expected) in cases {
3299        let stmts = mutate(sql);
3300        assert_eq!(expected, stmts, "sql: {sql}");
3301    }
3302}
3303
3304#[mz_ore::test]
3305fn test_strip_crdb_table_items() {
3306    let cases = vec![
3307        (
3308            "CREATE TABLE t (a INT, INDEX foo (a), b INT)",
3309            Some("CREATE TABLE t (a INT, b INT)"),
3310        ),
3311        (
3312            "CREATE TABLE t (a INT, INDEX(a))",
3313            Some("CREATE TABLE t (a INT)"),
3314        ),
3315        (
3316            "CREATE TABLE t (a INT, UNIQUE INDEX foo (a), INVERTED INDEX bar (a))",
3317            Some("CREATE TABLE t (a INT)"),
3318        ),
3319        (
3320            "CREATE TABLE t (k INT, v STRING, FAMILY \"primary\" (k, v))",
3321            Some("CREATE TABLE t (k INT, v STRING)"),
3322        ),
3323        // Partial and expression indexes have nested clauses.
3324        (
3325            "CREATE TABLE t (a INT, INDEX (a) WHERE a = 0, INDEX ((a + 1)))",
3326            Some("CREATE TABLE t (a INT)"),
3327        ),
3328        // Constraints are kept.
3329        (
3330            "CREATE TABLE t (a INT, UNIQUE (a), CHECK (a > 0))",
3331            Some("CREATE TABLE t (a INT, UNIQUE (a), CHECK (a > 0))"),
3332        ),
3333        // Column names that merely start with a stripped keyword are kept.
3334        (
3335            "CREATE TABLE t (index_col INT, \"index\" INT)",
3336            Some("CREATE TABLE t (index_col INT, \"index\" INT)"),
3337        ),
3338        // Parens inside strings do not confuse the item splitter.
3339        (
3340            "CREATE TABLE t (a TEXT DEFAULT 'a,(b', INDEX (a))",
3341            Some("CREATE TABLE t (a TEXT DEFAULT 'a,(b')"),
3342        ),
3343        ("SELECT 1", None),
3344    ];
3345    for (sql, expected) in cases {
3346        let stripped = strip_crdb_table_items(sql);
3347        assert_eq!(expected.map(|e| e.to_string()), stripped, "sql: {sql}");
3348    }
3349}