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

mz_adapter/
statement_logging.rs
