mz_adapter/catalog/command/
source.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
10pub(crate) mod alter;
11
12use std::collections::BTreeMap;
13
14use mz_ore::assert_none;
15use mz_repr::{CatalogItemId, GlobalId, RelationVersionSelector};
16use mz_sql::catalog::SessionCatalog;
17use mz_sql::names::{Aug, ResolvedItemName};
18use mz_sql::plan::{self, CreateSourcePlanBundle, Params, Plan, PlanContext, StatementContext};
19use mz_sql::pure::{
20    PurifiedStatement, generate_progress_subsource_statement, generate_subsource_statements,
21};
22use mz_sql_parser::ast::{
23    CreateSourceOptionName, CreateSubsourceOption, CreateSubsourceOptionName,
24    CreateSubsourceStatement, DeferredItemName, Statement, UnresolvedItemName,
25};
26
27use crate::AdapterError;
28use crate::catalog::ConnCatalog;
29
30/// Stable allocation keys, independent of metadata enumeration order.
31#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
32pub(crate) enum SourceObject {
33    Progress,
34    Primary,
35    Export(UnresolvedItemName),
36}
37
38#[derive(Debug)]
39pub(crate) enum SourcePlanning<T = Vec<CreateSourcePlanBundle>> {
40    NeedsIds(Vec<SourceObject>),
41    Ready(T),
42}
43
44pub(crate) fn plan_create_source(
45    catalog: &ConnCatalog<'_>,
46    pcx: &PlanContext,
47    params: &Params,
48    purified: PurifiedStatement,
49    ids: &BTreeMap<SourceObject, (CatalogItemId, GlobalId)>,
50) -> Result<SourcePlanning, AdapterError> {
51    let PurifiedStatement::PurifiedCreateSource {
52        create_progress_subsource,
53        mut create_source_stmt,
54        subsources,
55        available_source_references,
56    } = purified
57    else {
58        return Err(AdapterError::Unsupported(
59            "non-CREATE SOURCE source planning",
60        ));
61    };
62    let mut create_source_plans = Vec::with_capacity(subsources.len() + 2);
63
64    if create_progress_subsource {
65        let progress_stmt = {
66            let scx = StatementContext::new(None, catalog);
67            generate_progress_subsource_statement(&scx, &create_source_stmt)?
68        };
69        // The primary source references the progress collection, so emit it first.
70        assert_none!(progress_stmt.of_source);
71        let Some(&(item_id, global_id)) = ids.get(&SourceObject::Progress) else {
72            return Ok(SourcePlanning::NeedsIds(vec![SourceObject::Progress]));
73        };
74        let progress_plan =
75            super::plan_subsource(catalog, pcx, params, progress_stmt, item_id, global_id)?;
76        let progress_full_name = catalog.resolve_full_name(&progress_plan.plan.name);
77        let progress_subsource = ResolvedItemName::Item {
78            id: progress_plan.item_id,
79            qualifiers: progress_plan.plan.name.qualifiers.clone(),
80            full_name: progress_full_name,
81            print_id: true,
82            version: RelationVersionSelector::Latest,
83        };
84
85        create_source_plans.push(progress_plan);
86
87        create_source_stmt.progress_subsource = Some(DeferredItemName::Named(progress_subsource));
88    }
89
90    let resolved_ids = mz_sql::names::visit_dependencies(catalog, &create_source_stmt);
91
92    let propagated_with_options: Vec<_> = create_source_stmt
93        .with_options
94        .iter()
95        .filter_map(|opt| match opt.name {
96            CreateSourceOptionName::TimestampInterval => None,
97            CreateSourceOptionName::RetainHistory => Some(CreateSubsourceOption {
98                name: CreateSubsourceOptionName::RetainHistory,
99                value: opt.value.clone(),
100            }),
101        })
102        .collect();
103
104    let source_plan = match mz_sql::plan::plan(
105        Some(pcx),
106        catalog,
107        Statement::CreateSource(create_source_stmt),
108        params,
109        &resolved_ids,
110    )? {
111        (Plan::CreateSource(plan), _sql_impl_ids) => plan,
112        (p, _) => unreachable!("s must be CreateSourcePlan but got {:?}", p),
113    };
114
115    let Some(&(item_id, global_id)) = ids.get(&SourceObject::Primary) else {
116        return Ok(SourcePlanning::NeedsIds(vec![SourceObject::Primary]));
117    };
118
119    let source_full_name = catalog.resolve_full_name(&source_plan.name);
120    let of_source = ResolvedItemName::Item {
121        id: item_id,
122        qualifiers: source_plan.name.qualifiers.clone(),
123        full_name: source_full_name,
124        print_id: true,
125        version: RelationVersionSelector::Latest,
126    };
127
128    let conn_catalog = catalog.for_system_session();
129    let generated_pcx = plan::PlanContext::zero();
130    let scx = StatementContext::new(Some(&generated_pcx), &conn_catalog);
131
132    let mut subsource_stmts = generate_subsource_statements(&scx, of_source, subsources)?;
133
134    for subsource_stmt in subsource_stmts.iter_mut() {
135        subsource_stmt
136            .with_options
137            .extend(propagated_with_options.iter().cloned())
138    }
139
140    create_source_plans.push(CreateSourcePlanBundle {
141        item_id,
142        global_id,
143        plan: source_plan,
144        resolved_ids: resolved_ids.clone(),
145        available_source_references: Some(available_source_references),
146    });
147
148    match plan_source_exports(catalog, pcx, params, subsource_stmts, ids)? {
149        SourcePlanning::NeedsIds(keys) => return Ok(SourcePlanning::NeedsIds(keys)),
150        SourcePlanning::Ready(exports) => create_source_plans.extend(exports),
151    }
152    Ok(SourcePlanning::Ready(create_source_plans))
153}
154
155pub(crate) fn plan_alter_source_add_subsources(
156    catalog: &ConnCatalog<'_>,
157    pcx: &PlanContext,
158    params: &Params,
159    purified: PurifiedStatement,
160    ids: &BTreeMap<SourceObject, (CatalogItemId, GlobalId)>,
161) -> Result<SourcePlanning<Plan>, AdapterError> {
162    let PurifiedStatement::PurifiedAlterSourceAddSubsources {
163        source_name,
164        options,
165        subsources,
166    } = purified
167    else {
168        return Err(AdapterError::Unsupported(
169            "non-ADD SUBSOURCES source planning",
170        ));
171    };
172    let entry = catalog.state().get_entry(source_name.item_id());
173    let source = entry.source().ok_or_else(|| {
174        AdapterError::internal(
175            "plan alter source",
176            format!("expected Source found {entry:?}"),
177        )
178    })?;
179    let item_id = entry.id();
180    let ingestion_id = source.global_id();
181    let system_catalog = catalog.for_system_session();
182    let generated_pcx = PlanContext::zero();
183    let scx = StatementContext::new(Some(&generated_pcx), &system_catalog);
184    let statements = generate_subsource_statements(&scx, source_name, subsources)?;
185    let subsources = match plan_source_exports(catalog, pcx, params, statements, ids)? {
186        SourcePlanning::NeedsIds(keys) => return Ok(SourcePlanning::NeedsIds(keys)),
187        SourcePlanning::Ready(plans) => plans,
188    };
189    Ok(SourcePlanning::Ready(Plan::AlterSource(
190        plan::AlterSourcePlan {
191            item_id,
192            ingestion_id,
193            action: plan::AlterSourceAction::AddSubsourceExports {
194                subsources,
195                options,
196            },
197        },
198    )))
199}
200
201pub(crate) fn plan_alter_source_refresh_references(
202    catalog: &ConnCatalog<'_>,
203    source_name: ResolvedItemName,
204    references: plan::SourceReferences,
205) -> Result<Plan, AdapterError> {
206    let entry = catalog.state().get_entry(source_name.item_id());
207    let source = entry.source().ok_or_else(|| {
208        AdapterError::internal(
209            "plan alter source",
210            format!("expected Source found {entry:?}"),
211        )
212    })?;
213    Ok(Plan::AlterSource(plan::AlterSourcePlan {
214        item_id: entry.id(),
215        ingestion_id: source.global_id(),
216        action: plan::AlterSourceAction::RefreshReferences { references },
217    }))
218}
219
220fn plan_source_exports(
221    catalog: &ConnCatalog<'_>,
222    pcx: &PlanContext,
223    params: &Params,
224    statements: Vec<CreateSubsourceStatement<Aug>>,
225    ids: &BTreeMap<SourceObject, (CatalogItemId, GlobalId)>,
226) -> Result<SourcePlanning, AdapterError> {
227    let missing: Vec<_> = statements
228        .iter()
229        .map(|stmt| SourceObject::Export(stmt.name.clone()))
230        .filter(|key| !ids.contains_key(key))
231        .collect();
232    if !missing.is_empty() {
233        return Ok(SourcePlanning::NeedsIds(missing));
234    }
235    let plans = statements
236        .into_iter()
237        .map(|stmt| {
238            let (item_id, global_id) = ids[&SourceObject::Export(stmt.name.clone())];
239            super::plan_subsource(catalog, pcx, params, stmt, item_id, global_id)
240        })
241        .collect::<Result<_, _>>()?;
242    Ok(SourcePlanning::Ready(plans))
243}
mz_adapter/catalog/command/source.rs

mz_adapter/catalog/command/
source.rs
