mz_adapter/coord/command_handler.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//! Logic for processing client [`Command`]s. Each [`Command`] is initiated by a
11//! client via some external Materialize API (ex: HTTP and psql).
12
13use base64::prelude::*;
14use differential_dataflow::lattice::Lattice;
15use mz_adapter_types::dyncfgs::ALLOW_USER_SESSIONS;
16use mz_auth::AuthenticatorKind;
17use mz_auth::password::Password;
18use mz_repr::namespaces::MZ_INTERNAL_SCHEMA;
19use mz_sql::catalog::AutoProvisionSource;
20use mz_sql::session::metadata::SessionMetadata;
21use std::collections::{BTreeMap, BTreeSet};
22use std::net::IpAddr;
23use std::sync::Arc;
24
25use futures::future::LocalBoxFuture;
26use mz_adapter_types::connection::{ConnectionId, ConnectionIdType};
27use mz_catalog::SYSTEM_CONN_ID;
28use mz_catalog::memory::objects::{
29 CatalogItem, DataSourceDesc, Role, Source, Table, TableDataSource,
30};
31use mz_ore::task;
32use mz_ore::tracing::OpenTelemetryContext;
33use mz_ore::{instrument, soft_panic_or_log};
34use mz_repr::role_id::RoleId;
35use mz_repr::{Diff, GlobalId, SqlScalarType, Timestamp};
36use mz_sql::ast::{
37 AlterConnectionAction, AlterConnectionStatement, AlterSourceAction, AstInfo,
38 CreateSourceOptionName, Raw, Statement, StatementKind,
39};
40use mz_sql::catalog::RoleAttributesRaw;
41use mz_sql::names::{Aug, PartialItemName, ResolvedIds};
42use mz_sql::plan::{
43 AbortTransactionPlan, CommitTransactionPlan, CreateRolePlan, Params, Plan,
44 StatementClassification, TransactionType,
45};
46use mz_sql::pure::{
47 materialized_view_option_contains_temporal, purify_create_materialized_view_options,
48};
49use mz_sql::rbac;
50use mz_sql::session::user::User;
51use mz_sql::session::vars::{
52 EndTransactionAction, NETWORK_POLICY, OwnedVarInput, STATEMENT_LOGGING_SAMPLE_RATE,
53 TRANSACTION_ISOLATION_VAR_NAME, Value, Var, check_transaction_isolation_feature_flag,
54};
55use mz_sql_parser::ast::display::AstDisplay;
56use mz_sql_parser::ast::{
57 CreateMaterializedViewStatement, ExplainPlanStatement, Explainee, WithOptionValue,
58};
59use mz_storage_types::sources::Timeline;
60use opentelemetry::trace::TraceContextExt;
61use tokio::sync::{mpsc, oneshot, watch};
62use tracing::{Instrument, debug_span, info, warn};
63use tracing_opentelemetry::OpenTelemetrySpanExt;
64use uuid::Uuid;
65
66use crate::client::truncate_sql_for_logging;
67use crate::command::{
68 CatalogSnapshot, Command, ExecuteResponse, Response, SASLChallengeResponse,
69 SASLVerifyProofResponse, StartupResponse, SuperuserAttribute,
70};
71use crate::coord::appends::{PendingWriteTxn, UserWriteResponder, WriteResult};
72use crate::coord::peek::PendingPeek;
73use crate::coord::{
74 ConnMeta, Coordinator, DeferredDdl, Message, PendingTxn, PlanValidity, PurifiedStatementReady,
75 validate_ip_with_policy_rules,
76};
77use crate::error::{AdapterError, AuthenticationError};
78use crate::notice::AdapterNotice;
79use crate::session::{Session, TransactionOps, TransactionStatus};
80use crate::statement_logging::{StatementEndedExecutionReason, WatchSetCreation};
81use crate::util::{ClientTransmitter, CompletedClientTransmitter};
82use crate::webhook::{
83 AppendWebhookResponse, AppendWebhookValidator, WebhookAppender, WebhookAppenderInvalidator,
84};
85use crate::{AppendWebhookError, ExecuteContext, catalog, metrics};
86
87use super::{ExecuteContextExtra, ExecuteContextGuard};
88
89/// The login status of a role, used by authentication handlers to check role
90/// existence and login permission before proceeding to credential verification.
91enum RoleLoginStatus {
92 /// The role does not exist in the catalog.
93 NotFound,
94 /// The role exists and has the LOGIN attribute.
95 CanLogin,
96 /// The role exists but does not have the LOGIN attribute.
97 NonLogin,
98}
99
100fn role_login_status(role: Option<&Role>) -> RoleLoginStatus {
101 match role {
102 None => RoleLoginStatus::NotFound,
103 Some(role) => match role.attributes.login {
104 Some(login) if login => RoleLoginStatus::CanLogin,
105 _ => RoleLoginStatus::NonLogin,
106 },
107 }
108}
109
110impl Coordinator {
111 /// Dispatches a command in a size-bounded, heap-allocated future.
112 pub(crate) fn handle_command(&mut self, mut cmd: Command) -> LocalBoxFuture<'_, ()> {
113 let future = async move {
114 if let Some(session) = cmd.session_mut() {
115 session.require_coordinator_completion();
116 session.apply_external_metadata_updates();
117 }
118 match cmd {
119 Command::Startup {
120 tx,
121 user,
122 conn_id,
123 secret_key,
124 uuid,
125 client_ip,
126 application_name,
127 notice_tx,
128 } => {
129 // Note: A ClientTransmitter is not applicable here. Its purpose is to
130 // rescue the Session from an undeliverable Response, and startup
131 // responses carry no Session, the Session stays with the client. An
132 // undeliverable startup response is instead handled by handle_startup's
133 // failed-send path together with the cleanup guard in Client::startup.
134 self.handle_startup(
135 tx,
136 user,
137 conn_id,
138 secret_key,
139 uuid,
140 client_ip,
141 application_name,
142 notice_tx,
143 )
144 .await;
145 }
146
147 Command::AuthenticatePassword {
148 tx,
149 role_name,
150 password,
151 } => {
152 self.handle_authenticate_password(tx, role_name, password)
153 .await;
154 }
155
156 Command::AuthenticateGetSASLChallenge {
157 tx,
158 role_name,
159 nonce,
160 } => {
161 self.handle_generate_sasl_challenge(tx, role_name, nonce)
162 .await;
163 }
164
165 Command::AuthenticateVerifySASLProof {
166 tx,
167 role_name,
168 proof,
169 mock_hash,
170 auth_message,
171 } => {
172 self.handle_authenticate_verify_sasl_proof(
173 tx,
174 role_name,
175 proof,
176 auth_message,
177 mock_hash,
178 );
179 }
180
181 Command::CheckRoleCanLogin { tx, role_name } => {
182 self.handle_role_can_login(tx, role_name);
183 }
184
185 Command::Execute {
186 portal_name,
187 session,
188 tx,
189 outer_ctx_extra,
190 } => {
191 let tx = ClientTransmitter::new(tx, self.internal_cmd_tx.clone());
192
193 Box::pin(self.handle_execute(portal_name, session, tx, outer_ctx_extra)).await;
194 }
195
196 Command::StartCopyFromStdin {
197 target_id,
198 target_name,
199 columns,
200 row_desc,
201 params,
202 session,
203 tx,
204 } => {
205 let otel_ctx = OpenTelemetryContext::obtain();
206 let result = self.setup_copy_from_stdin(
207 &session,
208 target_id,
209 target_name,
210 columns,
211 row_desc,
212 params,
213 );
214 let _ = tx.send(Response {
215 result,
216 session,
217 otel_ctx,
218 });
219 }
220
221 Command::RetireExecute { data, reason } => self.retire_execution(reason, data),
222
223 Command::CancelRequest {
224 conn_id,
225 secret_key,
226 } => {
227 self.handle_cancel(conn_id, secret_key).await;
228 }
229
230 Command::PrivilegedCancelRequest { conn_id } => {
231 self.handle_privileged_cancel(conn_id).await;
232 }
233
234 Command::GetWebhook {
235 database,
236 schema,
237 name,
238 tx,
239 } => {
240 self.handle_get_webhook(database, schema, name, tx);
241 }
242
243 Command::GetSystemVars { tx } => {
244 let _ = tx.send(self.catalog.system_config().clone());
245 }
246
247 Command::SetSystemVars { vars, conn_id, tx } => {
248 let mut ops = Vec::with_capacity(vars.len());
249 let conn = &self.active_conns[&conn_id];
250
251 for (name, value) in vars {
252 if let Err(e) =
253 self.catalog().system_config().get(&name).and_then(|var| {
254 var.visible(conn.user(), self.catalog.system_config())
255 })
256 {
257 let _ = tx.send(Err(e.into()));
258 return;
259 }
260
261 ops.push(catalog::Op::UpdateSystemConfiguration {
262 name,
263 value: OwnedVarInput::Flat(value),
264 });
265 }
266
267 let result = self
268 .catalog_transact_with_context(Some(&conn_id), None, ops)
269 .await;
270 let _ = tx.send(result);
271 }
272
273 Command::UpdateScopedSystemParameters {
274 overrides,
275 prune_scope,
276 tx,
277 } => {
278 // Store the new working copy, persist it durably, and
279 // reconcile it into the per-scope resolution boundaries.
280 self.reconcile_scoped_system_parameters(overrides, prune_scope)
281 .await;
282 let _ = tx.send(());
283 }
284
285 Command::InstallScopedSystemParameterFrontend { frontend } => {
286 // Keep the shared frontend so create-cluster / create-replica
287 // can resolve scoped overrides synchronously at create time.
288 self.scoped_frontend = Some(frontend);
289 }
290
291 Command::InjectAuditEvents {
292 events,
293 conn_id,
294 tx,
295 } => {
296 let ops = vec![catalog::Op::InjectAuditEvents { events }];
297 let result = self
298 .catalog_transact_with_context(Some(&conn_id), None, ops)
299 .await;
300 let _ = tx.send(result);
301 }
302
303 Command::Terminate { conn_id, tx } => {
304 self.handle_terminate(conn_id).await;
305 // Note: We purposefully do not use a ClientTransmitter here because we're already
306 // terminating the provided session.
307 if let Some(tx) = tx {
308 let _ = tx.send(Ok(()));
309 }
310 }
311
312 Command::Commit {
313 action,
314 session,
315 tx,
316 } => {
317 let tx = ClientTransmitter::new(tx, self.internal_cmd_tx.clone());
318 // We reach here not through a statement execution, but from the
319 // "commit" pgwire command. Thus, we just generate a default statement
320 // execution context (once statement logging is implemented, this will cause nothing to be logged
321 // when the execution finishes.)
322 let ctx = ExecuteContext::from_parts(
323 tx,
324 self.internal_cmd_tx.clone(),
325 session,
326 Default::default(),
327 );
328 let plan = match action {
329 EndTransactionAction::Commit => {
330 Plan::CommitTransaction(CommitTransactionPlan {
331 transaction_type: TransactionType::Implicit,
332 })
333 }
334 EndTransactionAction::Rollback => {
335 Plan::AbortTransaction(AbortTransactionPlan {
336 transaction_type: TransactionType::Implicit,
337 })
338 }
339 };
340
341 let conn_id = ctx.session().conn_id().clone();
342 self.sequence_plan(ctx, plan, ResolvedIds::empty(), ResolvedIds::empty())
343 .await;
344 // Part of the Command::Commit contract is that the Coordinator guarantees that
345 // it has cleared its transaction state for the connection.
346 let retire_notify = self.clear_connection(&conn_id).await;
347 // `sequence_plan` has already handled the client response.
348 // This call only satisfies the internal cleanup contract.
349 drop(retire_notify);
350 }
351
352 Command::CatalogSnapshot { tx } => {
353 let _ = tx.send(CatalogSnapshot {
354 catalog: self.owned_catalog(),
355 });
356 }
357
358 Command::CheckConsistency { tx } => {
359 let _ = tx.send(self.check_consistency());
360 }
361
362 Command::Dump { tx } => {
363 let _ = tx.send(self.dump().await);
364 }
365
366 Command::GetComputeInstanceClient { instance_id, tx } => {
367 let _ = tx.send(self.controller.compute.instance_client(instance_id));
368 }
369
370 Command::GetOracle { timeline, tx } => {
371 let oracle = self
372 .global_timelines
373 .get(&timeline)
374 .map(|timeline_state| Arc::clone(&timeline_state.oracle))
375 .ok_or(AdapterError::ChangedPlan(
376 "timeline has disappeared during planning".to_string(),
377 ));
378 let _ = tx.send(oracle);
379 }
380
381 Command::DetermineRealTimeRecentTimestamp {
382 source_ids,
383 real_time_recency_timeout,
384 tx,
385 } => {
386 let result = self
387 .determine_real_time_recent_timestamp(
388 source_ids.iter().copied(),
389 real_time_recency_timeout,
390 )
391 .await;
392
393 match result {
394 Ok(Some(fut)) => {
395 let catalog = Arc::clone(&self.catalog);
396 task::spawn(|| "determine real time recent timestamp", async move {
397 let result =
398 Coordinator::await_real_time_recent_timestamp(catalog, fut)
399 .await
400 .map(Some);
401 let _ = tx.send(result);
402 });
403 }
404 Ok(None) => {
405 let _ = tx.send(Ok(None));
406 }
407 Err(e) => {
408 let _ = tx.send(Err(e));
409 }
410 }
411 }
412
413 Command::GetTransactionReadHoldsBundle { conn_id, tx } => {
414 let read_holds = self.txn_read_holds.get(&conn_id).cloned();
415 let _ = tx.send(read_holds);
416 }
417
418 Command::StoreTransactionReadHolds {
419 conn_id,
420 read_holds,
421 tx,
422 } => {
423 self.store_transaction_read_holds(conn_id, read_holds);
424 let _ = tx.send(());
425 }
426
427 Command::ExecuteSlowPathPeek {
428 dataflow_plan,
429 determination,
430 finishing,
431 compute_instance,
432 target_replica,
433 intermediate_result_type,
434 source_ids,
435 conn_id,
436 max_result_size,
437 max_query_result_size,
438 watch_set,
439 tx,
440 } => {
441 let result = self.implement_slow_path_peek(
442 *dataflow_plan,
443 determination,
444 finishing,
445 compute_instance,
446 target_replica,
447 intermediate_result_type,
448 source_ids,
449 conn_id,
450 max_result_size,
451 max_query_result_size,
452 watch_set,
453 );
454 let _ = tx.send(result);
455 }
456
457 Command::ExecuteSubscribe {
458 df_desc,
459 dependency_ids,
460 cluster_id,
461 replica_id,
462 conn_id,
463 session_uuid,
464 read_holds,
465 plan,
466 statement_logging_id,
467 tx,
468 } => {
469 match self
470 .implement_subscribe(
471 statement_logging_id,
472 df_desc,
473 dependency_ids,
474 cluster_id,
475 replica_id,
476 conn_id,
477 session_uuid,
478 read_holds,
479 plan,
480 )
481 .await
482 {
483 Ok((resp, write_notify)) => {
484 // Wait for the `mz_subscriptions` bookkeeping write off the
485 // coordinator loop before returning the `SUBSCRIBE` response to
486 // the subscribing session.
487 task::spawn(|| "execute_subscribe::await_bookkeeping", async move {
488 write_notify.await;
489 let _ = tx.send(Ok(resp));
490 });
491 }
492 Err(e) => {
493 let _ = tx.send(Err(e));
494 }
495 }
496 }
497
498 Command::CopyToPreflight {
499 s3_sink_connection,
500 sink_id,
501 tx,
502 } => {
503 // Spawn a background task to perform the slow S3 preflight operations.
504 // This avoids blocking the coordinator's main task.
505 let connection_context = self.connection_context().clone();
506 let enforce_external_addresses =
507 mz_storage_types::dyncfgs::ENFORCE_EXTERNAL_ADDRESSES
508 .get(self.controller.storage.config().config_set());
509 task::spawn(|| "copy_to_preflight", async move {
510 let result = mz_storage_types::sinks::s3_oneshot_sink::preflight(
511 connection_context,
512 &s3_sink_connection.aws_connection,
513 &s3_sink_connection.upload_info,
514 s3_sink_connection.connection_id,
515 sink_id,
516 enforce_external_addresses,
517 )
518 .await
519 .map_err(AdapterError::from);
520 let _ = tx.send(result);
521 });
522 }
523
524 Command::ExecuteCopyTo {
525 df_desc,
526 compute_instance,
527 target_replica,
528 source_ids,
529 conn_id,
530 watch_set,
531 tx,
532 } => {
533 // implement_copy_to spawns a background task that sends the response
534 // through tx when the COPY TO completes (or immediately if setup fails).
535 // We just call it and let it handle all response sending.
536 self.implement_copy_to(
537 *df_desc,
538 compute_instance,
539 target_replica,
540 source_ids,
541 conn_id,
542 watch_set,
543 tx,
544 )
545 .await;
546 }
547
548 Command::ExecuteSideEffectingFunc { plan, conn_id, tx } => {
549 let result = self.execute_side_effecting_func(plan, conn_id).await;
550 let _ = tx.send(result);
551 }
552 Command::LookupConnection { connection_id, tx } => {
553 let conn =
554 self.active_conns
555 .get_key_value(&connection_id)
556 .map(|(id_handle, meta)| {
557 (id_handle.clone(), *meta.authenticated_role_id())
558 });
559 let _ = tx.send(conn);
560 }
561 Command::RegisterPeek {
562 uuid,
563 conn_id,
564 cluster_id,
565 depends_on,
566 is_fast_path,
567 watch_set,
568 tx,
569 } => {
570 self.handle_register_peek(
571 uuid,
572 conn_id,
573 cluster_id,
574 depends_on,
575 is_fast_path,
576 watch_set,
577 tx,
578 );
579 }
580 Command::UnregisterPeek { uuid, reason, tx } => {
581 self.handle_unregister_peek(uuid, reason, tx);
582 }
583 Command::ExplainTimestamp {
584 conn_id,
585 session_wall_time,
586 cluster_id,
587 id_bundle,
588 determination,
589 tx,
590 } => {
591 let explanation = self.explain_timestamp(
592 &conn_id,
593 session_wall_time,
594 cluster_id,
595 &id_bundle,
596 determination,
597 );
598 let _ = tx.send(explanation);
599 }
600 Command::FrontendStatementLogging(event) => {
601 self.handle_frontend_statement_logging_event(event);
602 }
603 Command::RegisterConnectionCancelWatch { conn_id, tx } => {
604 // Always replace any existing entry. Another code path
605 // (e.g. `sequence_staged`) may have left a stale watch
606 // here, possibly already signaled `true` from a prior
607 // cancel. Reusing it via `or_insert_with` would hand out
608 // a `Receiver` that already reads `true`, causing the new
609 // operation to immediately return `Canceled` even though
610 // it hasn't been cancelled.
611 let (watch_tx, watch_rx) = watch::channel(false);
612 self.connection_cancel_watches
613 .insert(conn_id, (watch_tx, watch_rx.clone()));
614 let _ = tx.send(watch_rx);
615 }
616 Command::CreateInternalSubscribe {
617 df_desc,
618 cluster_id,
619 replica_id,
620 depends_on,
621 as_of,
622 arity,
623 sink_id,
624 owner,
625 start_time,
626 read_holds,
627 tx,
628 } => {
629 self.handle_create_internal_subscribe(
630 *df_desc, cluster_id, replica_id, depends_on, as_of, arity, sink_id, owner,
631 start_time, read_holds, tx,
632 )
633 .await;
634 }
635 Command::AttemptWrite {
636 attempt,
637 target_id,
638 target_global_id,
639 diffs,
640 tx,
641 } => {
642 self.handle_attempt_write(attempt, target_id, target_global_id, diffs, tx);
643 }
644 Command::DropInternalSubscribe { sink_id } => {
645 self.drop_internal_subscribe(sink_id).await;
646 }
647 }
648 }
649 .instrument(debug_span!("handle_command"));
650 crate::coord::box_dispatcher(future)
651 }
652
653 fn handle_role_can_login(
654 &self,
655 tx: oneshot::Sender<Result<(), AdapterError>>,
656 role_name: String,
657 ) {
658 let result =
659 match role_login_status(self.catalog().try_get_role_by_name(role_name.as_str())) {
660 RoleLoginStatus::NotFound => Err(AdapterError::AuthenticationError(
661 AuthenticationError::RoleNotFound,
662 )),
663 RoleLoginStatus::NonLogin => Err(AdapterError::AuthenticationError(
664 AuthenticationError::NonLogin,
665 )),
666 RoleLoginStatus::CanLogin => Ok(()),
667 };
668 let _ = tx.send(result);
669 }
670
671 fn handle_authenticate_verify_sasl_proof(
672 &self,
673 tx: oneshot::Sender<Result<SASLVerifyProofResponse, AdapterError>>,
674 role_name: String,
675 proof: String,
676 auth_message: String,
677 mock_hash: String,
678 ) {
679 let role = self.catalog().try_get_role_by_name(role_name.as_str());
680 let login_status = role_login_status(role);
681 let role_auth = role.and_then(|r| self.catalog().try_get_role_auth_by_id(&r.id));
682 let real_hash = role_auth
683 .as_ref()
684 .and_then(|auth| auth.password_hash.as_ref());
685 let hash_ref = real_hash.map(|s| s.as_str()).unwrap_or(&mock_hash);
686
687 match mz_auth::hash::sasl_verify(hash_ref, &proof, &auth_message) {
688 Ok(verifier) => {
689 // Success only if role exists, allows login, and a real password hash was used.
690 if matches!(login_status, RoleLoginStatus::CanLogin) && real_hash.is_some() {
691 let _ = tx.send(Ok(SASLVerifyProofResponse { verifier }));
692 } else {
693 let _ = tx.send(Err(AdapterError::AuthenticationError(match login_status {
694 RoleLoginStatus::NonLogin => AuthenticationError::NonLogin,
695 RoleLoginStatus::NotFound => AuthenticationError::RoleNotFound,
696 RoleLoginStatus::CanLogin => AuthenticationError::InvalidCredentials,
697 })));
698 }
699 }
700 Err(_) => {
701 let _ = tx.send(Err(AdapterError::AuthenticationError(
702 AuthenticationError::InvalidCredentials,
703 )));
704 }
705 }
706 }
707
708 #[mz_ore::instrument(level = "debug")]
709 async fn handle_generate_sasl_challenge(
710 &self,
711 tx: oneshot::Sender<Result<SASLChallengeResponse, AdapterError>>,
712 role_name: String,
713 client_nonce: String,
714 ) {
715 let role_auth = self
716 .catalog()
717 .try_get_role_by_name(&role_name)
718 .and_then(|role| self.catalog().try_get_role_auth_by_id(&role.id));
719
720 let nonce = match mz_auth::hash::generate_nonce(&client_nonce) {
721 Ok(n) => n,
722 Err(e) => {
723 let msg = format!(
724 "failed to generate nonce for client nonce {}: {}",
725 client_nonce, e
726 );
727 let _ = tx.send(Err(AdapterError::Internal(msg.clone())));
728 soft_panic_or_log!("{msg}");
729 return;
730 }
731 };
732
733 // It's important that the mock_nonce is deterministic per role, otherwise the purpose of
734 // doing mock authentication is defeated. We use a catalog-wide nonce, and combine that
735 // with the role name to get a per-role mock nonce.
736 let send_mock_challenge =
737 |role_name: String,
738 mock_nonce: String,
739 nonce: String,
740 tx: oneshot::Sender<Result<SASLChallengeResponse, AdapterError>>| {
741 let opts = mz_auth::hash::mock_sasl_challenge(
742 &role_name,
743 &mock_nonce,
744 &self.catalog().system_config().scram_iterations(),
745 );
746 let _ = tx.send(Ok(SASLChallengeResponse {
747 iteration_count: mz_ore::cast::u32_to_usize(opts.iterations.get()),
748 salt: BASE64_STANDARD.encode(opts.salt),
749 nonce,
750 }));
751 };
752
753 match role_auth {
754 Some(auth) if auth.password_hash.is_some() => {
755 let hash = auth.password_hash.as_ref().expect("checked above");
756 match mz_auth::hash::scram256_parse_opts(hash) {
757 Ok(opts) => {
758 let _ = tx.send(Ok(SASLChallengeResponse {
759 iteration_count: mz_ore::cast::u32_to_usize(opts.iterations.get()),
760 salt: BASE64_STANDARD.encode(opts.salt),
761 nonce,
762 }));
763 }
764 Err(_) => {
765 send_mock_challenge(
766 role_name,
767 self.catalog().state().mock_authentication_nonce(),
768 nonce,
769 tx,
770 );
771 }
772 }
773 }
774 _ => {
775 send_mock_challenge(
776 role_name,
777 self.catalog().state().mock_authentication_nonce(),
778 nonce,
779 tx,
780 );
781 }
782 }
783 }
784
785 #[mz_ore::instrument(level = "debug")]
786 async fn handle_authenticate_password(
787 &self,
788 tx: oneshot::Sender<Result<(), AdapterError>>,
789 role_name: String,
790 password: Option<Password>,
791 ) {
792 let Some(password) = password else {
793 // The user did not provide a password.
794 let _ = tx.send(Err(AdapterError::AuthenticationError(
795 AuthenticationError::PasswordRequired,
796 )));
797 return;
798 };
799 let role = self.catalog().try_get_role_by_name(role_name.as_str());
800
801 match role_login_status(role) {
802 RoleLoginStatus::NotFound => {
803 let _ = tx.send(Err(AdapterError::AuthenticationError(
804 AuthenticationError::RoleNotFound,
805 )));
806 return;
807 }
808 RoleLoginStatus::NonLogin => {
809 let _ = tx.send(Err(AdapterError::AuthenticationError(
810 AuthenticationError::NonLogin,
811 )));
812 return;
813 }
814 RoleLoginStatus::CanLogin => {}
815 }
816
817 let role_auth = role.and_then(|r| self.catalog().try_get_role_auth_by_id(&r.id));
818
819 if let Some(auth) = role_auth {
820 if let Some(hash) = &auth.password_hash {
821 let hash = hash.clone();
822 task::spawn_blocking(
823 || "auth-check-hash",
824 move || {
825 let _ = match mz_auth::hash::scram256_verify(&password, &hash) {
826 Ok(_) => tx.send(Ok(())),
827 Err(_) => tx.send(Err(AdapterError::AuthenticationError(
828 AuthenticationError::InvalidCredentials,
829 ))),
830 };
831 },
832 );
833 return;
834 }
835 }
836 // Authentication failed due to missing password hash.
837 let _ = tx.send(Err(AdapterError::AuthenticationError(
838 AuthenticationError::InvalidCredentials,
839 )));
840 }
841
842 #[mz_ore::instrument(level = "debug")]
843 async fn handle_startup(
844 &mut self,
845 tx: oneshot::Sender<Result<StartupResponse, AdapterError>>,
846 user: User,
847 conn_id: ConnectionId,
848 secret_key: u32,
849 uuid: uuid::Uuid,
850 client_ip: Option<IpAddr>,
851 application_name: String,
852 notice_tx: mpsc::UnboundedSender<AdapterNotice>,
853 ) {
854 // Early return if successful, otherwise cleanup any possible state.
855 match self
856 .handle_startup_inner(&user, &conn_id, &client_ip, ¬ice_tx)
857 .await
858 {
859 Ok((role_id, superuser_attribute, session_defaults)) => {
860 let (frontend_cancel_tx, frontend_cancel_rx) = watch::channel(());
861 let session_type = metrics::session_type_label_value(&user);
862 self.metrics
863 .active_sessions
864 .with_label_values(&[session_type])
865 .inc();
866 let conn = ConnMeta {
867 frontend_cancel_tx,
868 secret_key,
869 notice_tx,
870 drop_sinks: BTreeSet::new(),
871 connected_at: self.now(),
872 user,
873 application_name,
874 uuid,
875 client_ip,
876 conn_id: conn_id.clone(),
877 authenticated_role: role_id,
878 deferred_lock: None,
879 };
880 let update = self.catalog().state().pack_session_update(&conn, Diff::ONE);
881 let update = self.catalog().state().resolve_builtin_table_update(update);
882 self.begin_session_for_statement_logging(&conn);
883 self.active_conns.insert(conn_id.clone(), conn);
884
885 // Note: Do NOT await the notify here, we pass this back to
886 // whatever requested the startup to prevent blocking startup
887 // and the Coordinator on a builtin table update.
888 let updates = vec![update];
889 // It's not a hard error if our list is missing a builtin table, but we want to
890 // make sure these two things stay in-sync.
891 if mz_ore::assert::soft_assertions_enabled() {
892 let required_tables: BTreeSet<_> = super::appends::REQUIRED_BUILTIN_TABLES
893 .iter()
894 .map(|table| self.catalog().resolve_builtin_table(*table))
895 .collect();
896 let updates_tracked = updates
897 .iter()
898 .all(|update| required_tables.contains(&update.id));
899 let all_mz_internal = super::appends::REQUIRED_BUILTIN_TABLES
900 .iter()
901 .all(|table| table.schema == MZ_INTERNAL_SCHEMA);
902 mz_ore::soft_assert_or_log!(
903 updates_tracked,
904 "not tracking all required builtin table updates!"
905 );
906 // TODO(parkmycar): When checking if a query depends on these builtin table
907 // writes we do not check the transitive dependencies of the query, because
908 // we don't support creating views on mz_internal objects. If one of these
909 // tables is promoted out of mz_internal then we'll need to add this check.
910 mz_ore::soft_assert_or_log!(
911 all_mz_internal,
912 "not all builtin tables are in mz_internal! need to check transitive depends",
913 )
914 }
915 let notify = self.builtin_table_update().background(updates);
916
917 let catalog = self.owned_catalog();
918 let build_info_human_version =
919 catalog.state().config().build_info.human_version(None);
920
921 let statement_logging_frontend = self
922 .statement_logging
923 .create_frontend(build_info_human_version);
924
925 let resp = Ok(StartupResponse {
926 frontend_cancel_rx,
927 role_id,
928 write_notify: notify,
929 session_defaults,
930 catalog,
931 storage_collections: Arc::clone(&self.controller.storage_collections),
932 transient_id_gen: Arc::clone(&self.transient_id_gen),
933 optimizer_metrics: self.optimizer_metrics.clone(),
934 persist_client: self.persist_client.clone(),
935 statement_logging_frontend,
936 superuser_attribute,
937 occ_write_semaphore: Arc::clone(&self.occ_write_semaphore),
938 group_commit_notifier: self.group_commit_tx.clone(),
939 read_only: self.controller.read_only(),
940 });
941 if tx.send(resp).is_err() {
942 // Failed to send to adapter, but everything is setup so we can terminate
943 // normally.
944 self.handle_terminate(conn_id).await;
945 }
946 }
947 Err(e) => {
948 // Nothing to clean up and `handle_terminate` must not be called. Per the
949 // `handle_startup_inner` invariant, no per-connection state exists and the
950 // connection was never registered in `active_conns`. An auto-provisioned role
951 // stays, it is role-scoped rather than tied to this connection. Temporary
952 // schemas are created lazily, so none exists yet. For the same reason, a
953 // failure to send the error back to the client needs no handling.
954 let _ = tx.send(Err(e));
955 }
956 }
957 }
958
959 /// Fallible startup work.
960 ///
961 /// Invariant: when this returns an error, no per-connection coordinator
962 /// state exists. Per-connection state is registered only in
963 /// `handle_startup`'s Ok arm, after this function has succeeded. The
964 /// cleanup guard in `Client::startup` relies on this invariant by not
965 /// sending `Terminate` when startup fails, and `handle_terminate` panics
966 /// on connections it does not know about.
967 ///
968 /// Durable catalog changes made before an error (an auto-provisioned
969 /// role, synced role memberships) are role-scoped rather than
970 /// connection-scoped and are intentionally kept.
971 async fn handle_startup_inner(
972 &mut self,
973 user: &User,
974 _conn_id: &ConnectionId,
975 client_ip: &Option<IpAddr>,
976 notice_tx: &mpsc::UnboundedSender<AdapterNotice>,
977 ) -> Result<(RoleId, SuperuserAttribute, BTreeMap<String, OwnedVarInput>), AdapterError> {
978 if self.catalog().try_get_role_by_name(&user.name).is_none() {
979 // If the user has made it to this point, that means they have been fully authenticated.
980 // This includes preventing any user, except a pre-defined set of system users, from
981 // connecting to an internal port. Therefore it's ok to always create a new role for the
982 // user.
983 let mut attributes = RoleAttributesRaw::new();
984 // When auto-provisioning, we store the authenticator that was used to provision the role.
985 attributes.auto_provision_source = match user.authenticator_kind {
986 Some(AuthenticatorKind::Oidc) => Some(AutoProvisionSource::Oidc),
987 Some(AuthenticatorKind::Frontegg) => Some(AutoProvisionSource::Frontegg),
988 Some(AuthenticatorKind::None) => Some(AutoProvisionSource::None),
989 _ => {
990 warn!(
991 "auto-provisioning role with unexpected authenticator kind: {:?}",
992 user.authenticator_kind
993 );
994 None
995 }
996 };
997
998 // Auto-provision roles with the LOGIN attribute to distinguish
999 // them as users.
1000 attributes.login = Some(true);
1001
1002 let plan = CreateRolePlan {
1003 name: user.name.to_string(),
1004 attributes,
1005 };
1006 self.sequence_create_role_for_startup(plan).await?;
1007 }
1008 let role = self
1009 .catalog()
1010 .try_get_role_by_name(&user.name)
1011 .expect("created above");
1012 let role_id = role.id;
1013 let superuser_attribute = role.attributes.superuser;
1014
1015 // JWT group-to-role sync: reconcile role memberships with JWT group claims.
1016 // Missing groups claim (None) → skip sync; empty (Some([])) → revoke all.
1017 self.maybe_sync_jwt_groups(role_id, user.groups.as_deref(), notice_tx)
1018 .await?;
1019
1020 if role_id.is_user() && !ALLOW_USER_SESSIONS.get(self.catalog().system_config().dyncfgs()) {
1021 return Err(AdapterError::UserSessionsDisallowed);
1022 }
1023
1024 // Initialize the default session variables for this role.
1025 let mut session_defaults = BTreeMap::new();
1026 let system_config = self.catalog().state().system_config();
1027
1028 // Override the session with any system defaults.
1029 session_defaults.extend(
1030 system_config
1031 .iter_session()
1032 .map(|v| (v.name().to_string(), OwnedVarInput::Flat(v.value()))),
1033 );
1034 // Special case.
1035 let statement_logging_default = system_config
1036 .statement_logging_default_sample_rate()
1037 .format();
1038 session_defaults.insert(
1039 STATEMENT_LOGGING_SAMPLE_RATE.name().to_string(),
1040 OwnedVarInput::Flat(statement_logging_default),
1041 );
1042 // Override system defaults with role defaults.
1043 session_defaults.extend(
1044 self.catalog()
1045 .get_role(&role_id)
1046 .vars()
1047 .map(|(name, val)| (name.to_string(), val.clone())),
1048 );
1049
1050 // If the resolved `transaction_isolation` default names a feature-flagged
1051 // isolation level whose flag is now disabled (e.g. a role default set
1052 // while `strong session serializable` was enabled, then the flag turned
1053 // off), drop it so the session falls back to the built-in default rather
1054 // than silently using a gated level.
1055 if let Some(value) = session_defaults.get(TRANSACTION_ISOLATION_VAR_NAME) {
1056 if check_transaction_isolation_feature_flag(
1057 TRANSACTION_ISOLATION_VAR_NAME,
1058 value.borrow(),
1059 system_config,
1060 )
1061 .is_err()
1062 {
1063 session_defaults.remove(TRANSACTION_ISOLATION_VAR_NAME);
1064 }
1065 }
1066
1067 // Validate network policies for external users. Internal users can only connect on the
1068 // internal interfaces (internal HTTP/ pgwire). It is up to the person deploying the system
1069 // to ensure these internal interfaces are well secured.
1070 //
1071 // HACKY(parkmycar): We don't have a fully formed session yet for this role, but we want
1072 // the default network policy for this role, so we read directly out of what the session
1073 // will get initialized with.
1074 if !user.is_internal() {
1075 let network_policy_name = session_defaults
1076 .get(NETWORK_POLICY.name())
1077 .and_then(|value| match value {
1078 OwnedVarInput::Flat(name) => Some(name.clone()),
1079 OwnedVarInput::SqlSet(names) => {
1080 tracing::error!(?names, "found multiple network policies");
1081 None
1082 }
1083 })
1084 .unwrap_or_else(|| system_config.default_network_policy_name());
1085 let maybe_network_policy = self
1086 .catalog()
1087 .get_network_policy_by_name(&network_policy_name);
1088
1089 let Some(network_policy) = maybe_network_policy else {
1090 // We should prevent dropping the default network policy, or setting the policy
1091 // to something that doesn't exist, so complain loudly if this occurs.
1092 tracing::error!(
1093 network_policy_name,
1094 "default network policy does not exist. All user traffic will be blocked"
1095 );
1096 let reason = match client_ip {
1097 Some(ip) => super::NetworkPolicyError::AddressDenied(ip.clone()),
1098 None => super::NetworkPolicyError::MissingIp,
1099 };
1100 return Err(AdapterError::NetworkPolicyDenied(reason));
1101 };
1102
1103 if let Some(ip) = client_ip {
1104 match validate_ip_with_policy_rules(ip, &network_policy.rules) {
1105 Ok(_) => {}
1106 Err(e) => return Err(AdapterError::NetworkPolicyDenied(e)),
1107 }
1108 } else {
1109 // Only temporary and internal representation of a session
1110 // should be missing a client_ip. These sessions should not be
1111 // making requests or going through handle_startup.
1112 return Err(AdapterError::NetworkPolicyDenied(
1113 super::NetworkPolicyError::MissingIp,
1114 ));
1115 }
1116 }
1117
1118 // Temporary schemas are now created lazily when the first temporary object is created,
1119 // rather than eagerly on connection startup. This avoids expensive catalog_mut() calls
1120 // for the common case where connections never create temporary objects.
1121
1122 Ok((
1123 role_id,
1124 SuperuserAttribute(superuser_attribute),
1125 session_defaults,
1126 ))
1127 }
1128
1129 /// Handles an execute command.
1130 #[instrument(name = "coord::handle_execute", fields(session = session.uuid().to_string()))]
1131 pub(crate) async fn handle_execute(
1132 &mut self,
1133 portal_name: String,
1134 mut session: Session,
1135 tx: ClientTransmitter<ExecuteResponse>,
1136 // If this command was part of another execute command
1137 // (for example, executing a `FETCH` statement causes an execute to be
1138 // issued for the cursor it references),
1139 // then `outer_context` should be `Some`.
1140 // This instructs the coordinator that the
1141 // outer execute should be considered finished once the inner one is.
1142 outer_context: Option<ExecuteContextExtra>,
1143 ) {
1144 // Arm the obligation before anything can return: from here on, every
1145 // exit either retires it or hands it to a guard whose `Drop` does.
1146 // Nothing could have dropped it on the way here either, the command
1147 // travels a channel that only this loop drains.
1148 let outer_context = outer_context
1149 .map(|extra| ExecuteContextGuard::new(extra.retire(), self.internal_cmd_tx.clone()));
1150
1151 // A new statement is starting, so discard any cancellation that was signaled while no
1152 // statement was running. Such a cancellation targeted an earlier statement and must not
1153 // cancel the new one. (Like in PostgreSQL, a cancel request that arrives when nothing is
1154 // running has no effect.) The watch would otherwise retain a stale `true` within an
1155 // explicit transaction, because it is removed only when the transaction is cleared, not
1156 // at statement end.
1157 //
1158 // Don't do this for nested executes (e.g., FETCH executing its cursor's statement): the
1159 // outer statement is still running and a pending cancellation may target it.
1160 if outer_context.is_none() {
1161 self.connection_cancel_watches.remove(session.conn_id());
1162 }
1163
1164 if session.vars().emit_trace_id_notice() {
1165 let span_context = tracing::Span::current()
1166 .context()
1167 .span()
1168 .span_context()
1169 .clone();
1170 if span_context.is_valid() {
1171 session.add_notice(AdapterNotice::QueryTrace {
1172 trace_id: span_context.trace_id(),
1173 });
1174 }
1175 }
1176
1177 if let Err(err) = Self::verify_portal(self.catalog(), &mut session, &portal_name) {
1178 // If statement logging hasn't started yet, we don't need
1179 // to add any "end" event, so just make up a no-op
1180 // `ExecuteContextExtra` here, via `Default::default`.
1181 //
1182 // It's a bit unfortunate because the edge case of failed
1183 // portal verifications won't show up in statement
1184 // logging, but there seems to be nothing else we can do,
1185 // because we need access to the portal to begin logging.
1186 //
1187 // Another option would be to log a begin and end event, but just fill in NULLs
1188 // for everything we get from the portal (prepared statement id, params).
1189 let extra = outer_context.unwrap_or_else(Default::default);
1190 let ctx = ExecuteContext::from_parts(tx, self.internal_cmd_tx.clone(), session, extra);
1191 return ctx.retire(Err(err));
1192 }
1193
1194 // The reference to `portal` can't outlive `session`, which we
1195 // use to construct the context, so scope the reference to this block where we
1196 // get everything we need from the portal for later.
1197 let (stmt, ctx, params) = {
1198 let portal = session
1199 .get_portal_unverified(&portal_name)
1200 .expect("known to exist");
1201 let params = portal.parameters.clone();
1202 let stmt = portal.stmt.clone();
1203 let logging = Arc::clone(&portal.logging);
1204 let lifecycle_timestamps = portal.lifecycle_timestamps.clone();
1205
1206 let extra = if let Some(extra) = outer_context {
1207 // We are executing in the context of another SQL statement, so we don't
1208 // want to begin statement logging anew. The context of the actual statement
1209 // being executed is the one that should be retired once this finishes.
1210 extra
1211 } else {
1212 // This is a new statement, log it and return the context
1213 let maybe_uuid = self.begin_statement_execution(
1214 &mut session,
1215 ¶ms,
1216 &logging,
1217 lifecycle_timestamps,
1218 );
1219
1220 ExecuteContextGuard::new(maybe_uuid, self.internal_cmd_tx.clone())
1221 };
1222 let ctx = ExecuteContext::from_parts(tx, self.internal_cmd_tx.clone(), session, extra);
1223 (stmt, ctx, params)
1224 };
1225
1226 let stmt = match stmt {
1227 Some(stmt) => stmt,
1228 None => return ctx.retire(Ok(ExecuteResponse::EmptyQuery)),
1229 };
1230
1231 let session_type = metrics::session_type_label_value(ctx.session().user());
1232 let stmt_type = metrics::statement_type_label_value(&stmt);
1233 self.metrics
1234 .query_total
1235 .with_label_values(&[session_type, stmt_type])
1236 .inc();
1237 self.handle_execute_inner(stmt, params, ctx).await
1238 }
1239
1240 #[instrument(
1241 name = "coord::handle_execute_inner",
1242 fields(stmt = truncate_sql_for_logging(stmt.to_ast_string_redacted())),
1243 )]
1244 pub(crate) async fn handle_execute_inner(
1245 &mut self,
1246 stmt: Arc<Statement<Raw>>,
1247 params: Params,
1248 mut ctx: ExecuteContext,
1249 ) {
1250 // This comment describes the various ways DDL can execute (the ordered operations: name
1251 // resolve, purify, plan, sequence), all of which are managed by this function. DDL has
1252 // three notable properties that all partially interact.
1253 //
1254 // 1. Most DDL statements (and a few others) support single-statement transaction delayed
1255 // execution. This occurs when a session executes `BEGIN`, a single DDL, then `COMMIT`.
1256 // We announce success of the single DDL when it is executed, but do not attempt to plan
1257 // or sequence it until `COMMIT`, which is able to error if needed while sequencing the
1258 // DDL (this behavior is Postgres-compatible). The purpose of this is because some
1259 // drivers or tools wrap all statements in `BEGIN` and `COMMIT` and we would like them to
1260 // work. When the single DDL is announced as successful we also put the session's
1261 // transaction ops into `SingleStatement` which will produce an error if any other
1262 // statement is run in the transaction except `COMMIT`. Additionally, this will cause
1263 // `handle_execute_inner` to stop further processing (no planning, etc.) of the
1264 // statement.
1265 // 2. A few other DDL statements (`ALTER .. RENAME/SWAP`) enter the `DDL` ops which allows
1266 // any number of only these DDL statements to be executed in a transaction. During
1267 // sequencing we run an incremental catalog dry run against in-memory transaction state
1268 // and store the resulting `CatalogState` in `TransactionOps::DDL`, but nothing is yet
1269 // committed to the durable catalog. At `COMMIT`, all accumulated ops are applied in one
1270 // catalog transaction. The purpose of this is to allow multiple, atomic renames in the
1271 // same transaction.
1272 // 3. Some DDLs do off-thread work during purification or sequencing that is expensive or
1273 // makes network calls (interfacing with secrets, optimization of views/indexes, source
1274 // purification). These must guarantee correctness when they return to the main
1275 // coordinator thread because the catalog state could have changed while they were doing
1276 // the off-thread work. Previously we would use `PlanValidity::Checks` to specify a bunch
1277 // of IDs that we needed to exist. We discovered the way we were doing that was not
1278 // always correct. Instead of attempting to get that completely right, we have opted to
1279 // serialize DDL. Getting this right is difficult because catalog changes can affect name
1280 // resolution, planning, sequencing, and optimization. Correctly writing logic that is
1281 // aware of all possible catalog changes that would affect any of those parts is not
1282 // something our current code has been designed to be helpful at. Even if a DDL statement
1283 // is doing off-thread work, another DDL must not yet execute at all. Executing these
1284 // serially will guarantee that no off-thread work has affected the state of the catalog.
1285 // This is done by adding a VecDeque of deferred statements and a lock to the
1286 // Coordinator. When a DDL is run in `handle_execute_inner` (after applying whatever
1287 // transaction ops are needed to the session as described above), it attempts to own the
1288 // lock (a tokio Mutex). If acquired, it stashes the lock in the connection`s `ConnMeta`
1289 // struct in `active_conns` and proceeds. The lock is dropped at transaction end in
1290 // `clear_transaction` and a message sent to the Coordinator to execute the next queued
1291 // DDL. If the lock could not be acquired, the DDL is put into the VecDeque where it
1292 // awaits dequeuing caused by the lock being released.
1293
1294 if StatementKind::from(&*stmt).is_explain() {
1295 if let Err(error) = ctx.session_mut().transaction_mut().record_explain() {
1296 return ctx.retire(Err(error));
1297 }
1298 }
1299
1300 // For `Started`, this separates the first statement of an extended-protocol
1301 // pipeline from a later one that joins the ops staged before it.
1302 let txn_contains_ops = ctx.session().transaction().contains_ops();
1303
1304 // Verify that this statement type can be executed in the current
1305 // transaction state.
1306 match ctx.session().transaction() {
1307 // By this point we should be in a running transaction.
1308 TransactionStatus::Default => unreachable!(),
1309
1310 // Failed transactions have already been checked in pgwire for a safe statement
1311 // (COMMIT, ROLLBACK, etc.) and can proceed.
1312 TransactionStatus::Failed(_) => {}
1313
1314 // Started is a deceptive name, and means different things depending on which
1315 // protocol was used. It's either exactly one statement (known because this
1316 // is the simple protocol and the parser parsed the entire string, and it had
1317 // one statement). Or from the extended protocol, it means *some* query is
1318 // being executed, but there might be others after it before the Sync (commit)
1319 // message. Postgres handles this by teaching Started to eagerly commit certain
1320 // statements that can't be run in a transaction block.
1321 TransactionStatus::Started(_) if !txn_contains_ops => {
1322 if let Statement::Declare(_) = &*stmt {
1323 // Declare is an exception. Although it's not against any spec to execute
1324 // it, it will always result in nothing happening, since all portals will be
1325 // immediately closed. Users don't know this detail, so this error helps them
1326 // understand what's going wrong. Postgres does this too.
1327 return ctx.retire(Err(AdapterError::OperationRequiresTransaction(
1328 "DECLARE CURSOR".into(),
1329 )));
1330 }
1331 }
1332
1333 // Implicit or explicit transactions.
1334 //
1335 // Implicit transactions happen when a multi-statement query is executed
1336 // (a "simple query"). However if a "BEGIN" appears somewhere in there,
1337 // then the existing implicit transaction will be upgraded to an explicit
1338 // transaction. Thus, we should not separate what implicit and explicit
1339 // transactions can do unless there's some additional checking to make sure
1340 // something disallowed in explicit transactions did not previously take place
1341 // in the implicit portion.
1342 //
1343 // A `Started` transaction with staged ops belongs here too: this statement
1344 // runs alongside them, so a DDL or read-then-write that cannot see them must
1345 // be rejected rather than run against a state that lacks them.
1346 TransactionStatus::Started(_)
1347 | TransactionStatus::InTransactionImplicit(_)
1348 | TransactionStatus::InTransaction(_) => {
1349 match &*stmt {
1350 // Statements that are safe in a transaction. We still need to verify that we
1351 // don't interleave reads and writes since we can't perform those serializably.
1352 Statement::Close(_)
1353 | Statement::Commit(_)
1354 | Statement::Copy(_)
1355 | Statement::Deallocate(_)
1356 | Statement::Declare(_)
1357 | Statement::Discard(_)
1358 | Statement::Fetch(_)
1359 | Statement::Prepare(_)
1360 | Statement::Rollback(_)
1361 | Statement::SetTransaction(_)
1362 | Statement::Show(_)
1363 | Statement::SetVariable(_)
1364 | Statement::ResetVariable(_)
1365 | Statement::StartTransaction(_)
1366 | Statement::Raise(_)
1367 | Statement::ExplainPlan(_)
1368 | Statement::ExplainPushdown(_)
1369 | Statement::ExplainAnalyzeObject(_)
1370 | Statement::ExplainAnalyzeCluster(_)
1371 | Statement::ExplainSinkSchema(_) => {
1372 // Always safe.
1373 }
1374
1375 // `SessionClient::execute_attempts` unrolls `EXECUTE`, and the frontend takes
1376 // over the other statements, so none of them reach the coordinator. Planning one
1377 // below would hit `sequence_plan`'s soft-panicking arm.
1378 Statement::Delete(_)
1379 | Statement::Execute(_)
1380 | Statement::ExplainTimestamp(_)
1381 | Statement::Insert(_)
1382 | Statement::Select(_)
1383 | Statement::Subscribe(_)
1384 | Statement::Update(_) => {}
1385
1386 // These statements must be kept in-sync with `must_serialize_ddl()`.
1387 Statement::AlterObjectRename(_)
1388 | Statement::AlterObjectSwap(_)
1389 | Statement::CreateTableFromSource(_)
1390 | Statement::CreateSource(_) => {
1391 if !ctx.session().transaction().is_ddl() {
1392 // Preparation can yield before the first dry run. Capture both
1393 // baselines now so unrelated commits cannot mix their contents.
1394 let (state, snapshot) = self.catalog().ddl_transaction_snapshot().await;
1395
1396 // Initialize our transaction with a set of empty ops, or return an error
1397 // if we can't run a DDL transaction
1398 let txn_status = ctx.session_mut().transaction_mut();
1399 if let Err(err) = txn_status.add_ops(TransactionOps::DDL {
1400 commands: vec![],
1401 state,
1402 snapshot,
1403 }) {
1404 return ctx.retire(Err(err));
1405 }
1406 }
1407 }
1408
1409 // Statements below must by run singly (in Started).
1410 Statement::AlterCluster(_)
1411 | Statement::AlterConnection(_)
1412 | Statement::AlterDefaultPrivileges(_)
1413 | Statement::AlterIndex(_)
1414 | Statement::AlterMaterializedViewApplyReplacement(_)
1415 | Statement::AlterSetCluster(_)
1416 | Statement::AlterOwner(_)
1417 | Statement::AlterRetainHistory(_)
1418 | Statement::AlterRole(_)
1419 | Statement::AlterSecret(_)
1420 | Statement::AlterSink(_)
1421 | Statement::AlterSource(_)
1422 | Statement::AlterSystemReset(_)
1423 | Statement::AlterSystemResetAll(_)
1424 | Statement::AlterSystemSet(_)
1425 | Statement::AlterTableAddColumn(_)
1426 | Statement::AlterNetworkPolicy(_)
1427 | Statement::CreateCluster(_)
1428 | Statement::CreateClusterReplica(_)
1429 | Statement::CreateConnection(_)
1430 | Statement::CreateDatabase(_)
1431 | Statement::CreateIndex(_)
1432 | Statement::CreateMaterializedView(_)
1433 | Statement::CreateRole(_)
1434 | Statement::CreateSchema(_)
1435 | Statement::CreateSecret(_)
1436 | Statement::CreateSink(_)
1437 | Statement::CreateMetricSink(_)
1438 | Statement::CreateSubsource(_)
1439 | Statement::CreateTable(_)
1440 | Statement::CreateType(_)
1441 | Statement::CreateView(_)
1442 | Statement::CreateWebhookSource(_)
1443 | Statement::CreateNetworkPolicy(_)
1444 | Statement::DropObjects(_)
1445 | Statement::DropOwned(_)
1446 | Statement::GrantPrivileges(_)
1447 | Statement::GrantRole(_)
1448 | Statement::ReassignOwned(_)
1449 | Statement::RevokePrivileges(_)
1450 | Statement::RevokeRole(_)
1451 | Statement::ValidateConnection(_)
1452 | Statement::Comment(_)
1453 | Statement::ExecuteUnitTest(_) => {
1454 let txn_status = ctx.session_mut().transaction_mut();
1455
1456 // If we're not in an implicit transaction and we could generate exactly one
1457 // valid ExecuteResponse, we can delay execution until commit.
1458 if !txn_status.is_implicit() {
1459 // Statements whose tag is trivial (known only from an unexecuted statement) can
1460 // be run in a special single-statement explicit mode. In this mode (`BEGIN;
1461 // <stmt>; COMMIT`), we generate the expected tag from a successful <stmt>, but
1462 // delay execution until `COMMIT`.
1463 if let Ok(resp) = ExecuteResponse::try_from(&*stmt) {
1464 if let Err(err) = txn_status
1465 .add_ops(TransactionOps::SingleStatement { stmt, params })
1466 {
1467 ctx.retire(Err(err));
1468 return;
1469 }
1470 ctx.retire(Ok(resp));
1471 return;
1472 }
1473 }
1474
1475 // For statements that can carry sensitive material, redact
1476 // literals so they don't leak into the error message, which
1477 // is persisted in `mz_statement_execution_history` (matching
1478 // how their SQL text is redacted). Other statements keep
1479 // their literals for a clearer error.
1480 let op = if StatementKind::from(&*stmt).is_sensitive() {
1481 stmt.to_ast_string_redacted()
1482 } else {
1483 stmt.to_string()
1484 };
1485 return ctx.retire(Err(AdapterError::OperationProhibitsTransaction(op)));
1486 }
1487 }
1488 }
1489 }
1490
1491 // DDLs must be planned and sequenced serially. We do not rely on PlanValidity checking
1492 // various IDs because we have incorrectly done that in the past. Attempt to acquire the
1493 // ddl lock. The lock is stashed in the ConnMeta which is dropped at transaction end. If
1494 // acquired, proceed with sequencing. If not, enqueue and return. This logic assumes that
1495 // Coordinator::clear_transaction is correctly called when session transactions are ended
1496 // because that function will release the held lock from active_conns.
1497 if Self::must_serialize_ddl(&stmt, &ctx) {
1498 if let Ok(guard) = self.serialized_ddl.try_lock_owned() {
1499 let prev = self
1500 .active_conns
1501 .get_mut(ctx.session().conn_id())
1502 .expect("connection must exist")
1503 .deferred_lock
1504 .replace(guard);
1505 assert!(
1506 prev.is_none(),
1507 "connections should have at most one lock guard"
1508 );
1509 } else {
1510 if self
1511 .active_conns
1512 .get(ctx.session().conn_id())
1513 .expect("connection must exist")
1514 .deferred_lock
1515 .is_some()
1516 {
1517 // This session *already* has the lock, and incorrectly tried to execute another
1518 // DDL while still holding the lock, violating the assumption documented above.
1519 // This is an internal error, probably in some AdapterClient user (pgwire or
1520 // http). Because the session is now in some unexpected state, return an error
1521 // which should cause the AdapterClient user to fail the transaction.
1522 // (Terminating the connection is maybe what we would prefer to do, but is not
1523 // currently a thing we can do from the coordinator: calling handle_terminate
1524 // cleans up Coordinator state for the session but doesn't inform the
1525 // AdapterClient that the session should terminate.)
1526 soft_panic_or_log!(
1527 "session {} attempted to get ddl lock while already owning it",
1528 ctx.session().conn_id()
1529 );
1530 ctx.retire(Err(AdapterError::Internal(
1531 "session attempted to get ddl lock while already owning it".to_string(),
1532 )));
1533 return;
1534 }
1535 self.serialized_ddl
1536 .push_back(DeferredDdl::Statement { ctx, stmt, params });
1537 return;
1538 }
1539 }
1540
1541 if catalog::command::CatalogResultCommand::supports_statement(&stmt) {
1542 self.sequence_catalog_result_command(ctx, stmt, params);
1543 return;
1544 }
1545 if matches!(
1546 &*stmt,
1547 Statement::AlterObjectRename(_) | Statement::AlterObjectSwap(_)
1548 ) {
1549 self.sequence_rename_command(ctx, stmt, params).await;
1550 return;
1551 }
1552
1553 if matches!(
1554 &*stmt,
1555 Statement::CreateSource(_) | Statement::CreateTableFromSource(_)
1556 ) || matches!(&*stmt, Statement::AlterSource(stmt)
1557 if matches!(stmt.action,
1558 mz_sql::ast::AlterSourceAction::AddSubsources { .. }
1559 | mz_sql::ast::AlterSourceAction::RefreshReferences))
1560 {
1561 self.sequence_source_statement(ctx, stmt, params).await;
1562 return;
1563 }
1564
1565 let catalog = self.catalog();
1566 let catalog = catalog.for_session(ctx.session());
1567 let original_stmt = Arc::clone(&stmt);
1568 // `resolved_ids` should be derivable from `stmt`. If `stmt` is transformed to remove/add
1569 // IDs, then `resolved_ids` should be updated to also remove/add those IDs.
1570 let (stmt, mut resolved_ids) = match mz_sql::names::resolve(&catalog, (*stmt).clone()) {
1571 Ok(resolved) => resolved,
1572 Err(e) => return ctx.retire(Err(e.into())),
1573 };
1574 // N.B. The catalog can change during purification so we must validate that the dependencies still exist after
1575 // purification. This should be done back on the main thread.
1576 // We do the validation:
1577 // - In the handler for `Message::PurifiedStatementReady`, before we handle the purified statement.
1578 // If we add special handling for more types of `Statement`s, we'll need to ensure similar verification
1579 // occurs.
1580 let (stmt, resolved_ids) = match stmt {
1581 // Various statements must be purified off the main coordinator thread of control.
1582 stmt if Self::must_spawn_purification(&stmt) => {
1583 let internal_cmd_tx = self.internal_cmd_tx.clone();
1584 let conn_id = ctx.session().conn_id().clone();
1585 let catalog = self.owned_catalog();
1586 let now = self.now();
1587 let otel_ctx = OpenTelemetryContext::obtain();
1588 let current_storage_configuration = self.controller.storage.config().clone();
1589 task::spawn(|| format!("purify:{conn_id}"), async move {
1590 let conn_catalog = catalog.for_session(ctx.session());
1591
1592 let statement_source = mz_sql::pure::statement_source(&conn_catalog, &stmt);
1593
1594 // Authorization is usually checked after planning, but purification
1595 // happens before planning and may use connections and secrets, so it
1596 // gets its own check.
1597 if let Err(e) = rbac::check_purification(
1598 &conn_catalog,
1599 ctx.session(),
1600 statement_source,
1601 &resolved_ids,
1602 ) {
1603 return ctx.retire(Err(e.into()));
1604 }
1605
1606 let (result, cluster_id) = mz_sql::pure::purify_statement(
1607 conn_catalog,
1608 now,
1609 stmt,
1610 ¤t_storage_configuration,
1611 )
1612 .await;
1613 let result = result.map_err(|e| e.into());
1614 // `ALTER SOURCE` carries its target as an unresolved name, so name
1615 // resolution never records it and `resolved_ids` does not cover it.
1616 // Without it a source dropped while purification ran off-thread
1617 // passes the validity check below and then panics the coordinator
1618 // on the missing catalog entry ("catalog out of sync") during
1619 // planning.
1620 let mut dependency_ids: BTreeSet<_> = resolved_ids.items().copied().collect();
1621 dependency_ids.extend(statement_source.map(|source| source.id()));
1622 let plan_validity = PlanValidity::new(
1623 &catalog,
1624 dependency_ids,
1625 cluster_id,
1626 None,
1627 ctx.session().role_metadata().clone(),
1628 );
1629 // It is not an error for purification to complete after `internal_cmd_rx` is dropped.
1630 let result = internal_cmd_tx.send(Message::PurifiedStatementReady(
1631 PurifiedStatementReady {
1632 ctx,
1633 result,
1634 params,
1635 plan_validity,
1636 original_stmt,
1637 otel_ctx,
1638 },
1639 ));
1640 if let Err(e) = result {
1641 tracing::warn!("internal_cmd_rx dropped before we could send: {:?}", e);
1642 }
1643 });
1644 return;
1645 }
1646
1647 // `CREATE SUBSOURCE` statements are disallowed for users and are only generated
1648 // automatically as part of purification
1649 Statement::CreateSubsource(_) => {
1650 ctx.retire(Err(AdapterError::Unsupported(
1651 "CREATE SUBSOURCE statements",
1652 )));
1653 return;
1654 }
1655
1656 Statement::CreateMaterializedView(mut cmvs) => {
1657 // `CREATE MATERIALIZED VIEW ... AS OF ...` syntax is disallowed for users and is
1658 // only used for storing initial frontiers in the catalog.
1659 if cmvs.as_of.is_some() {
1660 return ctx.retire(Err(AdapterError::Unsupported(
1661 "CREATE MATERIALIZED VIEW ... AS OF statements",
1662 )));
1663 }
1664
1665 let mz_now = match self
1666 .resolve_mz_now_for_create_materialized_view(
1667 &cmvs,
1668 &resolved_ids,
1669 ctx.session_mut(),
1670 true,
1671 )
1672 .await
1673 {
1674 Ok(mz_now) => mz_now,
1675 Err(e) => return ctx.retire(Err(e)),
1676 };
1677
1678 let catalog = self.catalog().for_session(ctx.session());
1679
1680 purify_create_materialized_view_options(
1681 catalog,
1682 mz_now,
1683 &mut cmvs,
1684 &mut resolved_ids,
1685 );
1686
1687 let purified_stmt =
1688 Statement::CreateMaterializedView(CreateMaterializedViewStatement::<Aug> {
1689 if_exists: cmvs.if_exists,
1690 name: cmvs.name,
1691 columns: cmvs.columns,
1692 replacement_for: cmvs.replacement_for,
1693 in_cluster: cmvs.in_cluster,
1694 in_cluster_replica: cmvs.in_cluster_replica,
1695 query: cmvs.query,
1696 with_options: cmvs.with_options,
1697 as_of: None,
1698 });
1699
1700 // (Purifying CreateMaterializedView doesn't happen async, so no need to send
1701 // `Message::PurifiedStatementReady` here.)
1702 (purified_stmt, resolved_ids)
1703 }
1704
1705 Statement::ExplainPlan(ExplainPlanStatement {
1706 stage,
1707 with_options,
1708 format,
1709 explainee: Explainee::CreateMaterializedView(box_cmvs, broken),
1710 }) => {
1711 let mut cmvs = *box_cmvs;
1712 let mz_now = match self
1713 .resolve_mz_now_for_create_materialized_view(
1714 &cmvs,
1715 &resolved_ids,
1716 ctx.session_mut(),
1717 false,
1718 )
1719 .await
1720 {
1721 Ok(mz_now) => mz_now,
1722 Err(e) => return ctx.retire(Err(e)),
1723 };
1724
1725 let catalog = self.catalog().for_session(ctx.session());
1726
1727 purify_create_materialized_view_options(
1728 catalog,
1729 mz_now,
1730 &mut cmvs,
1731 &mut resolved_ids,
1732 );
1733
1734 let purified_stmt = Statement::ExplainPlan(ExplainPlanStatement {
1735 stage,
1736 with_options,
1737 format,
1738 explainee: Explainee::CreateMaterializedView(Box::new(cmvs), broken),
1739 });
1740
1741 (purified_stmt, resolved_ids)
1742 }
1743
1744 // All other statements are handled immediately.
1745 _ => (stmt, resolved_ids),
1746 };
1747
1748 match self.plan_statement(ctx.session(), stmt, ¶ms, &resolved_ids) {
1749 Ok((plan, sql_impl_ids)) => {
1750 self.sequence_plan(ctx, plan, resolved_ids, sql_impl_ids)
1751 .await
1752 }
1753 Err(e) => ctx.retire(Err(e)),
1754 }
1755 }
1756
1757 /// Whether the statement must be serialized and is DDL.
1758 fn must_serialize_ddl(stmt: &Statement<Raw>, ctx: &ExecuteContext) -> bool {
1759 // Non-DDL is not serialized here.
1760 if !StatementClassification::from(&*stmt).is_ddl() {
1761 return false;
1762 }
1763 // External preparation can perform arbitrarily slow network calls, so it must not
1764 // hold the DDL lock. Source commands re-execute against the current catalog at commit.
1765 // CREATE SINK and the remaining purified ALTER SOURCE forms use PlanValidity.
1766 if Self::must_spawn_purification(stmt) {
1767 return false;
1768 }
1769
1770 // Statements that support multiple DDLs in a single transaction aren't serialized here.
1771 // Their operations are serialized when applied to the catalog, guaranteeing that any
1772 // off-thread DDLs concurrent with a multiple DDL transaction will have a serial order.
1773 if ctx.session.transaction().is_ddl() {
1774 return false;
1775 }
1776
1777 // Some DDL is exempt. It is not great that we are matching on Statements here because
1778 // different plans can be produced from the same top-level statement type (i.e., `ALTER
1779 // CONNECTION ROTATE KEYS`). But the whole point of this is to prevent things from being
1780 // planned in the first place, so we accept the abstraction leak.
1781 match stmt {
1782 // Secrets have a small and understood set of dependencies, and their off-thread work
1783 // interacts with k8s.
1784 Statement::AlterSecret(_) => false,
1785 Statement::CreateSecret(_) => false,
1786 Statement::AlterConnection(AlterConnectionStatement { actions, .. })
1787 if actions
1788 .iter()
1789 .all(|action| matches!(action, AlterConnectionAction::RotateKeys)) =>
1790 {
1791 false
1792 }
1793
1794 // The off-thread work that altering a cluster may do (waiting for replicas to spin-up),
1795 // does not affect its catalog names or ids and so is safe to not serialize. This could
1796 // change the set of replicas that exist. For queries that name replicas or use the
1797 // current_replica session var, the `replica_id` field of `PlanValidity` serves to
1798 // ensure that those replicas exist during the query finish stage. Additionally, that
1799 // work can take hours (configured by the user), so would also be a bad experience for
1800 // users.
1801 Statement::AlterCluster(_) => false,
1802
1803 // `ALTER SINK` waits for the sink to make enough progress for a clean cutover to the
1804 // new configuration. If the sink is stalled, it may block forever. Checks in
1805 // sequencing ensure that the operation fails if any one of these happens concurrently:
1806 // * the sink is dropped
1807 // * the source relation is dropped
1808 // * another `ALTER SINK` for the same sink is applied first
1809 Statement::AlterSink(_) => false,
1810
1811 // `ALTER MATERIALIZED VIEW ... APPLY REPLACEMENT` waits for the target MV to make
1812 // enough progress for a clean cutover. If the target MV is stalled, it may block
1813 // forever. Checks in sequencing ensure the operation fails if any of these happens
1814 // concurrently:
1815 // * the target MV is dropped
1816 // * the replacement MV is dropped
1817 Statement::AlterMaterializedViewApplyReplacement(_) => false,
1818
1819 // Everything else must be serialized.
1820 _ => true,
1821 }
1822 }
1823
1824 /// Whether the statement must be purified off of the Coordinator thread.
1825 ///
1826 /// Every statement listed here is authorized by [`rbac::check_purification`]
1827 /// before purification runs, against the source that
1828 /// [`mz_sql::pure::statement_source`] resolves for it.
1829 fn must_spawn_purification<A: AstInfo>(stmt: &Statement<A>) -> bool {
1830 // `CREATE` and `ALTER` `SOURCE` and `SINK` statements must be purified off the main
1831 // coordinator thread.
1832 if !matches!(
1833 stmt,
1834 Statement::CreateSource(_)
1835 | Statement::AlterSource(_)
1836 | Statement::CreateSink(_)
1837 | Statement::CreateTableFromSource(_)
1838 ) {
1839 return false;
1840 }
1841
1842 // However `ALTER SOURCE RETAIN HISTORY` should be excluded from off-thread purification.
1843 if let Statement::AlterSource(stmt) = stmt {
1844 let names: Vec<CreateSourceOptionName> = match &stmt.action {
1845 AlterSourceAction::SetOptions(options) => {
1846 options.iter().map(|o| o.name.clone()).collect()
1847 }
1848 AlterSourceAction::ResetOptions(names) => names.clone(),
1849 _ => vec![],
1850 };
1851 if !names.is_empty()
1852 && names
1853 .iter()
1854 .all(|n| matches!(n, CreateSourceOptionName::RetainHistory))
1855 {
1856 return false;
1857 }
1858 }
1859
1860 true
1861 }
1862
1863 /// Chooses a timestamp for `mz_now()`, if `mz_now()` occurs in a REFRESH option of the
1864 /// materialized view. Additionally, if `acquire_read_holds` is true and the MV has any REFRESH
1865 /// option, this function grabs read holds at the earliest possible time on input collections
1866 /// that might be involved in the MV.
1867 ///
1868 /// Note that this is NOT what handles `mz_now()` in the query part of the MV. (handles it only
1869 /// in `with_options`).
1870 ///
1871 /// (Note that the chosen timestamp won't be the same timestamp as the system table inserts,
1872 /// unfortunately.)
1873 async fn resolve_mz_now_for_create_materialized_view(
1874 &mut self,
1875 cmvs: &CreateMaterializedViewStatement<Aug>,
1876 resolved_ids: &ResolvedIds,
1877 session: &Session,
1878 acquire_read_holds: bool,
1879 ) -> Result<Option<Timestamp>, AdapterError> {
1880 if cmvs
1881 .with_options
1882 .iter()
1883 .any(|wo| matches!(wo.value, Some(WithOptionValue::Refresh(..))))
1884 {
1885 let catalog = self.catalog().for_session(session);
1886 let cluster = mz_sql::plan::resolve_cluster_for_materialized_view(&catalog, cmvs)?;
1887 let ids = self
1888 .index_oracle(cluster)
1889 .sufficient_collections(resolved_ids.collections().copied());
1890
1891 // If there is any REFRESH option, then acquire read holds. (Strictly speaking, we'd
1892 // need this only if there is a `REFRESH AT`, not for `REFRESH EVERY`, because later
1893 // we want to check the AT times against the read holds that we acquire here. But
1894 // we do it for any REFRESH option, to avoid having so many code paths doing different
1895 // things.)
1896 //
1897 // It's important that we acquire read holds _before_ we determine the least valid read.
1898 // Otherwise, we're not guaranteed that the since frontier doesn't
1899 // advance forward from underneath us.
1900 let read_holds = self.acquire_read_holds(&ids);
1901
1902 // Does `mz_now()` occur?
1903 let mz_now_ts = if cmvs
1904 .with_options
1905 .iter()
1906 .any(materialized_view_option_contains_temporal)
1907 {
1908 let timeline_context = self
1909 .catalog()
1910 .validate_timeline_context(resolved_ids.collections().copied())?;
1911
1912 // We default to EpochMilliseconds, similarly to `determine_timestamp_for`,
1913 // but even in the TimestampIndependent case.
1914 // Note that we didn't accurately decide whether we are TimestampDependent
1915 // or TimestampIndependent, because for this we'd need to also check whether
1916 // `query.contains_temporal()`, similarly to how `peek_stage_validate` does.
1917 // However, this doesn't matter here, as we are just going to default to
1918 // EpochMilliseconds in both cases.
1919 let timeline = timeline_context
1920 .timeline()
1921 .unwrap_or(&Timeline::EpochMilliseconds);
1922
1923 // Let's start with the timestamp oracle read timestamp.
1924 let mut timestamp = self.get_timestamp_oracle(timeline).read_ts().await;
1925
1926 // If `least_valid_read` is later than the oracle, then advance to that time.
1927 // If we didn't do this, then there would be a danger of missing the first refresh,
1928 // which might cause the materialized view to be unreadable for hours. This might
1929 // be what was happening here:
1930 // https://github.com/MaterializeInc/database-issues/issues/7265#issuecomment-1931856361
1931 //
1932 // In the long term, it would be good to actually block the MV creation statement
1933 // until `least_valid_read`. https://github.com/MaterializeInc/database-issues/issues/7504
1934 // Without blocking, we have the problem that a REFRESH AT CREATION is not linearized
1935 // with the CREATE MATERIALIZED VIEW statement, in the sense that a query from the MV
1936 // after its creation might see input changes that happened after the CRATE MATERIALIZED
1937 // VIEW statement returned.
1938 let oracle_timestamp = timestamp;
1939 let least_valid_read = read_holds.least_valid_read();
1940 timestamp.advance_by(least_valid_read.borrow());
1941
1942 if oracle_timestamp != timestamp {
1943 warn!(%cmvs.name, %oracle_timestamp, %timestamp, "REFRESH MV's inputs are not readable at the oracle read ts");
1944 }
1945
1946 info!("Resolved `mz_now()` to {timestamp} for REFRESH MV");
1947 Ok(Some(timestamp))
1948 } else {
1949 Ok(None)
1950 };
1951
1952 // NOTE: The Drop impl of ReadHolds makes sure that the hold is
1953 // released when we don't use it.
1954 if acquire_read_holds {
1955 self.store_transaction_read_holds(session.conn_id().clone(), read_holds);
1956 }
1957
1958 mz_now_ts
1959 } else {
1960 Ok(None)
1961 }
1962 }
1963
1964 /// Instruct the dataflow layer to cancel any ongoing, interactive work for
1965 /// the named `conn_id` if the correct secret key is specified.
1966 ///
1967 /// Note: Here we take a [`ConnectionIdType`] as opposed to an owned
1968 /// `ConnectionId` because this method gets called by external clients when
1969 /// they request to cancel a request.
1970 #[mz_ore::instrument(level = "debug")]
1971 async fn handle_cancel(&mut self, conn_id: ConnectionIdType, secret_key: u32) {
1972 if let Some((id_handle, conn_meta)) = self.active_conns.get_key_value(&conn_id) {
1973 // If the secret key specified by the client doesn't match the
1974 // actual secret key for the target connection, we treat this as a
1975 // rogue cancellation request and ignore it.
1976 if conn_meta.secret_key != secret_key {
1977 return;
1978 }
1979
1980 // Now that we've verified the secret key, this is a privileged
1981 // cancellation request. We can upgrade the raw connection ID to a
1982 // proper `IdHandle`.
1983 self.handle_privileged_cancel(id_handle.clone()).await;
1984 }
1985 }
1986
1987 /// Unconditionally instructs the dataflow layer to cancel any ongoing,
1988 /// interactive work for the named `conn_id`.
1989 #[mz_ore::instrument(level = "debug")]
1990 pub(crate) async fn handle_privileged_cancel(&mut self, conn_id: ConnectionId) {
1991 if let Some(conn) = self.active_conns.get(&conn_id) {
1992 conn.frontend_cancel_tx.send_replace(());
1993 }
1994 let mut maybe_ctx = None;
1995 // A COMMIT waiting for group commit or for its read to be linearized.
1996 let mut maybe_commit_ctx = None;
1997
1998 // Cancel all pending writes for this connection:
1999 // - At most one session-bound write (`UserWriteResponder::Session`),
2000 // retired via its `ExecuteContext` below.
2001 // - Any frontend blind write that has not entered the committer. The
2002 // waiter receives `WriteResult::Canceled`. Timestamped writes already
2003 // in the committer have an indeterminate outcome until it responds.
2004 let (cancelled, kept): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pending_writes)
2005 .into_iter()
2006 .partition(|pending_write_txn| match pending_write_txn {
2007 PendingWriteTxn::User {
2008 responder: UserWriteResponder::Session(PendingTxn { ctx, .. }),
2009 ..
2010 } => *ctx.session().conn_id() == conn_id,
2011 PendingWriteTxn::User {
2012 responder: UserWriteResponder::Internal { conn_id: c, .. },
2013 ..
2014 } => *c == conn_id,
2015 PendingWriteTxn::System { .. } => false,
2016 });
2017 self.pending_writes = kept;
2018 for pending in cancelled {
2019 match pending {
2020 PendingWriteTxn::User {
2021 responder: UserWriteResponder::Session(PendingTxn { ctx, .. }),
2022 ..
2023 } => {
2024 maybe_commit_ctx = Some(ctx);
2025 }
2026 PendingWriteTxn::User {
2027 responder: UserWriteResponder::Internal { result, .. },
2028 ..
2029 } => {
2030 result.send(WriteResult::Canceled);
2031 }
2032 PendingWriteTxn::System { .. } => unreachable!("filtered out above"),
2033 }
2034 }
2035
2036 // Cancel deferred statements.
2037 let deferred_ddl_idx = self
2038 .serialized_ddl
2039 .iter()
2040 .position(|deferred| *deferred.conn_id() == conn_id);
2041 if let Some(idx) = deferred_ddl_idx {
2042 let deferred = self
2043 .serialized_ddl
2044 .remove(idx)
2045 .expect("known to exist from call to `position` above");
2046 match deferred {
2047 DeferredDdl::Statement {
2048 ctx,
2049 stmt: _,
2050 params: _,
2051 } => maybe_ctx = Some(ctx),
2052 DeferredDdl::Commit(pending) => {
2053 self.finish_ddl_commit(pending, Err(AdapterError::Canceled));
2054 }
2055 }
2056 }
2057
2058 // Cancel reads waiting on being linearized. There is at most one linearized read per
2059 // session.
2060 if let Some(pending_read_txn) = self.pending_linearize_read_txns.remove(&conn_id) {
2061 maybe_commit_ctx = Some(pending_read_txn.take_context());
2062 }
2063
2064 if let Some(ctx) = maybe_commit_ctx {
2065 // The COMMIT already cleared the session's transaction status, so the transaction's
2066 // `SET LOCAL` and staged variable values are the only transaction state left. Without
2067 // ending them here they apply to the session's next transaction.
2068 //
2069 // NOTE: The error response cannot carry the changed parameters, so a reverted
2070 // parameter in the notify set is not announced to the client.
2071 let (ctx, result) = CompletedClientTransmitter::new(
2072 ctx,
2073 Err(AdapterError::Canceled),
2074 EndTransactionAction::Rollback,
2075 )
2076 .finalize();
2077 ctx.retire(result);
2078 }
2079 if let Some(ctx) = maybe_ctx {
2080 ctx.retire(Err(AdapterError::Canceled));
2081 }
2082
2083 self.cancel_pending_peeks(&conn_id);
2084 self.cancel_pending_watchsets(&conn_id);
2085 let retire_notify = self.cancel_compute_sinks_for_conn(&conn_id).await;
2086 // SQL cancellation has no success response to delay. Each subscribe
2087 // still waits for its own retraction before it observes retirement.
2088 drop(retire_notify);
2089 self.cancel_pending_copy(&conn_id);
2090 if let Some((tx, _rx)) = self.connection_cancel_watches.get_mut(&conn_id) {
2091 let _ = tx.send(true);
2092 }
2093 }
2094
2095 /// Handle termination of a client session.
2096 ///
2097 /// This cleans up any state in the coordinator associated with the session.
2098 ///
2099 /// Must only be called for connections that completed a successful
2100 /// startup, i.e. that are present in `active_conns`. A failed startup
2101 /// leaves no state behind (see `handle_startup_inner`), so no Terminate
2102 /// may be sent for it.
2103 #[mz_ore::instrument(level = "debug")]
2104 async fn handle_terminate(&mut self, conn_id: ConnectionId) {
2105 // Checked up front so that an unknown connection panics before the cleanup below runs.
2106 assert!(
2107 self.active_conns.contains_key(&conn_id),
2108 "unknown connection: {conn_id:?}"
2109 );
2110
2111 // We do not need to call clear_transaction here because there are no side effects to run
2112 // based on any session transaction state.
2113 let retire_notify = self.clear_connection(&conn_id).await;
2114 // Termination has no statement response to delay. Each subscribe still
2115 // waits for its own retraction before it observes retirement.
2116 drop(retire_notify);
2117
2118 self.drop_temp_items(&conn_id).await;
2119 // Only call catalog_mut() if a temporary namespace actually exists for this connection.
2120 // This avoids an expensive Arc::make_mut clone for the common case where the connection
2121 // never created any temporary objects.
2122 if self.catalog().state().has_temporary_namespace(&conn_id) {
2123 self.catalog_mut().drop_temporary_namespace(&conn_id);
2124 }
2125 let conn = self.active_conns.remove(&conn_id).expect("conn must exist");
2126 let session_type = metrics::session_type_label_value(conn.user());
2127 self.metrics
2128 .active_sessions
2129 .with_label_values(&[session_type])
2130 .dec();
2131 self.cancel_pending_peeks(conn.conn_id());
2132 self.cancel_pending_watchsets(&conn_id);
2133 self.cancel_pending_copy(&conn_id);
2134 self.end_session_for_statement_logging(conn.uuid());
2135
2136 // Queue the retraction for the next group commit without waiting for it, so closing many
2137 // connections at once doesn't block the Coordinator. Forcing a commit here (`defer`)
2138 // would hand every `mz_catalog_server` dataflow a new timestamp per closed session,
2139 // making connection churn cost catalog server CPU per connection.
2140 let update = self
2141 .catalog()
2142 .state()
2143 .pack_session_update(&conn, Diff::MINUS_ONE);
2144 let update = self.catalog().state().resolve_builtin_table_update(update);
2145
2146 let _builtin_update_notify = self.builtin_table_update().background(vec![update]);
2147 }
2148
2149 /// Returns the necessary metadata for appending to a webhook source, and a channel to send
2150 /// rows.
2151 #[mz_ore::instrument(level = "debug")]
2152 fn handle_get_webhook(
2153 &mut self,
2154 database: String,
2155 schema: String,
2156 name: String,
2157 tx: oneshot::Sender<Result<AppendWebhookResponse, AppendWebhookError>>,
2158 ) {
2159 /// Attempts to resolve a Webhook source from a provided `database.schema.name` path.
2160 ///
2161 /// Returns a struct that can be used to append data to the underlying storate collection, and the
2162 /// types we should cast the request to.
2163 fn resolve(
2164 coord: &mut Coordinator,
2165 database: String,
2166 schema: String,
2167 name: String,
2168 ) -> Result<AppendWebhookResponse, PartialItemName> {
2169 // Resolve our collection.
2170 let name = PartialItemName {
2171 database: Some(database),
2172 schema: Some(schema),
2173 item: name,
2174 };
2175 let Ok(entry) = coord
2176 .catalog()
2177 .resolve_entry(None, &vec![], &name, &SYSTEM_CONN_ID)
2178 else {
2179 return Err(name);
2180 };
2181
2182 // Webhooks can be created with `CREATE SOURCE` or `CREATE TABLE`.
2183 let (data_source, desc, global_id) = match entry.item() {
2184 CatalogItem::Source(Source {
2185 data_source: data_source @ DataSourceDesc::Webhook { .. },
2186 desc,
2187 global_id,
2188 ..
2189 }) => (data_source, desc.clone(), *global_id),
2190 CatalogItem::Table(
2191 table @ Table {
2192 desc,
2193 data_source:
2194 TableDataSource::DataSource {
2195 desc: data_source @ DataSourceDesc::Webhook { .. },
2196 ..
2197 },
2198 ..
2199 },
2200 ) => (data_source, desc.latest(), table.global_id_writes()),
2201 _ => return Err(name),
2202 };
2203
2204 let DataSourceDesc::Webhook {
2205 validate_using,
2206 body_format,
2207 headers,
2208 ..
2209 } = data_source
2210 else {
2211 mz_ore::soft_panic_or_log!("programming error! checked above for webhook");
2212 return Err(name);
2213 };
2214 let body_format = body_format.clone();
2215 let header_tys = headers.clone();
2216
2217 // Assert we have one column for the body, and how ever many are required for
2218 // the headers.
2219 let num_columns = headers.num_columns() + 1;
2220 mz_ore::soft_assert_or_log!(
2221 desc.arity() <= num_columns,
2222 "expected at most {} columns, but got {}",
2223 num_columns,
2224 desc.arity()
2225 );
2226
2227 // Double check that the body column of the webhook source matches the type
2228 // we're about to deserialize as.
2229 let body_column = desc
2230 .get_by_name(&"body".into())
2231 .map(|(_idx, ty)| ty.clone())
2232 .ok_or_else(|| name.clone())?;
2233 assert!(!body_column.nullable, "webhook body column is nullable!?");
2234 assert_eq!(body_column.scalar_type, SqlScalarType::from(body_format));
2235
2236 // Create a validator that can be called to validate a webhook request.
2237 let validator = validate_using.as_ref().map(|v| {
2238 let validation = v.clone();
2239 AppendWebhookValidator::new(validation, coord.caching_secrets_reader.clone())
2240 });
2241
2242 // Get a channel so we can queue updates to be written.
2243 let row_tx = coord
2244 .controller
2245 .storage
2246 .monotonic_appender(global_id)
2247 .map_err(|_| name.clone())?;
2248 let stats = coord
2249 .controller
2250 .storage
2251 .webhook_statistics(global_id)
2252 .map_err(|_| name)?;
2253 let invalidator = coord
2254 .active_webhooks
2255 .entry(entry.id())
2256 .or_insert_with(WebhookAppenderInvalidator::new);
2257 let tx = WebhookAppender::new(row_tx, invalidator.guard(), stats);
2258
2259 Ok(AppendWebhookResponse {
2260 tx,
2261 body_format,
2262 header_tys,
2263 validator,
2264 })
2265 }
2266
2267 let response = resolve(self, database, schema, name).map_err(|name| {
2268 AppendWebhookError::UnknownWebhook {
2269 database: name.database.expect("provided"),
2270 schema: name.schema.expect("provided"),
2271 name: name.item,
2272 }
2273 });
2274 let _ = tx.send(response);
2275 }
2276
2277 /// Handle registration of a frontend peek, for statement logging and query cancellation
2278 /// handling.
2279 fn handle_register_peek(
2280 &mut self,
2281 uuid: Uuid,
2282 conn_id: ConnectionId,
2283 cluster_id: mz_controller_types::ClusterId,
2284 depends_on: BTreeSet<GlobalId>,
2285 is_fast_path: bool,
2286 watch_set: Option<WatchSetCreation>,
2287 tx: oneshot::Sender<Result<(), AdapterError>>,
2288 ) {
2289 let statement_logging_id = watch_set.as_ref().map(|ws| ws.logging_id);
2290 if let Some(ws) = watch_set {
2291 if let Err(e) = self.install_peek_watch_sets(conn_id.clone(), ws) {
2292 let _ = tx.send(Err(
2293 AdapterError::concurrent_dependency_drop_from_compute_lookup_error(e),
2294 ));
2295 return;
2296 }
2297 }
2298
2299 // Store the peek in pending_peeks for later retrieval when results arrive
2300 self.pending_peeks.insert(
2301 uuid,
2302 PendingPeek {
2303 conn_id: conn_id.clone(),
2304 cluster_id,
2305 depends_on,
2306 ctx_extra: ExecuteContextGuard::new(
2307 statement_logging_id,
2308 self.internal_cmd_tx.clone(),
2309 ),
2310 is_fast_path,
2311 },
2312 );
2313
2314 // Also track it by connection ID for cancellation support
2315 self.client_pending_peeks
2316 .entry(conn_id)
2317 .or_default()
2318 .insert(uuid, cluster_id);
2319
2320 let _ = tx.send(Ok(()));
2321 }
2322
2323 /// Handles [`Command::UnregisterPeek`]; see its documentation for
2324 /// the end-of-execution ownership contract.
2325 fn handle_unregister_peek(
2326 &mut self,
2327 uuid: Uuid,
2328 reason: StatementEndedExecutionReason,
2329 tx: oneshot::Sender<()>,
2330 ) {
2331 // A peek missing from `pending_peeks` was already retired, and its end
2332 // logged, by a concurrent teardown.
2333 if let Some(pending_peek) = self.remove_pending_peek(&uuid) {
2334 self.retire_execution(reason, pending_peek.ctx_extra.defuse());
2335 }
2336 let _ = tx.send(());
2337 }
2338}