mz_adapter/coord/
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::collections::BTreeMap;
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14use mz_adapter_types::connection::ConnectionId;
15use mz_compute_client::controller::error::CollectionLookupError;
16use mz_controller_types::ClusterId;
17use mz_ore::now::{EpochMillis, NowFn, epoch_to_uuid_v7, to_datetime};
18use mz_ore::task::spawn;
19use mz_ore::{cast::CastFrom, cast::CastInto};
20use mz_repr::adt::timestamp::TimestampLike;
21use mz_repr::{Datum, Diff, GlobalId, Row, Timestamp};
22use mz_sql::plan::Params;
23use mz_sql::session::metadata::SessionMetadata;
24use mz_storage_client::controller::IntrospectionType;
25use qcell::QCell;
26use rand::SeedableRng;
27use sha2::{Digest, Sha256};
28use tokio::time::MissedTickBehavior;
29use uuid::Uuid;
30
31use crate::coord::{ConnMeta, Coordinator, WatchSetResponse};
32use crate::session::{LifecycleTimestamps, Session};
33use crate::statement_logging::{
34    FrontendStatementLoggingEvent, PreparedStatementEvent, PreparedStatementLoggingInfo,
35    SessionHistoryEvent, StatementBeganExecutionRecord, StatementEndedExecutionReason,
36    StatementEndedExecutionRecord, StatementLifecycleEvent, StatementLoggingFrontend,
37    StatementLoggingId, StatementPreparedRecord, ThrottlingState, WatchSetCreation,
38    create_began_execution_record, effective_sample_rate, pack_statement_began_execution_update,
39    pack_statement_execution_inner, pack_statement_prepared_update, should_sample_statement,
40};
41
42use super::Message;
43
44/// Statement logging state in the Coordinator.
45#[derive(Debug)]
46pub(crate) struct StatementLogging {
47    /// Information about statement executions that have been logged
48    /// but not finished.
49    ///
50    /// This map needs to have enough state left over to later retract
51    /// the system table entries (so that we can update them when the
52    /// execution finished.)
53    executions_begun: BTreeMap<Uuid, StatementBeganExecutionRecord>,
54
55    /// Information about sessions that have been started, but which
56    /// have not yet been logged in `mz_session_history`.
57    /// They may be logged as part of a statement being executed (and chosen for logging).
58    unlogged_sessions: BTreeMap<Uuid, SessionHistoryEvent>,
59
60    /// A reproducible RNG for deciding whether to sample statement executions.
61    /// Only used by tests; otherwise, `rand::rng()` is used.
62    /// Controlled by the system var `statement_logging_use_reproducible_rng`.
63    /// This same instance will be used by all frontend tasks.
64    reproducible_rng: Arc<Mutex<rand::rngs::ChaCha8Rng>>,
65
66    /// Events to be persisted periodically.
67    pending_statement_execution_events: Vec<(Row, Diff)>,
68    pending_prepared_statement_events: Vec<PreparedStatementEvent>,
69    pending_session_events: Vec<Row>,
70    pending_statement_lifecycle_events: Vec<Row>,
71
72    /// Shared throttling state for rate-limiting statement logging.
73    pub(crate) throttling_state: Arc<ThrottlingState>,
74
75    /// Function to get the current time.
76    pub(crate) now: NowFn,
77}
78
79impl StatementLogging {
80    const REPRODUCIBLE_RNG_SEED: u64 = 42;
81
82    pub(crate) fn new(now: NowFn) -> Self {
83        Self {
84            executions_begun: BTreeMap::new(),
85            unlogged_sessions: BTreeMap::new(),
86            reproducible_rng: Arc::new(Mutex::new(rand::rngs::ChaCha8Rng::seed_from_u64(
87                Self::REPRODUCIBLE_RNG_SEED,
88            ))),
89            pending_statement_execution_events: Vec::new(),
90            pending_prepared_statement_events: Vec::new(),
91            pending_session_events: Vec::new(),
92            pending_statement_lifecycle_events: Vec::new(),
93            throttling_state: Arc::new(ThrottlingState::new(&now)),
94            now,
95        }
96    }
97
98    /// Create a `StatementLoggingFrontend` for use by frontend peek sequencing.
99    ///
100    /// This provides the frontend with all the state it needs to perform statement
101    /// logging without direct access to the Coordinator.
102    pub(crate) fn create_frontend(
103        &self,
104        build_info_human_version: String,
105    ) -> StatementLoggingFrontend {
106        StatementLoggingFrontend {
107            throttling_state: Arc::clone(&self.throttling_state),
108            reproducible_rng: Arc::clone(&self.reproducible_rng),
109            build_info_human_version,
110            now: self.now.clone(),
111        }
112    }
113}
114
115impl Coordinator {
116    /// Helper to write began execution events to pending buffers.
117    fn write_began_execution_events(
118        &mut self,
119        record: StatementBeganExecutionRecord,
120        mseh_update: Row,
121        prepared_statement: Option<PreparedStatementEvent>,
122    ) {
123        // `mz_statement_execution_history`
124        self.statement_logging
125            .pending_statement_execution_events
126            .push((mseh_update, Diff::ONE));
127
128        // Track the execution for later updates
129        self.statement_logging
130            .executions_begun
131            .insert(record.id, record);
132
133        // If we have a prepared statement, log it and possibly its session
134        if let Some(ps_event) = prepared_statement {
135            let session_id = ps_event.session_id;
136            self.statement_logging
137                .pending_prepared_statement_events
138                .push(ps_event);
139
140            // Check if we need to log the session for this prepared statement
141            if let Some(sh) = self.statement_logging.unlogged_sessions.remove(&session_id) {
142                let sh_update = Self::pack_session_history_update(&sh);
143                self.statement_logging
144                    .pending_session_events
145                    .push(sh_update);
146            }
147        }
148    }
149
150    /// Handle a statement logging event from frontend peek sequencing.
151    pub(crate) fn handle_frontend_statement_logging_event(
152        &mut self,
153        event: FrontendStatementLoggingEvent,
154    ) {
155        match event {
156            FrontendStatementLoggingEvent::BeganExecution {
157                record,
158                mseh_update,
159                prepared_statement,
160            } => {
161                self.record_statement_lifecycle_event(
162                    &StatementLoggingId(record.id),
163                    &StatementLifecycleEvent::ExecutionBegan,
164                    record.began_at,
165                );
166                self.write_began_execution_events(record, mseh_update, prepared_statement);
167            }
168            FrontendStatementLoggingEvent::EndedExecution(ended_record) => {
169                self.end_statement_execution(
170                    StatementLoggingId(ended_record.id),
171                    ended_record.reason,
172                    ended_record.ended_at,
173                );
174            }
175            FrontendStatementLoggingEvent::SetCluster {
176                id,
177                cluster_id,
178                cluster_name,
179            } => {
180                self.set_statement_execution_cluster(id, cluster_id, cluster_name);
181            }
182            FrontendStatementLoggingEvent::SetTimestamp { id, timestamp } => {
183                self.set_statement_execution_timestamp(id, timestamp);
184            }
185            FrontendStatementLoggingEvent::SetTransientIndex {
186                id,
187                transient_index_id,
188            } => {
189                self.set_transient_index_id(id, transient_index_id);
190            }
191            FrontendStatementLoggingEvent::Lifecycle { id, event, when } => {
192                self.record_statement_lifecycle_event(&id, &event, when);
193            }
194        }
195    }
196
197    // TODO[btv] make this configurable via LD?
198    // Although... Logging every 5 seconds seems like it
199    // should have acceptable cost for now, since we do a
200    // group commit for tables every 1s anyway.
201    const STATEMENT_LOGGING_WRITE_INTERVAL: Duration = Duration::from_secs(5);
202
203    pub(crate) fn spawn_statement_logging_task(&self) {
204        let internal_cmd_tx = self.internal_cmd_tx.clone();
205        spawn(|| "statement_logging", async move {
206            let mut interval = tokio::time::interval(Coordinator::STATEMENT_LOGGING_WRITE_INTERVAL);
207            interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
208            loop {
209                interval.tick().await;
210                let _ = internal_cmd_tx.send(Message::DrainStatementLog);
211            }
212        });
213    }
214
215    #[mz_ore::instrument(level = "debug")]
216    pub(crate) fn drain_statement_log(&mut self) {
217        let session_updates = std::mem::take(&mut self.statement_logging.pending_session_events)
218            .into_iter()
219            .map(|update| (update, Diff::ONE))
220            .collect();
221        let (prepared_statement_updates, sql_text_updates) =
222            std::mem::take(&mut self.statement_logging.pending_prepared_statement_events)
223                .into_iter()
224                .map(
225                    |PreparedStatementEvent {
226                         prepared_statement,
227                         sql_text,
228                         ..
229                     }| {
230                        ((prepared_statement, Diff::ONE), (sql_text, Diff::ONE))
231                    },
232                )
233                .unzip::<_, _, Vec<_>, Vec<_>>();
234        let statement_execution_updates =
235            std::mem::take(&mut self.statement_logging.pending_statement_execution_events);
236        let statement_lifecycle_updates =
237            std::mem::take(&mut self.statement_logging.pending_statement_lifecycle_events)
238                .into_iter()
239                .map(|update| (update, Diff::ONE))
240                .collect();
241
242        use IntrospectionType::*;
243        for (type_, updates) in [
244            (SessionHistory, session_updates),
245            (PreparedStatementHistory, prepared_statement_updates),
246            (StatementExecutionHistory, statement_execution_updates),
247            (StatementLifecycleHistory, statement_lifecycle_updates),
248            (SqlText, sql_text_updates),
249        ] {
250            if !updates.is_empty() && !self.controller.read_only() {
251                self.controller
252                    .storage
253                    .append_introspection_updates(type_, updates);
254            }
255        }
256    }
257
258    /// Check whether we need to do throttling (i.e., whether `STATEMENT_LOGGING_TARGET_DATA_RATE` is set).
259    /// If so, actually do the check.
260    ///
261    /// We expect `rows` to be the list of rows we intend to record and calculate the cost by summing the
262    /// byte lengths of the rows.
263    ///
264    /// Returns `false` if we must throttle this statement, and `true` otherwise.
265    fn statement_logging_throttling_check<'a, I>(&self, rows: I) -> bool
266    where
267        I: IntoIterator<Item = Option<&'a Row>>,
268    {
269        let cost = rows
270            .into_iter()
271            .filter_map(|row_opt| row_opt.map(|row| row.byte_len()))
272            .fold(0_usize, |acc, x| acc.saturating_add(x));
273
274        let Some(target_data_rate) = self
275            .catalog
276            .system_config()
277            .statement_logging_target_data_rate()
278        else {
279            return true;
280        };
281        let max_data_credit = self
282            .catalog
283            .system_config()
284            .statement_logging_max_data_credit();
285
286        self.statement_logging.throttling_state.throttling_check(
287            cost.cast_into(),
288            target_data_rate.cast_into(),
289            max_data_credit.map(CastInto::cast_into),
290            &self.statement_logging.now,
291        )
292    }
293
294    /// Marks a prepared statement as "already logged".
295    /// Mutates the `PreparedStatementLoggingInfo` metadata.
296    fn record_prepared_statement_as_logged(
297        &self,
298        uuid: Uuid,
299        session: &mut Session,
300        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
301    ) {
302        let logging = session.qcell_rw(&*logging);
303        if let PreparedStatementLoggingInfo::StillToLog { kind, .. } = logging {
304            let kind = *kind;
305            *logging = PreparedStatementLoggingInfo::AlreadyLogged { uuid, kind };
306        }
307    }
308
309    /// Returns any statement logging events needed for a particular
310    /// prepared statement. This is a read-only operation that does not mutate
311    /// the `PreparedStatementLoggingInfo` metadata.
312    ///
313    /// This function does not do a sampling check, and assumes we did so in a higher layer.
314    /// It also does not do a throttling check - that is done separately in `begin_statement_execution`.
315    ///
316    /// Returns a tuple containing:
317    /// - `Option<(StatementPreparedRecord, PreparedStatementEvent)>`: If the prepared statement
318    ///   has not yet been logged, returns the prepared statement record and the packed rows.
319    /// - `Uuid`: The UUID of the prepared statement.
320    pub(crate) fn get_prepared_statement_info(
321        &self,
322        session: &Session,
323        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
324    ) -> (
325        Option<(StatementPreparedRecord, PreparedStatementEvent)>,
326        Uuid,
327    ) {
328        let logging = session.qcell_ro(&*logging);
329
330        match logging {
331            PreparedStatementLoggingInfo::AlreadyLogged { uuid, .. } => (None, *uuid),
332            PreparedStatementLoggingInfo::StillToLog {
333                sql,
334                redacted_sql,
335                prepared_at,
336                name,
337                session_id,
338                accounted,
339                kind,
340                _sealed: _,
341            } => {
342                assert!(
343                    *accounted,
344                    "accounting for logging should be done in `begin_statement_execution`"
345                );
346                let uuid = epoch_to_uuid_v7(prepared_at);
347                let sql_hash: [u8; 32] = Sha256::digest(sql.as_bytes()).into();
348                let record = StatementPreparedRecord {
349                    id: uuid,
350                    sql_hash,
351                    name: name.to_string(),
352                    session_id: *session_id,
353                    prepared_at: *prepared_at,
354                    kind: *kind,
355                };
356
357                // `mz_prepared_statement_history`
358                let mut mpsh_row = Row::default();
359                let mut mpsh_packer = mpsh_row.packer();
360                pack_statement_prepared_update(&record, &mut mpsh_packer);
361                let throttled_count = self
362                    .statement_logging
363                    .throttling_state
364                    .get_throttled_count();
365                mpsh_packer.push(Datum::UInt64(CastFrom::cast_from(throttled_count)));
366
367                let sql_row = Row::pack([
368                    Datum::TimestampTz(
369                        to_datetime(*prepared_at)
370                            .truncate_day()
371                            .try_into()
372                            .expect("must fit"),
373                    ),
374                    Datum::Bytes(sql_hash.as_slice()),
375                    Datum::String(sql.as_str()),
376                    Datum::String(redacted_sql.as_str()),
377                ]);
378
379                (
380                    Some((
381                        record,
382                        PreparedStatementEvent {
383                            prepared_statement: mpsh_row,
384                            sql_text: sql_row,
385                            session_id: *session_id,
386                        },
387                    )),
388                    uuid,
389                )
390            }
391        }
392    }
393
394    /// Record the end of statement execution for a statement whose beginning
395    /// was logged. Ends are idempotent: the first end wins and later ends for
396    /// the same statement are ignored.
397    ///
398    /// `ended_at` is when execution finished, which is not the same as when this
399    /// runs. A statement sequenced off the coordinator loop finishes in its
400    /// session task and reports the end as a message, so taking the clock here
401    /// would charge it for however long that message sat in the queue.
402    pub(crate) fn end_statement_execution(
403        &mut self,
404        id: StatementLoggingId,
405        reason: StatementEndedExecutionReason,
406        ended_at: EpochMillis,
407    ) {
408        let StatementLoggingId(uuid) = id;
409        let ended_record = StatementEndedExecutionRecord {
410            id: uuid,
411            reason,
412            ended_at,
413        };
414
415        let Some(began_record) = self.statement_logging.executions_begun.remove(&uuid) else {
416            // The statement was already ended; the first end wins.
417            //
418            // A missing entry can only mean a duplicate end, never an end that
419            // overtook its begin: a StatementLoggingId is only minted when a
420            // begin is logged, and begins travel the same FIFO command channel
421            // as the commands that hand statements to the coordinator, so the
422            // coordinator never ends a statement before processing its begin.
423            //
424            // Duplicate ends are legitimate, if rare. Ownership of the end is
425            // handed from the frontend to the coordinator while a statement is
426            // dispatched, and async cancellation can strike mid-handoff: if a
427            // client disconnect drops the frontend future after the
428            // coordinator registered a peek but before the frontend defused
429            // its logging guard, both sides own the end and both emit one.
430            tracing::warn!(
431                statement_uuid = %uuid,
432                reason = ?ended_record.reason,
433                "duplicate end_statement_execution, keeping the first end",
434            );
435            return;
436        };
437        for (row, diff) in
438            Self::pack_statement_ended_execution_updates(&began_record, &ended_record)
439        {
440            self.statement_logging
441                .pending_statement_execution_events
442                .push((row, diff));
443        }
444        self.record_statement_lifecycle_event(
445            &id,
446            &StatementLifecycleEvent::ExecutionFinished,
447            ended_at,
448        );
449    }
450
451    fn pack_session_history_update(event: &SessionHistoryEvent) -> Row {
452        let SessionHistoryEvent {
453            id,
454            connected_at,
455            application_name,
456            authenticated_user,
457        } = event;
458        Row::pack_slice(&[
459            Datum::Uuid(*id),
460            Datum::TimestampTz(to_datetime(*connected_at).try_into().expect("must fit")),
461            Datum::String(&*application_name),
462            Datum::String(&*authenticated_user),
463        ])
464    }
465
466    fn pack_statement_lifecycle_event(
467        StatementLoggingId(uuid): &StatementLoggingId,
468        event: &StatementLifecycleEvent,
469        when: EpochMillis,
470    ) -> Row {
471        Row::pack_slice(&[
472            Datum::Uuid(*uuid),
473            Datum::String(event.as_str()),
474            Datum::TimestampTz(to_datetime(when).try_into().expect("must fit")),
475        ])
476    }
477
478    fn pack_full_statement_execution_update(
479        began_record: &StatementBeganExecutionRecord,
480        ended_record: &StatementEndedExecutionRecord,
481    ) -> Row {
482        let mut row = Row::default();
483        let mut packer = row.packer();
484        pack_statement_execution_inner(began_record, &mut packer);
485        let (status, error_message, result_size, rows_returned, execution_strategy) =
486            match &ended_record.reason {
487                StatementEndedExecutionReason::Success {
488                    result_size,
489                    rows_returned,
490                    execution_strategy,
491                } => (
492                    "success",
493                    None,
494                    result_size.map(|rs| i64::try_from(rs).expect("must fit")),
495                    rows_returned.map(|rr| i64::try_from(rr).expect("must fit")),
496                    execution_strategy.map(|es| es.name()),
497                ),
498                StatementEndedExecutionReason::Canceled => ("canceled", None, None, None, None),
499                StatementEndedExecutionReason::Errored { error } => {
500                    // Backstop for SQL-435: `CREATE SECRET`/`ALTER SECRET` errors
501                    // can embed secret material (the rejected statement text, or a
502                    // value-bearing eval error). The per-site fixes redact the known
503                    // cases, but we never persist an unredacted error for these
504                    // kinds, regardless of which error variant fired. The client
505                    // still receives the real error over pgwire.
506                    let error = if began_record.kind.is_some_and(|kind| kind.is_secret()) {
507                        "<error redacted for secret statement>"
508                    } else {
509                        error.as_str()
510                    };
511                    ("error", Some(error), None, None, None)
512                }
513                StatementEndedExecutionReason::Aborted => ("aborted", None, None, None, None),
514            };
515        packer.extend([
516            Datum::TimestampTz(
517                to_datetime(ended_record.ended_at)
518                    .try_into()
519                    .expect("Sane system time"),
520            ),
521            status.into(),
522            error_message.into(),
523            result_size.into(),
524            rows_returned.into(),
525            execution_strategy.into(),
526        ]);
527        row
528    }
529
530    fn pack_statement_ended_execution_updates(
531        began_record: &StatementBeganExecutionRecord,
532        ended_record: &StatementEndedExecutionRecord,
533    ) -> [(Row, Diff); 2] {
534        let retraction = pack_statement_began_execution_update(began_record);
535        let new = Self::pack_full_statement_execution_update(began_record, ended_record);
536        [(retraction, Diff::MINUS_ONE), (new, Diff::ONE)]
537    }
538
539    /// Mutate a statement execution record via the given function `f`.
540    fn mutate_record<F: FnOnce(&mut StatementBeganExecutionRecord)>(
541        &mut self,
542        StatementLoggingId(id): StatementLoggingId,
543        f: F,
544    ) {
545        let record = self
546            .statement_logging
547            .executions_begun
548            .get_mut(&id)
549            .expect("mutate_record must not be called after execution ends");
550        let retraction = pack_statement_began_execution_update(record);
551        self.statement_logging
552            .pending_statement_execution_events
553            .push((retraction, Diff::MINUS_ONE));
554        f(record);
555        let update = pack_statement_began_execution_update(record);
556        self.statement_logging
557            .pending_statement_execution_events
558            .push((update, Diff::ONE));
559    }
560
561    /// Set the `cluster_id` and `cluster_name` for a statement, once they're known.
562    ///
563    /// The name is resolved by the caller (from the same catalog snapshot that selected the
564    /// cluster), rather than re-resolved here from `cluster_id`. This avoids a panic when the
565    /// cluster was concurrently dropped, and logs the name as it was at selection time.
566    ///
567    /// TODO(peek-seq): We could do the packing in the frontend task, and just send over the rows.
568    pub(crate) fn set_statement_execution_cluster(
569        &mut self,
570        id: StatementLoggingId,
571        cluster_id: ClusterId,
572        cluster_name: String,
573    ) {
574        self.mutate_record(id, |record| {
575            record.cluster_name = Some(cluster_name);
576            record.cluster_id = Some(cluster_id);
577        });
578    }
579
580    /// Sets the execution timestamp if the statement is still active.
581    ///
582    /// A duplicate end can race the asynchronous group-commit update, so an ended statement is
583    /// skipped rather than treated as corruption.
584    pub(crate) fn set_statement_execution_timestamp(
585        &mut self,
586        id: StatementLoggingId,
587        timestamp: Timestamp,
588    ) {
589        let StatementLoggingId(uuid) = id;
590        if !self.statement_logging.executions_begun.contains_key(&uuid) {
591            tracing::warn!(
592                statement_uuid = %uuid,
593                "execution already ended, skipping execution timestamp update",
594            );
595            return;
596        }
597        self.mutate_record(id, |record| {
598            record.execution_timestamp = Some(u64::from(timestamp));
599        });
600    }
601
602    pub(crate) fn set_transient_index_id(
603        &mut self,
604        id: StatementLoggingId,
605        transient_index_id: GlobalId,
606    ) {
607        self.mutate_record(id, |record| {
608            record.transient_index_id = Some(transient_index_id)
609        });
610    }
611
612    /// Possibly record the beginning of statement execution, depending on a randomly-chosen value.
613    /// If the execution beginning was indeed logged, returns a `StatementLoggingId` that must be
614    /// passed to `end_statement_execution` to record when it ends.
615    ///
616    /// `lifecycle_timestamps` has timestamps that come from the Adapter frontend (`mz-pgwire`) part
617    /// of the lifecycle.
618    pub(crate) fn begin_statement_execution(
619        &mut self,
620        session: &mut Session,
621        params: &Params,
622        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
623        lifecycle_timestamps: Option<LifecycleTimestamps>,
624    ) -> Option<StatementLoggingId> {
625        let enable_internal_statement_logging = self
626            .catalog()
627            .system_config()
628            .enable_internal_statement_logging();
629        if session.user().is_internal() && !enable_internal_statement_logging {
630            return None;
631        }
632
633        let sample_rate = effective_sample_rate(session, self.catalog().system_config());
634        let use_reproducible_rng = self
635            .catalog()
636            .system_config()
637            .statement_logging_use_reproducible_rng();
638        // Only lock the RNG when we actually need reproducible sampling (tests only)
639        let sample = if use_reproducible_rng {
640            let mut rng = self
641                .statement_logging
642                .reproducible_rng
643                .lock()
644                .expect("rng lock poisoned");
645            should_sample_statement(sample_rate, Some(&mut *rng))
646        } else {
647            should_sample_statement(sample_rate, None)
648        };
649
650        // Figure out the cost of everything before we log.
651
652        // Track how many statements we're recording.
653        let sampled_label = sample.then_some("true").unwrap_or("false");
654        self.metrics
655            .statement_logging_records
656            .with_label_values(&[sampled_label])
657            .inc_by(1);
658
659        if let Some((sql, accounted)) = match session.qcell_rw(logging) {
660            PreparedStatementLoggingInfo::AlreadyLogged { .. } => None,
661            PreparedStatementLoggingInfo::StillToLog { sql, accounted, .. } => {
662                Some((sql, accounted))
663            }
664        } {
665            if !*accounted {
666                self.metrics
667                    .statement_logging_unsampled_bytes
668                    .inc_by(u64::cast_from(sql.len()));
669                if sample {
670                    self.metrics
671                        .statement_logging_actual_bytes
672                        .inc_by(u64::cast_from(sql.len()));
673                }
674                *accounted = true;
675            }
676        }
677        if !sample {
678            return None;
679        }
680
681        let (maybe_ps, ps_uuid) = self.get_prepared_statement_info(session, logging);
682
683        let began_at = if let Some(lifecycle_timestamps) = lifecycle_timestamps {
684            lifecycle_timestamps.received
685        } else {
686            self.now()
687        };
688        let now = self.now();
689        let execution_uuid = epoch_to_uuid_v7(&now);
690
691        let build_info_version = self
692            .catalog()
693            .state()
694            .config()
695            .build_info
696            .human_version(None);
697        let kind = session.qcell_ro(logging).kind();
698        let record = create_began_execution_record(
699            execution_uuid,
700            ps_uuid,
701            sample_rate,
702            params,
703            session,
704            began_at,
705            build_info_version,
706            kind,
707        );
708
709        // `mz_statement_execution_history`
710        let mseh_update = pack_statement_began_execution_update(&record);
711
712        let (maybe_ps_event, maybe_sh_event) = if let Some((ps_record, ps_event)) = maybe_ps {
713            if let Some(sh) = self
714                .statement_logging
715                .unlogged_sessions
716                .get(&ps_record.session_id)
717            {
718                (
719                    Some(ps_event),
720                    Some((Self::pack_session_history_update(sh), ps_record.session_id)),
721                )
722            } else {
723                (Some(ps_event), None)
724            }
725        } else {
726            (None, None)
727        };
728
729        let maybe_ps_prepared_statement = maybe_ps_event.as_ref().map(|e| &e.prepared_statement);
730        let maybe_ps_sql_text = maybe_ps_event.as_ref().map(|e| &e.sql_text);
731
732        if !self.statement_logging_throttling_check([
733            Some(&mseh_update),
734            maybe_ps_prepared_statement,
735            maybe_ps_sql_text,
736            maybe_sh_event.as_ref().map(|(row, _)| row),
737        ]) {
738            // Increment throttled_count in shared state
739            self.statement_logging
740                .throttling_state
741                .increment_throttled_count();
742            return None;
743        }
744        // When we successfully log the first instance of a prepared statement
745        // (i.e., it is not throttled), we also capture the number of previously
746        // throttled statement executions in the builtin prepared statement history table above,
747        // and then reset the throttled count for future tracking.
748        else if let PreparedStatementLoggingInfo::StillToLog { .. } = session.qcell_ro(logging) {
749            self.statement_logging
750                .throttling_state
751                .reset_throttled_count();
752        }
753
754        self.record_prepared_statement_as_logged(ps_uuid, session, logging);
755
756        self.record_statement_lifecycle_event(
757            &StatementLoggingId(execution_uuid),
758            &StatementLifecycleEvent::ExecutionBegan,
759            began_at,
760        );
761
762        self.statement_logging
763            .pending_statement_execution_events
764            .push((mseh_update, Diff::ONE));
765        self.statement_logging
766            .executions_begun
767            .insert(execution_uuid, record);
768
769        if let Some((sh_update, session_id)) = maybe_sh_event {
770            self.statement_logging
771                .pending_session_events
772                .push(sh_update);
773            // Mark the session as logged to avoid logging it again in the future
774            self.statement_logging.unlogged_sessions.remove(&session_id);
775        }
776        if let Some(ps_event) = maybe_ps_event {
777            self.statement_logging
778                .pending_prepared_statement_events
779                .push(ps_event);
780        }
781
782        Some(StatementLoggingId(execution_uuid))
783    }
784
785    /// Record a new connection event
786    pub(crate) fn begin_session_for_statement_logging(&mut self, session: &ConnMeta) {
787        let id = session.uuid();
788        let session_role = session.authenticated_role_id();
789        let event = SessionHistoryEvent {
790            id,
791            connected_at: session.connected_at(),
792            application_name: session.application_name().to_owned(),
793            authenticated_user: self.catalog.get_role(session_role).name.clone(),
794        };
795        self.statement_logging.unlogged_sessions.insert(id, event);
796    }
797
798    pub(crate) fn end_session_for_statement_logging(&mut self, uuid: Uuid) {
799        self.statement_logging.unlogged_sessions.remove(&uuid);
800    }
801
802    pub(crate) fn record_statement_lifecycle_event(
803        &mut self,
804        id: &StatementLoggingId,
805        event: &StatementLifecycleEvent,
806        when: EpochMillis,
807    ) {
808        if mz_adapter_types::dyncfgs::ENABLE_STATEMENT_LIFECYCLE_LOGGING
809            .get(self.catalog().system_config().dyncfgs())
810        {
811            let row = Self::pack_statement_lifecycle_event(id, event, when);
812            self.statement_logging
813                .pending_statement_lifecycle_events
814                .push(row);
815        }
816    }
817
818    /// Install watch sets for statement lifecycle logging.
819    ///
820    /// This installs both storage and compute watch sets that will fire
821    /// `StatementLifecycleEvent::StorageDependenciesFinished` and
822    /// `StatementLifecycleEvent::ComputeDependenciesFinished` respectively
823    /// when the dependencies are ready at the given timestamp.
824    pub(crate) fn install_peek_watch_sets(
825        &mut self,
826        conn_id: ConnectionId,
827        watch_set: WatchSetCreation,
828    ) -> Result<(), CollectionLookupError> {
829        let WatchSetCreation {
830            logging_id,
831            timestamp,
832            storage_ids,
833            compute_ids,
834        } = watch_set;
835
836        self.install_storage_watch_set(
837            conn_id.clone(),
838            storage_ids,
839            timestamp,
840            WatchSetResponse::StatementDependenciesReady(
841                logging_id,
842                StatementLifecycleEvent::StorageDependenciesFinished,
843            ),
844        )?;
845        self.install_compute_watch_set(
846            conn_id,
847            compute_ids,
848            timestamp,
849            WatchSetResponse::StatementDependenciesReady(
850                logging_id,
851                StatementLifecycleEvent::ComputeDependenciesFinished,
852            ),
853        )?;
854        Ok(())
855    }
856}
mz_adapter/coord/statement_logging.rs

mz_adapter/coord/
statement_logging.rs
