mz_sql/pure/
source_preparation.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::collections::BTreeMap;
11use std::sync::{Arc, OnceLock};
12
13use mz_storage_types::configuration::StorageConfiguration;
14
15use super::{
16    AlterSourcePreparation, CreateSourcePreparation, CsrAvroPreparation, CsrProtobufPreparation,
17    GlueAvroPreparation, PreparedCreateSource, PurifiedSourceExport, PurifiedStatement,
18    SourceFormatOptions, TableFromSourcePreparation, execute_alter_source_preparation,
19    execute_create_source_preparation, execute_table_from_source_preparation,
20    finish_table_from_source, prepare_alter_source, prepare_csr_connection_avro,
21    prepare_csr_connection_proto, prepare_glue_connection_avro, prepare_table_from_source,
22};
23use crate::ast::{
24    AvroSchema, CsrSeedAvro, CsrSeedProtobuf, Format, FormatSpecifier, GlueAvroSeed,
25    ProtobufSchema, SourceEnvelope, Statement,
26};
27use crate::catalog::SessionCatalog;
28use crate::names::Aug;
29use crate::plan::PlanError;
30
31/// Statement-local external results, reusable only for matching current requests.
32#[derive(Clone, Debug)]
33pub struct SourcePreparation {
34    now: u64,
35    clock_as_of: OnceLock<u64>,
36    inputs: BTreeMap<usize, PreparedSourceInput>,
37}
38
39/// The result of resolving source preparation against a catalog.
40#[derive(Debug)]
41pub enum SourcePreparationStatus {
42    Ready(Box<PurifiedStatement>),
43    Pending(SourcePreparationRequest),
44}
45
46/// An external request derived from the current catalog and submitted statement.
47#[derive(Clone, Debug, PartialEq, Eq)]
48pub struct SourcePreparationRequest {
49    position: usize,
50    kind: RequestKind,
51}
52
53/// Catalog-derived inputs used by external source preparation.
54///
55/// Equal requests reuse cached results, so each variant must include every
56/// catalog-derived input its external reads depend on.
57/// Each request carries every catalog-derived input its upstream reads use. Equal requests reuse a cached result,
58/// so an input missing here would let a stale result be reused.
59#[derive(Clone, Debug, PartialEq, Eq)]
60enum RequestKind {
61    Alter(Box<AlterSourcePreparation>),
62    Source(Box<CreateSourcePreparation>),
63    Table(Box<TableFromSourcePreparation>),
64    Protobuf(Box<CsrProtobufPreparation>),
65    Avro(Box<CsrAvroPreparation>),
66    Glue(Box<GlueAvroPreparation>),
67}
68
69/// An external result paired with the exact request that produced it.
70#[derive(Clone, Debug)]
71struct PreparedSourceInput {
72    request: SourcePreparationRequest,
73    value: Arc<PreparationValue>,
74}
75
76/// An external completion whose success or failure must be checked against current inputs.
77#[derive(Debug)]
78#[must_use = "pass completions to SourcePreparation::finish to validate their inputs"]
79pub struct SourcePreparationResponse {
80    request: SourcePreparationRequest,
81    result: Result<Arc<PreparationValue>, PlanError>,
82}
83
84#[derive(Debug)]
85enum PreparationValue {
86    Alter(Box<PurifiedStatement>),
87    Source(Box<PreparedCreateSource>),
88    Table(Box<(PurifiedSourceExport, SourceFormatOptions)>),
89    Protobuf(CsrSeedProtobuf),
90    Avro(CsrSeedAvro),
91    Glue(GlueAvroSeed),
92}
93
94enum ResolveError {
95    Plan(Box<PlanError>),
96    Pending(SourcePreparationRequest),
97}
98
99impl From<PlanError> for ResolveError {
100    fn from(error: PlanError) -> Self {
101        Self::Plan(Box::new(error))
102    }
103}
104
105impl SourcePreparation {
106    pub fn new(now: u64) -> Self {
107        Self {
108            now,
109            clock_as_of: OnceLock::new(),
110            inputs: BTreeMap::new(),
111        }
112    }
113
114    /// Recompute catalog decisions, requesting external work when its inputs do not match.
115    pub fn resolve(
116        &self,
117        catalog: &dyn SessionCatalog,
118        statement: Statement<Aug>,
119    ) -> Result<SourcePreparationStatus, PlanError> {
120        let mut cursor = Cursor {
121            preparation: self,
122            next: 0,
123        };
124        match cursor.resolve(catalog, statement) {
125            Ok(statement) => Ok(SourcePreparationStatus::Ready(Box::new(statement))),
126            Err(ResolveError::Pending(request)) => Ok(SourcePreparationStatus::Pending(request)),
127            Err(ResolveError::Plan(error)) => Err(*error),
128        }
129    }
130
131    /// Accept completed work, reporting a failure only if its request is still required.
132    pub fn finish(
133        &mut self,
134        catalog: &dyn SessionCatalog,
135        statement: Statement<Aug>,
136        response: SourcePreparationResponse,
137    ) -> Result<SourcePreparationStatus, PlanError> {
138        let status = self.resolve(catalog, statement.clone())?;
139        let required = matches!(
140            &status,
141            SourcePreparationStatus::Pending(request) if *request == response.request
142        );
143        if !required {
144            return Ok(status);
145        }
146        let value = response.result?;
147        self.inputs.insert(
148            response.request.position,
149            PreparedSourceInput {
150                request: response.request,
151                value,
152            },
153        );
154        self.resolve(catalog, statement)
155    }
156}
157
158impl SourcePreparationRequest {
159    /// Execute external reads without access to the catalog.
160    pub async fn execute(
161        &self,
162        storage_configuration: &StorageConfiguration,
163    ) -> SourcePreparationResponse {
164        SourcePreparationResponse {
165            request: self.clone(),
166            result: self.execute_inner(storage_configuration).await,
167        }
168    }
169
170    async fn execute_inner(
171        &self,
172        storage_configuration: &StorageConfiguration,
173    ) -> Result<Arc<PreparationValue>, PlanError> {
174        let value = match &self.kind {
175            RequestKind::Alter(request) => PreparationValue::Alter(Box::new(
176                execute_alter_source_preparation((**request).clone(), storage_configuration)
177                    .await?,
178            )),
179            RequestKind::Source(request) => PreparationValue::Source(Box::new(
180                execute_create_source_preparation((**request).clone(), storage_configuration)
181                    .await?,
182            )),
183            RequestKind::Table(request) => PreparationValue::Table(Box::new(
184                execute_table_from_source_preparation((**request).clone(), storage_configuration)
185                    .await?,
186            )),
187            RequestKind::Protobuf(request) => PreparationValue::Protobuf(
188                (**request).clone().execute(storage_configuration).await?,
189            ),
190            RequestKind::Avro(request) => {
191                PreparationValue::Avro((**request).clone().execute(storage_configuration).await?)
192            }
193            RequestKind::Glue(request) => {
194                PreparationValue::Glue((**request).clone().execute(storage_configuration).await?)
195            }
196        };
197        Ok(Arc::new(value))
198    }
199}
200
201struct Cursor<'a> {
202    preparation: &'a SourcePreparation,
203    next: usize,
204}
205
206impl<'a> Cursor<'a> {
207    fn request(&mut self, kind: RequestKind) -> Result<&'a PreparationValue, ResolveError> {
208        let request = SourcePreparationRequest {
209            position: self.next,
210            kind,
211        };
212        self.next += 1;
213        match self.preparation.inputs.get(&request.position) {
214            Some(input) if input.request == request => Ok(input.value.as_ref()),
215            _ => Err(ResolveError::Pending(request)),
216        }
217    }
218
219    fn resolve(
220        &mut self,
221        catalog: &dyn SessionCatalog,
222        statement: Statement<Aug>,
223    ) -> Result<PurifiedStatement, ResolveError> {
224        match statement {
225            Statement::AlterSource(statement) => {
226                let Some(request) = prepare_alter_source(catalog, &statement)? else {
227                    return Ok(PurifiedStatement::PurifiedAlterSource {
228                        alter_source_stmt: statement,
229                    });
230                };
231                let PreparationValue::Alter(prepared) =
232                    self.request(RequestKind::Alter(Box::new(request)))?
233                else {
234                    unreachable!("the result kind is determined by its request");
235                };
236                Ok((**prepared).clone())
237            }
238            Statement::CreateSource(statement) => {
239                let request = CreateSourcePreparation::new(
240                    catalog,
241                    self.preparation.now,
242                    || *self.preparation.clock_as_of.get_or_init(|| catalog.now()),
243                    statement,
244                )?;
245                let PreparationValue::Source(prepared) =
246                    self.request(RequestKind::Source(Box::new(request)))?
247                else {
248                    unreachable!("the result kind is determined by its request")
249                };
250                let PreparedCreateSource {
251                    create_progress_subsource,
252                    mut create_source_stmt,
253                    subsources,
254                    available_source_references,
255                    format_options,
256                } = (**prepared).clone();
257                self.format(
258                    catalog,
259                    &mut create_source_stmt.format,
260                    &format_options,
261                    &create_source_stmt.envelope,
262                )?;
263                Ok(PurifiedStatement::PurifiedCreateSource {
264                    create_progress_subsource,
265                    create_source_stmt,
266                    subsources,
267                    available_source_references,
268                })
269            }
270            Statement::CreateTableFromSource(mut statement) => {
271                let request = prepare_table_from_source(catalog, &statement)?;
272                let PreparationValue::Table(prepared) =
273                    self.request(RequestKind::Table(Box::new(request)))?
274                else {
275                    unreachable!("the result kind is determined by its request")
276                };
277                let (export, options) = (**prepared).clone();
278                self.format(
279                    catalog,
280                    &mut statement.format,
281                    &options,
282                    &statement.envelope,
283                )?;
284                Ok(finish_table_from_source(catalog, statement, export)?)
285            }
286            _ => Err(internal_err!(
287                "source preparation requires CREATE SOURCE, CREATE TABLE FROM SOURCE, or ALTER SOURCE"
288            )
289            .into()),
290        }
291    }
292
293    fn format(
294        &mut self,
295        catalog: &dyn SessionCatalog,
296        format: &mut Option<FormatSpecifier<Aug>>,
297        options: &SourceFormatOptions,
298        envelope: &Option<SourceEnvelope>,
299    ) -> Result<(), ResolveError> {
300        if matches!(format, Some(FormatSpecifier::KeyValue { .. }))
301            && !matches!(options, SourceFormatOptions::Kafka { .. })
302        {
303            return Err(sql_err!(
304                "Kafka sources are the only source type that can provide KEY/VALUE formats"
305            )
306            .into());
307        }
308        match format {
309            None => {}
310            Some(FormatSpecifier::Bare(format)) => {
311                self.single_format(catalog, format, options, envelope)?;
312            }
313            Some(FormatSpecifier::KeyValue { key, value }) => {
314                self.single_format(catalog, key, options, envelope)?;
315                self.single_format(catalog, value, options, envelope)?;
316            }
317        }
318        Ok(())
319    }
320
321    fn single_format(
322        &mut self,
323        catalog: &dyn SessionCatalog,
324        format: &mut Format<Aug>,
325        options: &SourceFormatOptions,
326        envelope: &Option<SourceEnvelope>,
327    ) -> Result<(), ResolveError> {
328        match format {
329            Format::Avro(AvroSchema::Csr { csr_connection }) => {
330                if let Some(request) =
331                    prepare_csr_connection_avro(catalog, options, csr_connection, envelope)?
332                {
333                    let PreparationValue::Avro(seed) =
334                        self.request(RequestKind::Avro(Box::new(request)))?
335                    else {
336                        unreachable!("the result kind is determined by its request")
337                    };
338                    csr_connection.seed = Some(seed.clone());
339                }
340            }
341            Format::Avro(AvroSchema::Glue {
342                connection,
343                with_options,
344                seed,
345            }) => {
346                if let Some(request) =
347                    prepare_glue_connection_avro(catalog, options, connection, with_options, seed)?
348                {
349                    let PreparationValue::Glue(prepared) =
350                        self.request(RequestKind::Glue(Box::new(request)))?
351                    else {
352                        unreachable!("the result kind is determined by its request")
353                    };
354                    *seed = Some(prepared.clone());
355                }
356            }
357            Format::Protobuf(ProtobufSchema::Csr { csr_connection }) => {
358                if let Some(request) =
359                    prepare_csr_connection_proto(catalog, options, csr_connection, envelope)?
360                {
361                    let PreparationValue::Protobuf(seed) =
362                        self.request(RequestKind::Protobuf(Box::new(request)))?
363                    else {
364                        unreachable!("the result kind is determined by its request")
365                    };
366                    csr_connection.seed = Some(seed.clone());
367                }
368            }
369            Format::Avro(AvroSchema::InlineSchema { .. })
370            | Format::Protobuf(ProtobufSchema::InlineSchema { .. })
371            | Format::Bytes
372            | Format::Regex(_)
373            | Format::Json { .. }
374            | Format::Text
375            | Format::Csv { .. } => {}
376        }
377        Ok(())
378    }
379}
380
381pub(super) async fn purify(
382    catalog: &dyn SessionCatalog,
383    now: u64,
384    statement: Statement<Aug>,
385    storage_configuration: &StorageConfiguration,
386) -> Result<PurifiedStatement, PlanError> {
387    let mut preparation = SourcePreparation::new(now);
388    let mut status = preparation.resolve(catalog, statement.clone())?;
389    loop {
390        match status {
391            SourcePreparationStatus::Ready(statement) => return Ok(*statement),
392            SourcePreparationStatus::Pending(request) => {
393                let response = request.execute(storage_configuration).await;
394                status = preparation.finish(catalog, statement.clone(), response)?;
395            }
396        }
397    }
398}
mz_sql/pure/source_preparation.rs

mz_sql/pure/
source_preparation.rs
