mz_adapter/
statement_logging.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::borrow::Cow;
11use std::collections::BTreeSet;
12use std::sync::atomic::Ordering;
13use std::sync::{Arc, Mutex};
14
15use bytes::BytesMut;
16use mz_catalog::memory::objects::CatalogItem;
17use mz_controller_types::ClusterId;
18use mz_ore::cast::{CastFrom, CastInto};
19use mz_ore::now::{EpochMillis, NowFn, epoch_to_uuid_v7, to_datetime};
20use mz_ore::soft_panic_or_log;
21use mz_repr::adt::array::ArrayDimension;
22use mz_repr::adt::timestamp::TimestampLike;
23use mz_repr::{Datum, GlobalId, Row, RowIterator, RowPacker, Timestamp};
24use mz_sql::ast::display::AstDisplay;
25use mz_sql::ast::{AstInfo, Statement};
26use mz_sql::plan::Params;
27use mz_sql::session::metadata::SessionMetadata;
28use mz_sql::session::vars::SystemVars;
29use mz_sql_parser::ast::{StatementKind, statement_kind_label_value};
30use qcell::QCell;
31use rand::distr::{Bernoulli, Distribution};
32use sha2::{Digest, Sha256};
33use uuid::Uuid;
34
35use crate::catalog::CatalogState;
36use crate::session::{LifecycleTimestamps, Session, TransactionId};
37use crate::{AdapterError, CollectionIdBundle, ExecuteResponse};
38
39#[derive(Clone, Debug)]
40pub enum StatementLifecycleEvent {
41    ExecutionBegan,
42    OptimizationFinished,
43    StorageDependenciesFinished,
44    ComputeDependenciesFinished,
45    ExecutionFinished,
46}
47
48impl StatementLifecycleEvent {
49    pub fn as_str(&self) -> &str {
50        match self {
51            Self::ExecutionBegan => "execution-began",
52            Self::OptimizationFinished => "optimization-finished",
53            Self::StorageDependenciesFinished => "storage-dependencies-finished",
54            Self::ComputeDependenciesFinished => "compute-dependencies-finished",
55            Self::ExecutionFinished => "execution-finished",
56        }
57    }
58}
59
60/// Contains all the information necessary to generate the initial
61/// entry in `mz_statement_execution_history`. We need to keep this
62/// around in order to modify the entry later once the statement finishes executing.
63#[derive(Clone, Debug)]
64pub struct StatementBeganExecutionRecord {
65    pub id: Uuid,
66    pub prepared_statement_id: Uuid,
67    pub sample_rate: f64,
68    pub params: Vec<Option<String>>,
69    pub began_at: EpochMillis,
70    pub cluster_id: Option<ClusterId>,
71    pub cluster_name: Option<String>,
72    pub database_name: String,
73    pub search_path: Vec<String>,
74    pub application_name: String,
75    pub transaction_isolation: String,
76    pub execution_timestamp: Option<EpochMillis>,
77    pub transaction_id: TransactionId,
78    pub transient_index_id: Option<GlobalId>,
79    pub mz_version: String,
80    /// The kind of statement being executed, if known. Used to redact
81    /// `error_message` for kinds that can carry secret material.
82    pub kind: Option<StatementKind>,
83}
84
85#[derive(Clone, Copy, Debug)]
86pub enum StatementExecutionStrategy {
87    /// The statement was executed by spinning up a dataflow.
88    Standard,
89    /// The statement was executed by reading from an existing
90    /// arrangement.
91    FastPath,
92    /// Experimental: The statement was executed by reading from an existing
93    /// persist collection.
94    PersistFastPath,
95    /// The statement was determined to be constant by
96    /// environmentd, and not sent to a cluster.
97    Constant,
98}
99
100impl StatementExecutionStrategy {
101    pub fn name(&self) -> &'static str {
102        match self {
103            Self::Standard => "standard",
104            Self::FastPath => "fast-path",
105            Self::PersistFastPath => "persist-fast-path",
106            Self::Constant => "constant",
107        }
108    }
109}
110
111#[derive(Clone, Debug)]
112pub enum StatementEndedExecutionReason {
113    Success {
114        result_size: Option<u64>,
115        rows_returned: Option<u64>,
116        execution_strategy: Option<StatementExecutionStrategy>,
117    },
118    Canceled,
119    Errored {
120        error: String,
121    },
122    /// Should only be emitted by `impl Drop for ExecuteContextGuard`.
123    /// Code paths that explicitly complete execution should use
124    /// `retire_execute` with `Success`, `Canceled`, or `Errored`.
125    Aborted,
126}
127
128#[derive(Clone, Debug)]
129pub struct StatementEndedExecutionRecord {
130    pub id: Uuid,
131    pub reason: StatementEndedExecutionReason,
132    pub ended_at: EpochMillis,
133}
134
135/// Contains all the information necessary to generate an entry in
136/// `mz_prepared_statement_history`
137#[derive(Clone, Debug)]
138pub(crate) struct StatementPreparedRecord {
139    pub id: Uuid,
140    pub sql_hash: [u8; 32],
141    pub name: String,
142    pub session_id: Uuid,
143    pub prepared_at: EpochMillis,
144    pub kind: Option<StatementKind>,
145}
146
147#[derive(Clone, Debug)]
148pub(crate) struct SessionHistoryEvent {
149    pub id: Uuid,
150    pub connected_at: EpochMillis,
151    pub application_name: String,
152    pub authenticated_user: String,
153}
154
155impl From<&Result<ExecuteResponse, AdapterError>> for StatementEndedExecutionReason {
156    fn from(value: &Result<ExecuteResponse, AdapterError>) -> StatementEndedExecutionReason {
157        match value {
158            Ok(resp) => resp.into(),
159            Err(e) => StatementEndedExecutionReason::Errored {
160                error: e.to_string(),
161            },
162        }
163    }
164}
165
166impl From<&ExecuteResponse> for StatementEndedExecutionReason {
167    fn from(value: &ExecuteResponse) -> StatementEndedExecutionReason {
168        match value {
169            ExecuteResponse::CopyTo { resp, .. } => match resp.as_ref() {
170                // NB [btv]: It's not clear that this combination
171                // can ever actually happen.
172                ExecuteResponse::SendingRowsImmediate { rows, .. } => {
173                    // Note(parkmycar): It potentially feels bad here to iterate over the entire
174                    // iterator _just_ to get the encoded result size. As noted above, it's not
175                    // entirely clear this case ever happens, so the simplicity is worth it.
176                    let result_size: usize = rows.box_clone().map(|row| row.byte_len()).sum();
177                    StatementEndedExecutionReason::Success {
178                        result_size: Some(u64::cast_from(result_size)),
179                        rows_returned: Some(u64::cast_from(rows.count())),
180                        execution_strategy: Some(StatementExecutionStrategy::Constant),
181                    }
182                }
183                ExecuteResponse::SendingRowsStreaming { .. } => {
184                    panic!("SELECTs terminate on peek finalization, not here.")
185                }
186                ExecuteResponse::Subscribing { .. } => {
187                    panic!("SUBSCRIBEs terminate in the protocol layer, not here.")
188                }
189                _ => panic!("Invalid COPY response type"),
190            },
191            ExecuteResponse::CopyFrom { .. } => {
192                panic!("COPY FROMs terminate in the protocol layer, not here.")
193            }
194            ExecuteResponse::Fetch { .. } => {
195                panic!("FETCHes terminate after a follow-up message is sent.")
196            }
197            ExecuteResponse::SendingRowsStreaming { .. } => {
198                panic!("SELECTs terminate on peek finalization, not here.")
199            }
200            ExecuteResponse::Subscribing { .. } => {
201                panic!("SUBSCRIBEs terminate in the protocol layer, not here.")
202            }
203
204            ExecuteResponse::SendingRowsImmediate { rows, .. } => {
205                // Note(parkmycar): It potentially feels bad here to iterate over the entire
206                // iterator _just_ to get the encoded result size, the number of Rows returned here
207                // shouldn't be too large though. An alternative is to pre-compute some of the
208                // result size, but that would require always decoding Rows to handle projecting
209                // away columns, which has a negative impact for much larger response sizes.
210                let result_size: usize = rows.box_clone().map(|row| row.byte_len()).sum();
211                StatementEndedExecutionReason::Success {
212                    result_size: Some(u64::cast_from(result_size)),
213                    rows_returned: Some(u64::cast_from(rows.count())),
214                    execution_strategy: Some(StatementExecutionStrategy::Constant),
215                }
216            }
217
218            ExecuteResponse::AlteredDefaultPrivileges
219            | ExecuteResponse::AlteredObject(_)
220            | ExecuteResponse::AlteredRole
221            | ExecuteResponse::AlteredSystemConfiguration
222            | ExecuteResponse::ClosedCursor
223            | ExecuteResponse::Comment
224            | ExecuteResponse::Copied(_)
225            | ExecuteResponse::CreatedConnection
226            | ExecuteResponse::CreatedDatabase
227            | ExecuteResponse::CreatedSchema
228            | ExecuteResponse::CreatedRole
229            | ExecuteResponse::CreatedCluster
230            | ExecuteResponse::CreatedClusterReplica
231            | ExecuteResponse::CreatedIndex
232            | ExecuteResponse::CreatedIntrospectionSubscribe
233            | ExecuteResponse::CreatedMetricSink
234            | ExecuteResponse::CreatedSecret
235            | ExecuteResponse::CreatedSink
236            | ExecuteResponse::CreatedSource
237            | ExecuteResponse::CreatedTable
238            | ExecuteResponse::CreatedView
239            | ExecuteResponse::CreatedViews
240            | ExecuteResponse::CreatedMaterializedView
241            | ExecuteResponse::CreatedType
242            | ExecuteResponse::CreatedNetworkPolicy
243            | ExecuteResponse::Deallocate { .. }
244            | ExecuteResponse::DeclaredCursor
245            | ExecuteResponse::Deleted(_)
246            | ExecuteResponse::DiscardedTemp
247            | ExecuteResponse::DiscardedAll
248            | ExecuteResponse::DroppedObject(_)
249            | ExecuteResponse::DroppedOwned
250            | ExecuteResponse::EmptyQuery
251            | ExecuteResponse::GrantedPrivilege
252            | ExecuteResponse::GrantedRole
253            | ExecuteResponse::Inserted(_)
254            | ExecuteResponse::Prepare
255            | ExecuteResponse::Raised
256            | ExecuteResponse::ReassignOwned
257            | ExecuteResponse::RevokedPrivilege
258            | ExecuteResponse::RevokedRole
259            | ExecuteResponse::SetVariable { .. }
260            | ExecuteResponse::StartedTransaction
261            | ExecuteResponse::TransactionCommitted { .. }
262            | ExecuteResponse::TransactionRolledBack { .. }
263            | ExecuteResponse::Updated(_)
264            | ExecuteResponse::ValidatedConnection { .. } => {
265                StatementEndedExecutionReason::Success {
266                    result_size: None,
267                    rows_returned: None,
268                    execution_strategy: None,
269                }
270            }
271        }
272    }
273}
274
275mod sealed {
276    /// A struct that is purposefully private so folks are forced to use the constructor of an
277    /// enum.
278    #[derive(Debug, Copy, Clone)]
279    pub struct Private;
280}
281
282/// Metadata required for logging a prepared statement.
283#[derive(Debug)]
284pub enum PreparedStatementLoggingInfo {
285    /// The statement has already been logged; we don't need to log it
286    /// again if a future execution hits the sampling rate; we merely
287    /// need to reference the corresponding UUID.
288    AlreadyLogged {
289        uuid: Uuid,
290        kind: Option<StatementKind>,
291    },
292    /// The statement has not yet been logged; if a future execution
293    /// hits the sampling rate, we need to log it at that point.
294    StillToLog {
295        /// The SQL text of the statement.
296        sql: String,
297        /// The SQL text of the statement, redacted to follow our data management
298        /// policy
299        redacted_sql: String,
300        /// When the statement was prepared
301        prepared_at: EpochMillis,
302        /// The name with which the statement was prepared
303        name: String,
304        /// The ID of the session that prepared the statement
305        session_id: Uuid,
306        /// Whether we have already recorded this in the "would have logged" metric
307        accounted: bool,
308        /// The top-level kind of the statement (e.g., `Select`), or `None` for an empty statement
309        kind: Option<StatementKind>,
310
311        /// Private type that forces use of the [`PreparedStatementLoggingInfo::still_to_log`]
312        /// constructor.
313        _sealed: sealed::Private,
314    },
315}
316
317impl PreparedStatementLoggingInfo {
318    /// The kind of the prepared statement, if known. Available regardless of
319    /// whether the statement has already been logged.
320    pub fn kind(&self) -> Option<StatementKind> {
321        match self {
322            PreparedStatementLoggingInfo::StillToLog { kind, .. } => *kind,
323            PreparedStatementLoggingInfo::AlreadyLogged { kind, .. } => *kind,
324        }
325    }
326
327    /// Constructor for the [`PreparedStatementLoggingInfo::StillToLog`] variant that ensures SQL
328    /// statements are properly redacted.
329    pub fn still_to_log<A: AstInfo>(
330        raw_sql: String,
331        stmt: Option<&Statement<A>>,
332        prepared_at: EpochMillis,
333        name: String,
334        session_id: Uuid,
335        accounted: bool,
336    ) -> Self {
337        let kind = stmt.map(StatementKind::from);
338        let sql = match kind {
339            // Redact the SQL text of statements that can carry sensitive material:
340            // secret values (`CREATE`/`ALTER SECRET`), or bulk/PII user data
341            // (`INSERT`/`UPDATE`/`EXECUTE`). See `StatementKind::is_sensitive`.
342            Some(kind) if kind.is_sensitive() => {
343                stmt.map(|s| s.to_ast_string_redacted()).unwrap_or_default()
344            }
345            _ => raw_sql,
346        };
347
348        PreparedStatementLoggingInfo::StillToLog {
349            sql,
350            redacted_sql: stmt.map(|s| s.to_ast_string_redacted()).unwrap_or_default(),
351            prepared_at,
352            name,
353            session_id,
354            accounted,
355            kind,
356            _sealed: sealed::Private,
357        }
358    }
359}
360
361#[derive(Copy, Clone, Debug, Ord, Eq, PartialOrd, PartialEq)]
362pub struct StatementLoggingId(pub Uuid);
363
364/// Rows to be written to `mz_prepared_statement_history` and `mz_sql_text`, with the session id.
365#[derive(Debug, Clone)]
366pub struct PreparedStatementEvent {
367    pub prepared_statement: Row,
368    pub sql_text: Row,
369    pub session_id: Uuid,
370}
371
372/// Throttling state for statement logging, shared across multiple frontend tasks (and currently
373/// also shared with the old peek sequencing).
374#[derive(Debug)]
375pub struct ThrottlingState {
376    /// Inner state protected by a mutex for rate-limiting, because the two inner fields have to be
377    /// manipulated together atomically.
378    /// This mutex is locked once per unsampled query. (There is both sampling and throttling.
379    /// Sampling happens before throttling.) This should be ok for now: Our QPS will not be more
380    /// than 10000s for now, and a mutex should be able to do 100000s of lockings per second, even
381    /// with some contention. If this ever becomes an issue, then we could redesign throttling to be
382    /// per-session/per-tokio-worker-thread.
383    inner: Mutex<ThrottlingStateInner>,
384    /// The number of statements that have been throttled since the last successfully logged
385    /// statement. This is not needed for the throttling decision itself, so it can be a separate
386    /// atomic to allow reading/writing without acquiring the inner mutex.
387    throttled_count: std::sync::atomic::AtomicUsize,
388}
389
390#[derive(Debug)]
391struct ThrottlingStateInner {
392    /// The number of bytes that we are allowed to emit for statement logging without being throttled.
393    /// Increases at a rate of [`mz_sql::session::vars::STATEMENT_LOGGING_TARGET_DATA_RATE`] per second,
394    /// up to a max value of [`mz_sql::session::vars::STATEMENT_LOGGING_MAX_DATA_CREDIT`].
395    tokens: u64,
396    /// The last time at which a statement was logged.
397    last_logged_ts_seconds: u64,
398}
399
400impl ThrottlingState {
401    /// Create a new throttling state.
402    pub fn new(now: &NowFn) -> Self {
403        Self {
404            inner: Mutex::new(ThrottlingStateInner {
405                tokens: 0,
406                last_logged_ts_seconds: now() / 1000,
407            }),
408            throttled_count: std::sync::atomic::AtomicUsize::new(0),
409        }
410    }
411
412    /// Check if we need to drop a statement due to throttling, and update the number of available
413    /// tokens appropriately.
414    ///
415    /// Returns `false` if we must throttle this statement, and `true` otherwise.
416    /// Note: `throttled_count` is NOT modified by this method - callers are responsible
417    /// for incrementing it on throttle failure and resetting it when appropriate.
418    pub fn throttling_check(
419        &self,
420        cost: u64,
421        target_data_rate: u64,
422        max_data_credit: Option<u64>,
423        now: &NowFn,
424    ) -> bool {
425        let ts = now() / 1000;
426        let mut inner = self.inner.lock().expect("throttling state lock poisoned");
427        // We use saturating_sub here because system time isn't monotonic, causing cases
428        // when last_logged_ts_seconds is greater than ts.
429        let elapsed = ts.saturating_sub(inner.last_logged_ts_seconds);
430        inner.last_logged_ts_seconds = ts;
431        inner.tokens = inner
432            .tokens
433            .saturating_add(target_data_rate.saturating_mul(elapsed));
434        if let Some(max_data_credit) = max_data_credit {
435            inner.tokens = inner.tokens.min(max_data_credit);
436        }
437        if let Some(remaining) = inner.tokens.checked_sub(cost) {
438            tracing::debug!("throttling check passed. tokens remaining: {remaining}; cost: {cost}");
439            inner.tokens = remaining;
440            true
441        } else {
442            tracing::debug!(
443                "throttling check failed. tokens available: {}; cost: {cost}",
444                inner.tokens
445            );
446            false
447        }
448    }
449
450    pub fn get_throttled_count(&self) -> usize {
451        self.throttled_count.load(Ordering::Relaxed)
452    }
453
454    pub fn increment_throttled_count(&self) {
455        self.throttled_count.fetch_add(1, Ordering::Relaxed);
456    }
457
458    pub fn reset_throttled_count(&self) {
459        self.throttled_count.store(0, Ordering::Relaxed);
460    }
461}
462
463/// Encapsulates statement logging state needed by the frontend peek sequencing.
464///
465/// This struct bundles together all the statement logging-related state that
466/// the frontend peek sequencing needs to perform statement logging independently
467/// of the Coordinator's main task.
468#[derive(Debug, Clone)]
469pub struct StatementLoggingFrontend {
470    /// Shared throttling state for rate-limiting statement logging.
471    pub throttling_state: Arc<ThrottlingState>,
472    /// Reproducible RNG for statement sampling (only used in tests).
473    pub reproducible_rng: Arc<Mutex<rand_chacha::ChaCha8Rng>>,
474    /// Cached human version string from build info.
475    pub build_info_human_version: String,
476    /// Function to get current time for statement logging.
477    pub now: NowFn,
478}
479
480impl StatementLoggingFrontend {
481    /// Get prepared statement info for frontend peek sequencing.
482    ///
483    /// This function processes prepared statement logging info and builds the event rows.
484    /// It does NOT do throttling - that is handled externally by the caller in `begin_statement_execution`.
485    /// This is a read-only operation that does not mutate
486    /// the `PreparedStatementLoggingInfo` metadata.
487    ///
488    /// # Arguments
489    /// * `session` - The session executing the statement
490    /// * `logging` - Prepared statement logging info
491    ///
492    /// # Returns
493    /// A tuple containing:
494    /// - `Option<PreparedStatementEvent>`: If the prepared statement has not yet been logged,
495    ///   returns the packed rows for the prepared statement.
496    /// - `Uuid`: The UUID of the prepared statement.
497    fn get_prepared_statement_info(
498        &self,
499        session: &Session,
500        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
501    ) -> (Option<PreparedStatementEvent>, Uuid) {
502        let logging_ref = session.qcell_ro(&*logging);
503
504        match logging_ref {
505            PreparedStatementLoggingInfo::AlreadyLogged { uuid, .. } => (None, *uuid),
506            PreparedStatementLoggingInfo::StillToLog {
507                sql,
508                redacted_sql,
509                prepared_at,
510                name,
511                session_id,
512                accounted,
513                kind,
514                _sealed: _,
515            } => {
516                assert!(
517                    *accounted,
518                    "accounting for logging should be done in `begin_statement_execution`"
519                );
520                let uuid = epoch_to_uuid_v7(prepared_at);
521                let sql_hash: [u8; 32] = Sha256::digest(sql.as_bytes()).into();
522
523                let record = StatementPreparedRecord {
524                    id: uuid,
525                    sql_hash,
526                    name: name.clone(),
527                    session_id: *session_id,
528                    prepared_at: *prepared_at,
529                    kind: *kind,
530                };
531
532                // `mz_prepared_statement_history`
533                let mut mpsh_row = Row::default();
534                let mut mpsh_packer = mpsh_row.packer();
535                pack_statement_prepared_update(&record, &mut mpsh_packer);
536
537                // Read throttled_count from shared state
538                let throttled_count = self.throttling_state.get_throttled_count();
539                mpsh_packer.push(Datum::UInt64(CastFrom::cast_from(throttled_count)));
540
541                let sql_row = Row::pack([
542                    Datum::TimestampTz(
543                        to_datetime(*prepared_at)
544                            .truncate_day()
545                            .try_into()
546                            .expect("must fit"),
547                    ),
548                    Datum::Bytes(sql_hash.as_slice()),
549                    Datum::String(sql.as_str()),
550                    Datum::String(redacted_sql.as_str()),
551                ]);
552
553                let prepared_statement_event = PreparedStatementEvent {
554                    prepared_statement: mpsh_row,
555                    sql_text: sql_row,
556                    session_id: *session_id,
557                };
558
559                (Some(prepared_statement_event), uuid)
560            }
561        }
562    }
563
564    /// Marks a prepared statement as "already logged", so future executions only reference its
565    /// UUID rather than logging it again.
566    ///
567    /// This must only be called once we are committed to logging the prepared statement, i.e.
568    /// after the sampling and throttling checks in [`Self::begin_statement_execution`] have
569    /// passed. Mirrors `Coordinator::record_prepared_statement_as_logged` used by the old peek
570    /// sequencing.
571    fn record_prepared_statement_as_logged(
572        &self,
573        uuid: Uuid,
574        session: &mut Session,
575        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
576    ) {
577        let logging = session.qcell_rw(&*logging);
578        if let PreparedStatementLoggingInfo::StillToLog { kind, .. } = logging {
579            let kind = *kind;
580            *logging = PreparedStatementLoggingInfo::AlreadyLogged { uuid, kind };
581        }
582    }
583
584    /// Begin statement execution logging from the frontend. (Corresponds to
585    /// `Coordinator::begin_statement_execution`, which is used by the old peek sequencing.)
586    ///
587    /// This encapsulates all the statement logging setup:
588    /// - Retrieves system config values
589    /// - Performs sampling and throttling checks
590    /// - Creates statement logging records
591    /// - Attends to metrics.
592    ///
593    /// Returns None if the statement should not be logged (due to sampling or throttling), or the
594    /// info required to proceed with statement logging.
595    /// The `Row` is the pre-packed row for `mz_statement_execution_history`.
596    /// The `Option<PreparedStatementEvent>` is None when we have already logged the prepared
597    /// statement before, and this is just a subsequent execution.
598    pub fn begin_statement_execution(
599        &self,
600        session: &mut Session,
601        params: &Params,
602        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
603        system_config: &SystemVars,
604        lifecycle_timestamps: Option<LifecycleTimestamps>,
605    ) -> Option<(
606        StatementLoggingId,
607        StatementBeganExecutionRecord,
608        Row,
609        Option<PreparedStatementEvent>,
610    )> {
611        // Skip logging for internal users unless explicitly enabled
612        let enable_internal_statement_logging = system_config.enable_internal_statement_logging();
613        if session.user().is_internal() && !enable_internal_statement_logging {
614            return None;
615        }
616
617        let sample_rate = effective_sample_rate(session, system_config);
618
619        let use_reproducible_rng = system_config.statement_logging_use_reproducible_rng();
620        let target_data_rate: Option<u64> = system_config
621            .statement_logging_target_data_rate()
622            .map(|rate| rate.cast_into());
623        let max_data_credit: Option<u64> = system_config
624            .statement_logging_max_data_credit()
625            .map(|credit| credit.cast_into());
626
627        // Only lock the RNG when we actually need reproducible sampling (tests only)
628        let sample = if use_reproducible_rng {
629            let mut rng = self.reproducible_rng.lock().expect("rng lock poisoned");
630            should_sample_statement(sample_rate, Some(&mut *rng))
631        } else {
632            should_sample_statement(sample_rate, None)
633        };
634
635        let sampled_label = sample.then_some("true").unwrap_or("false");
636        session
637            .metrics()
638            .statement_logging_records(&[sampled_label])
639            .inc_by(1);
640
641        // Clone only the metrics needed below, before the mutable borrow of session.
642        let unsampled_bytes_metric = session
643            .metrics()
644            .statement_logging_unsampled_bytes()
645            .clone();
646        let actual_bytes_metric = session.metrics().statement_logging_actual_bytes().clone();
647
648        // Handle the accounted flag and record byte metrics
649        let is_new_prepared_statement = if let Some((sql, accounted)) =
650            match session.qcell_rw(logging) {
651                PreparedStatementLoggingInfo::AlreadyLogged { .. } => None,
652                PreparedStatementLoggingInfo::StillToLog { sql, accounted, .. } => {
653                    Some((sql, accounted))
654                }
655            } {
656            if !*accounted {
657                unsampled_bytes_metric.inc_by(u64::cast_from(sql.len()));
658                if sample {
659                    actual_bytes_metric.inc_by(u64::cast_from(sql.len()));
660                }
661                *accounted = true;
662            }
663            true
664        } else {
665            false
666        };
667
668        if !sample {
669            return None;
670        }
671
672        // Capture the statement kind for the began-execution record, before
673        // `record_prepared_statement_as_logged` transitions the logging info to
674        // `AlreadyLogged`.
675        let kind = session.qcell_ro(logging).kind();
676
677        // Get prepared statement info.
678        let (prepared_statement_event, ps_uuid) =
679            self.get_prepared_statement_info(session, logging);
680
681        let began_at = if let Some(lifecycle_timestamps) = lifecycle_timestamps {
682            lifecycle_timestamps.received
683        } else {
684            (self.now)()
685        };
686
687        let current_time = (self.now)();
688        let execution_uuid = epoch_to_uuid_v7(&current_time);
689
690        // Create the execution record
691        let began_execution = create_began_execution_record(
692            execution_uuid,
693            ps_uuid,
694            sample_rate,
695            params,
696            session,
697            began_at,
698            self.build_info_human_version.clone(),
699            kind,
700        );
701
702        // Build rows to calculate cost for throttling
703        let mseh_update = pack_statement_began_execution_update(&began_execution);
704        let maybe_ps_prepared_statement = prepared_statement_event
705            .as_ref()
706            .map(|e| &e.prepared_statement);
707        let maybe_ps_sql_text = prepared_statement_event.as_ref().map(|e| &e.sql_text);
708
709        // Calculate cost of all rows we intend to log
710        let cost: usize = [
711            Some(&mseh_update),
712            maybe_ps_prepared_statement,
713            maybe_ps_sql_text,
714        ]
715        .into_iter()
716        .filter_map(|row_opt| row_opt.map(|row| row.byte_len()))
717        .fold(0_usize, |acc, x| acc.saturating_add(x));
718
719        // Do throttling check
720        let passed = if let Some(target_data_rate) = target_data_rate {
721            self.throttling_state.throttling_check(
722                cost.cast_into(),
723                target_data_rate,
724                max_data_credit,
725                &self.now,
726            )
727        } else {
728            true // No throttling configured
729        };
730
731        if !passed {
732            // Increment throttled_count in shared state
733            self.throttling_state.increment_throttled_count();
734            return None;
735        }
736
737        // Throttling passed, so we are now committed to mark the
738        // prepared statement as logged.
739        self.record_prepared_statement_as_logged(ps_uuid, session, logging);
740
741        // When we successfully log the first instance of a prepared statement
742        // (i.e., it is not throttled), reset the throttled count for future tracking.
743        if is_new_prepared_statement {
744            self.throttling_state.reset_throttled_count();
745        }
746
747        Some((
748            StatementLoggingId(execution_uuid),
749            began_execution,
750            mseh_update,
751            prepared_statement_event,
752        ))
753    }
754}
755
756/// The effective rate at which statement execution should be sampled.
757/// This is the value of the session var `statement_logging_sample_rate`,
758/// constrained by the system var `statement_logging_max_sample_rate`.
759pub(crate) fn effective_sample_rate(session: &Session, system_vars: &SystemVars) -> f64 {
760    let system_max: f64 = system_vars
761        .statement_logging_max_sample_rate()
762        .try_into()
763        .expect("value constrained to be convertible to f64");
764    let user_rate: f64 = session
765        .vars()
766        .get_statement_logging_sample_rate()
767        .try_into()
768        .expect("value constrained to be convertible to f64");
769    f64::min(system_max, user_rate)
770}
771
772/// Helper function to decide whether to sample a statement execution.
773/// Returns `true` if the statement should be sampled based on the sample rate.
774///
775/// If `reproducible_rng` is `Some`, uses the provided RNG for reproducible sampling (used in tests).
776/// If `reproducible_rng` is `None`, uses the thread-local RNG.
777pub(crate) fn should_sample_statement(
778    sample_rate: f64,
779    reproducible_rng: Option<&mut rand_chacha::ChaCha8Rng>,
780) -> bool {
781    let distribution = Bernoulli::new(sample_rate).unwrap_or_else(|_| {
782        soft_panic_or_log!("statement_logging_sample_rate is out of range [0, 1]");
783        Bernoulli::new(0.0).expect("0.0 is valid for Bernoulli")
784    });
785    if let Some(rng) = reproducible_rng {
786        distribution.sample(rng)
787    } else {
788        distribution.sample(&mut rand::rng())
789    }
790}
791
792/// Serializes statement parameters for logging as UTF-8 strings.
793///
794/// Non-UTF-8 wire bytes are lossily replaced with `U+FFFD`.
795fn serialize_params(params: &Params) -> Vec<Option<String>> {
796    std::iter::zip(params.execute_types.iter(), params.datums.iter())
797        .enumerate()
798        .map(|(index, (r#type, datum))| {
799            mz_pgrepr::Value::from_datum(datum, r#type).map(|val| {
800                let mut buf = BytesMut::new();
801                val.encode_text(&mut buf, mz_pgrepr::TextEncodeSettings::STABLE);
802                // NOTE: `encode_text` can emit non-UTF-8 bytes for `"char"`
803                // (`PgLegacyChar`) params, which write their byte verbatim.
804                // Log the raw bytes so operators can recover what came in.
805                match String::from_utf8_lossy(&buf) {
806                    Cow::Borrowed(s) => s.to_owned(),
807                    Cow::Owned(s) => {
808                        let bytes_hex: String = buf.iter().map(|b| format!("{:02x}", b)).collect();
809                        tracing::warn!(
810                            index,
811                            ty = ?r#type,
812                            bytes_hex = %bytes_hex,
813                            "non-UTF-8 bytes in statement-logging param, replaced with U+FFFD"
814                        );
815                        s
816                    }
817                }
818            })
819        })
820        .collect()
821}
822
823/// Helper function to create a `StatementBeganExecutionRecord`.
824pub(crate) fn create_began_execution_record(
825    execution_uuid: Uuid,
826    prepared_statement_uuid: Uuid,
827    sample_rate: f64,
828    params: &Params,
829    session: &Session,
830    began_at: EpochMillis,
831    build_info_version: String,
832    kind: Option<StatementKind>,
833) -> StatementBeganExecutionRecord {
834    let params = serialize_params(params);
835    StatementBeganExecutionRecord {
836        id: execution_uuid,
837        prepared_statement_id: prepared_statement_uuid,
838        sample_rate,
839        params,
840        began_at,
841        application_name: session.application_name().to_string(),
842        transaction_isolation: session.vars().transaction_isolation().to_string(),
843        transaction_id: session
844            .transaction()
845            .inner()
846            .map(|t| t.id)
847            .unwrap_or_else(|| {
848                // This should never happen because every statement runs in an explicit or implicit
849                // transaction.
850                soft_panic_or_log!(
851                    "Statement logging got a statement with no associated transaction"
852                );
853                9999999
854            }),
855        mz_version: build_info_version,
856        kind,
857        // These are not known yet; we'll fill them in later.
858        cluster_id: None,
859        cluster_name: None,
860        execution_timestamp: None,
861        transient_index_id: None,
862        database_name: session.vars().database().into(),
863        search_path: session
864            .vars()
865            .search_path()
866            .iter()
867            .map(|s| s.as_str().to_string())
868            .collect(),
869    }
870}
871
872/// Represents a single statement logging event that can be sent from the frontend
873/// peek sequencing to the Coordinator via an mpsc channel.
874#[derive(Debug, Clone)]
875pub enum FrontendStatementLoggingEvent {
876    /// Statement execution began, possibly with an associated prepared statement
877    /// if this is the first time the prepared statement is being logged
878    BeganExecution {
879        record: StatementBeganExecutionRecord,
880        /// `mz_statement_execution_history`
881        mseh_update: Row,
882        prepared_statement: Option<PreparedStatementEvent>,
883    },
884    /// Statement execution ended
885    EndedExecution(StatementEndedExecutionRecord),
886    /// Set the cluster for a statement execution
887    SetCluster {
888        id: StatementLoggingId,
889        cluster_id: ClusterId,
890        cluster_name: String,
891    },
892    /// Set the execution timestamp for a statement
893    SetTimestamp {
894        id: StatementLoggingId,
895        timestamp: Timestamp,
896    },
897    /// Set the transient index ID for a statement
898    SetTransientIndex {
899        id: StatementLoggingId,
900        transient_index_id: GlobalId,
901    },
902    /// Record a statement lifecycle event
903    Lifecycle {
904        id: StatementLoggingId,
905        event: StatementLifecycleEvent,
906        when: EpochMillis,
907    },
908}
909
910pub(crate) fn pack_statement_execution_inner(
911    record: &StatementBeganExecutionRecord,
912    packer: &mut RowPacker,
913) {
914    let StatementBeganExecutionRecord {
915        id,
916        prepared_statement_id,
917        sample_rate,
918        params,
919        began_at,
920        cluster_id,
921        cluster_name,
922        database_name,
923        search_path,
924        application_name,
925        transaction_isolation,
926        execution_timestamp,
927        transaction_id,
928        transient_index_id,
929        mz_version,
930        // Not packed into a column; only used to redact `error_message`.
931        kind: _,
932    } = record;
933
934    let cluster = cluster_id.map(|id| id.to_string());
935    let transient_index_id = transient_index_id.map(|id| id.to_string());
936    packer.extend([
937        Datum::Uuid(*id),
938        Datum::Uuid(*prepared_statement_id),
939        Datum::Float64((*sample_rate).into()),
940        match &cluster {
941            None => Datum::Null,
942            Some(cluster_id) => Datum::String(cluster_id),
943        },
944        Datum::String(&*application_name),
945        cluster_name.as_ref().map(String::as_str).into(),
946        Datum::String(database_name),
947    ]);
948    packer.push_list(search_path.iter().map(|s| Datum::String(s)));
949    packer.extend([
950        Datum::String(&*transaction_isolation),
951        (*execution_timestamp).into(),
952        Datum::UInt64(*transaction_id),
953        match &transient_index_id {
954            None => Datum::Null,
955            Some(transient_index_id) => Datum::String(transient_index_id),
956        },
957    ]);
958    packer
959        .try_push_array(
960            &[ArrayDimension {
961                lower_bound: 1,
962                length: params.len(),
963            }],
964            params
965                .iter()
966                .map(|p| Datum::from(p.as_ref().map(String::as_str))),
967        )
968        .expect("correct array dimensions");
969    packer.push(Datum::from(mz_version.as_str()));
970    packer.push(Datum::TimestampTz(
971        to_datetime(*began_at).try_into().expect("Sane system time"),
972    ));
973}
974
975pub(crate) fn pack_statement_began_execution_update(record: &StatementBeganExecutionRecord) -> Row {
976    let mut row = Row::default();
977    let mut packer = row.packer();
978    pack_statement_execution_inner(record, &mut packer);
979    packer.extend([
980        // finished_at
981        Datum::Null,
982        // finished_status
983        Datum::Null,
984        // error_message
985        Datum::Null,
986        // result_size
987        Datum::Null,
988        // rows_returned
989        Datum::Null,
990        // execution_status
991        Datum::Null,
992    ]);
993    row
994}
995
996pub(crate) fn pack_statement_prepared_update(
997    record: &StatementPreparedRecord,
998    packer: &mut RowPacker,
999) {
1000    let StatementPreparedRecord {
1001        id,
1002        session_id,
1003        name,
1004        sql_hash,
1005        prepared_at,
1006        kind,
1007    } = record;
1008    packer.extend([
1009        Datum::Uuid(*id),
1010        Datum::Uuid(*session_id),
1011        Datum::String(name.as_str()),
1012        Datum::Bytes(sql_hash.as_slice()),
1013        Datum::TimestampTz(to_datetime(*prepared_at).try_into().expect("must fit")),
1014        kind.map(statement_kind_label_value).into(),
1015    ]);
1016}
1017
1018/// Bundles all information needed to install watch sets for statement lifecycle logging.
1019/// This includes the statement logging ID and the transitive dependencies to watch.
1020#[derive(Debug)]
1021pub struct WatchSetCreation {
1022    /// The statement logging ID for this execution.
1023    pub logging_id: StatementLoggingId,
1024    /// The timestamp at which to watch for dependencies becoming ready.
1025    pub timestamp: Timestamp,
1026    /// Transitive storage dependencies (tables, sources) to watch.
1027    pub storage_ids: BTreeSet<GlobalId>,
1028    /// Transitive compute dependencies (materialized views, indexes) to watch.
1029    pub compute_ids: BTreeSet<GlobalId>,
1030}
1031
1032impl WatchSetCreation {
1033    /// Compute transitive dependencies for watch sets from an input ID bundle, categorized into
1034    /// storage and compute IDs.
1035    pub fn new(
1036        logging_id: StatementLoggingId,
1037        catalog_state: &CatalogState,
1038        input_id_bundle: &CollectionIdBundle,
1039        timestamp: Timestamp,
1040    ) -> Self {
1041        let mut storage_ids = BTreeSet::new();
1042        let mut compute_ids = BTreeSet::new();
1043
1044        for item_id in input_id_bundle
1045            .iter()
1046            .map(|gid| catalog_state.get_entry_by_global_id(&gid).id())
1047            .flat_map(|id| catalog_state.transitive_uses(id))
1048        {
1049            let entry = catalog_state.get_entry(&item_id);
1050            match entry.item() {
1051                // TODO(alter_table): Adding all of the GlobalIds for an object is incorrect.
1052                // For example, this peek may depend on just a single version of a table, but
1053                // we would add dependencies on all versions of said table. Doing this is okay
1054                // for now since we can't yet version tables, but should get fixed.
1055                CatalogItem::Table(_) | CatalogItem::Source(_) => {
1056                    storage_ids.extend(entry.global_ids());
1057                }
1058                // Each catalog item is computed by at most one compute collection at a time,
1059                // which is also the most recent one.
1060                CatalogItem::MaterializedView(_) | CatalogItem::Index(_) => {
1061                    compute_ids.insert(entry.latest_global_id());
1062                }
1063                _ => {}
1064            }
1065        }
1066
1067        Self {
1068            logging_id,
1069            timestamp,
1070            storage_ids,
1071            compute_ids,
1072        }
1073    }
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078    use mz_repr::{Datum, Row, SqlScalarType};
1079    use mz_sql::plan::Params;
1080
1081    use super::serialize_params;
1082
1083    /// A `"char"` param whose byte is `>= 0x80` used to panic
1084    /// `String::from_utf8` on the statement-logging path. It should now
1085    /// round-trip through `from_utf8_lossy` and emit the replacement
1086    /// character instead.
1087    #[mz_ore::test]
1088    fn serialize_params_replaces_non_utf8_char() {
1089        let params = Params {
1090            datums: Row::pack_slice(&[Datum::UInt8(0xFF)]),
1091            execute_types: vec![SqlScalarType::PgLegacyChar],
1092            expected_types: vec![SqlScalarType::PgLegacyChar],
1093        };
1094        let out = serialize_params(&params);
1095        assert_eq!(out.len(), 1);
1096        assert_eq!(out[0].as_deref(), Some("\u{FFFD}"));
1097    }
1098
1099    /// An ASCII `"char"` param must render as its ASCII character, unchanged.
1100    #[mz_ore::test]
1101    fn serialize_params_ascii_char_unchanged() {
1102        let params = Params {
1103            datums: Row::pack_slice(&[Datum::UInt8(b'A')]),
1104            execute_types: vec![SqlScalarType::PgLegacyChar],
1105            expected_types: vec![SqlScalarType::PgLegacyChar],
1106        };
1107        let out = serialize_params(&params);
1108        assert_eq!(out[0].as_deref(), Some("A"));
1109    }
1110}
mz_adapter/statement_logging.rs

mz_adapter/
statement_logging.rs
