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