risingwave_frontend/handler/
util.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use core::str::FromStr;
16use std::pin::Pin;
17use std::sync::Arc;
18use std::task::{Context, Poll};
19
20use anyhow::Context as _;
21use bytes::{Bytes, BytesMut};
22use futures::Stream;
23use itertools::Itertools;
24use pgwire::pg_field_descriptor::PgFieldDescriptor;
25use pgwire::pg_response::RowSetResult;
26use pgwire::pg_server::BoxedError;
27use pgwire::types::{Format, FormatIterator, Row};
28use pin_project_lite::pin_project;
29use risingwave_common::array::DataChunk;
30use risingwave_common::catalog::Field;
31use risingwave_common::config::FrontendConfig;
32use risingwave_common::id::ObjectId;
33use risingwave_common::row::Row as _;
34use risingwave_common::types::{
35    DataType, Interval, ScalarRefImpl, Timestamptz, write_date_time_tz,
36};
37use risingwave_common::util::epoch::Epoch;
38use risingwave_common::util::iter_util::ZipEqFast;
39use risingwave_connector::sink::elasticsearch_opensearch::elasticsearch::ES_SINK;
40use risingwave_connector::sink::file_sink::fs::FS_SINK;
41use risingwave_connector::source::iceberg::ICEBERG_CONNECTOR;
42use risingwave_connector::source::{BATCH_POSIX_FS_CONNECTOR, KAFKA_CONNECTOR, POSIX_FS_CONNECTOR};
43use risingwave_pb::catalog::connection_params::PbConnectionType;
44use risingwave_sqlparser::ast::{
45    CompatibleFormatEncode, FormatEncodeOptions, ObjectName, Query, Select, SelectItem, SetExpr,
46    TableFactor, TableWithJoins,
47};
48use thiserror_ext::AsReport;
49use tokio::select;
50use tokio::time::{Duration, sleep};
51
52use crate::catalog::root_catalog::SchemaPath;
53use crate::error::ErrorCode::ProtocolError;
54use crate::error::{ErrorCode, Result as RwResult, RwError};
55use crate::session::SessionImpl;
56use crate::{Binder, HashSet, TableCatalog};
57
58pub fn ensure_local_fs_connector_allowed(session: &SessionImpl, connector: &str) -> RwResult<()> {
59    let is_local_fs_connector = is_local_fs_connector(connector);
60
61    if !is_local_fs_connector
62        || is_local_fs_connector_enabled(session.env().frontend_config(), connector)
63    {
64        return Ok(());
65    }
66
67    Err(RwError::from(ProtocolError(format!(
68        "local filesystem connector '{}' is disabled. Set `frontend.unsafe_enable_local_fs_connector = true` in `risingwave.toml` to enable it.",
69        connector
70    ))))
71}
72
73fn is_local_fs_connector_enabled(frontend_config: &FrontendConfig, connector: &str) -> bool {
74    !is_local_fs_connector(connector) || frontend_config.unsafe_enable_local_fs_connector
75}
76
77fn is_local_fs_connector(connector: &str) -> bool {
78    connector.eq_ignore_ascii_case(POSIX_FS_CONNECTOR)
79        || connector.eq_ignore_ascii_case(BATCH_POSIX_FS_CONNECTOR)
80        || connector.eq_ignore_ascii_case(FS_SINK)
81}
82
83pin_project! {
84    /// Wrapper struct that converts a stream of DataChunk to a stream of RowSet based on formatting
85    /// parameters.
86    ///
87    /// This is essentially `StreamExt::map(self, move |res| res.map(|chunk| to_pg_rows(chunk,
88    /// format)))` but we need a nameable type as part of [`super::PgResponseStream`], but we cannot
89    /// name the type of a closure.
90    pub struct DataChunkToRowSetAdapter<VS>
91    where
92        VS: Stream<Item = Result<DataChunk, BoxedError>>,
93    {
94        #[pin]
95        chunk_stream: VS,
96        column_types: Vec<DataType>,
97        pub formats: Vec<Format>,
98        session_data: StaticSessionData,
99    }
100}
101
102// Static session data frozen at the time of the creation of the stream
103pub struct StaticSessionData {
104    pub timezone: String,
105}
106
107impl<VS> DataChunkToRowSetAdapter<VS>
108where
109    VS: Stream<Item = Result<DataChunk, BoxedError>>,
110{
111    pub fn new(
112        chunk_stream: VS,
113        column_types: Vec<DataType>,
114        formats: Vec<Format>,
115        session: Arc<SessionImpl>,
116    ) -> Self {
117        let session_data = StaticSessionData {
118            timezone: session.config().timezone(),
119        };
120        Self {
121            chunk_stream,
122            column_types,
123            formats,
124            session_data,
125        }
126    }
127}
128
129impl<VS> Stream for DataChunkToRowSetAdapter<VS>
130where
131    VS: Stream<Item = Result<DataChunk, BoxedError>>,
132{
133    type Item = RowSetResult;
134
135    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
136        let mut this = self.project();
137        match this.chunk_stream.as_mut().poll_next(cx) {
138            Poll::Pending => Poll::Pending,
139            Poll::Ready(chunk) => match chunk {
140                Some(chunk_result) => match chunk_result {
141                    Ok(chunk) => Poll::Ready(Some(
142                        to_pg_rows(this.column_types, chunk, this.formats, this.session_data)
143                            .map_err(|err| err.into()),
144                    )),
145                    Err(err) => Poll::Ready(Some(Err(err))),
146                },
147                None => Poll::Ready(None),
148            },
149        }
150    }
151}
152
153/// Format scalars according to postgres convention.
154pub fn pg_value_format(
155    data_type: &DataType,
156    d: ScalarRefImpl<'_>,
157    format: Format,
158    session_data: &StaticSessionData,
159) -> RwResult<Bytes> {
160    // format == false means TEXT format
161    // format == true means BINARY format
162    match format {
163        Format::Text => {
164            if *data_type == DataType::Timestamptz {
165                Ok(timestamptz_to_string_with_session_data(d, session_data))
166            } else {
167                Ok(d.text_format(data_type).into())
168            }
169        }
170        Format::Binary => Ok(d
171            .binary_format(data_type)
172            .context("failed to format binary value")?),
173    }
174}
175
176fn timestamptz_to_string_with_session_data(
177    d: ScalarRefImpl<'_>,
178    session_data: &StaticSessionData,
179) -> Bytes {
180    let tz = d.into_timestamptz();
181    let time_zone = Timestamptz::lookup_time_zone(&session_data.timezone).unwrap();
182    let instant_local = tz.to_datetime_in_zone(time_zone);
183    let mut result_string = BytesMut::new();
184    write_date_time_tz(instant_local, &mut result_string).unwrap();
185    result_string.into()
186}
187
188fn to_pg_rows(
189    column_types: &[DataType],
190    chunk: DataChunk,
191    formats: &[Format],
192    session_data: &StaticSessionData,
193) -> RwResult<Vec<Row>> {
194    assert_eq!(chunk.dimension(), column_types.len());
195    if cfg!(debug_assertions) {
196        let chunk_data_types = chunk.data_types();
197        for (ty1, ty2) in chunk_data_types.iter().zip_eq_fast(column_types) {
198            debug_assert!(
199                ty1.equals_datatype(ty2),
200                "chunk_data_types: {chunk_data_types:?}, column_types: {column_types:?}"
201            )
202        }
203    }
204
205    chunk
206        .rows()
207        .map(|r| {
208            let format_iter = FormatIterator::new(formats, chunk.dimension())
209                .map_err(ErrorCode::InternalError)?;
210            let row = r
211                .iter()
212                .zip_eq_fast(column_types)
213                .zip_eq_fast(format_iter)
214                .map(|((data, t), format)| match data {
215                    Some(data) => Some(pg_value_format(t, data, format, session_data)).transpose(),
216                    None => Ok(None),
217                })
218                .try_collect()?;
219            Ok(Row::new(row))
220        })
221        .try_collect()
222}
223
224/// Convert from [`Field`] to [`PgFieldDescriptor`].
225pub fn to_pg_field(f: &Field) -> PgFieldDescriptor {
226    PgFieldDescriptor::new(
227        f.name.clone(),
228        f.data_type().to_oid(),
229        f.data_type().type_len(),
230    )
231}
232
233#[easy_ext::ext(SourceSchemaCompatExt)]
234impl CompatibleFormatEncode {
235    /// Convert `self` to [`FormatEncodeOptions`] and warn the user if the syntax is deprecated.
236    pub fn into_v2_with_warning(self) -> FormatEncodeOptions {
237        match self {
238            CompatibleFormatEncode::V2(inner) => inner,
239        }
240    }
241}
242
243pub fn gen_query_from_table_name(from_name: ObjectName) -> Query {
244    let table_factor = TableFactor::Table {
245        name: from_name,
246        alias: None,
247        as_of: None,
248    };
249    let from = vec![TableWithJoins {
250        relation: table_factor,
251        joins: vec![],
252    }];
253    let select = Select {
254        from,
255        projection: vec![SelectItem::Wildcard(None)],
256        ..Default::default()
257    };
258    let body = SetExpr::Select(Box::new(select));
259    Query {
260        with: None,
261        body,
262        order_by: vec![],
263        limit: None,
264        offset: None,
265        fetch: None,
266    }
267}
268
269pub fn convert_unix_millis_to_logstore_u64(unix_millis: u64) -> u64 {
270    Epoch::from_unix_millis(unix_millis).0
271}
272
273pub fn convert_logstore_u64_to_unix_millis(logstore_u64: u64) -> u64 {
274    Epoch::from(logstore_u64).as_unix_millis()
275}
276
277pub fn convert_interval_to_u64_seconds(interval: &String) -> RwResult<u64> {
278    let seconds = (Interval::from_str(interval)
279        .map_err(|err| {
280            ErrorCode::InternalError(format!(
281                "Convert interval to u64 error, please check format, error: {:?}",
282                err.to_report_string()
283            ))
284        })?
285        .epoch_in_micros()
286        / 1000000) as u64;
287    Ok(seconds)
288}
289
290pub fn ensure_connection_type_allowed(
291    connection_type: PbConnectionType,
292    allowed_types: &HashSet<PbConnectionType>,
293) -> RwResult<()> {
294    if !allowed_types.contains(&connection_type) {
295        return Err(RwError::from(ProtocolError(format!(
296            "connection type {:?} is not allowed, allowed types: {:?}",
297            connection_type, allowed_types
298        ))));
299    }
300    Ok(())
301}
302
303fn connection_type_to_connector(connection_type: &PbConnectionType) -> &str {
304    match connection_type {
305        PbConnectionType::Kafka => KAFKA_CONNECTOR,
306        PbConnectionType::Iceberg => ICEBERG_CONNECTOR,
307        PbConnectionType::Elasticsearch => ES_SINK,
308        _ => unreachable!(),
309    }
310}
311
312pub fn check_connector_match_connection_type(
313    connector: &str,
314    connection_type: &PbConnectionType,
315) -> RwResult<()> {
316    if !connector.eq(connection_type_to_connector(connection_type)) {
317        return Err(RwError::from(ProtocolError(format!(
318            "connector {} and connection type {:?} are not compatible",
319            connector, connection_type
320        ))));
321    }
322    Ok(())
323}
324
325pub fn get_table_catalog_by_table_name(
326    session: &SessionImpl,
327    table_name: &ObjectName,
328) -> RwResult<(Arc<TableCatalog>, String)> {
329    let db_name = &session.database();
330    let (schema_name, real_table_name) =
331        Binder::resolve_schema_qualified_name(db_name, table_name)?;
332    let search_path = session.config().search_path();
333    let user_name = &session.user_name();
334
335    let schema_path = SchemaPath::new(schema_name.as_deref(), &search_path, user_name);
336    let reader = session.env().catalog_reader().read_guard();
337    match reader.get_created_table_by_name(db_name, schema_path, &real_table_name) {
338        Ok((table, schema_name)) => Ok((table.clone(), schema_name.to_owned())),
339        Err(err) => {
340            if let Some(table) = session
341                .staging_catalog_manager()
342                .get_table(&real_table_name)
343            {
344                // During CREATE TABLE (iceberg engine), the table is only in staging, but we
345                // still need a stable schema name for internal sink planning.
346                let schema_name = reader
347                    .get_schema_by_id(table.database_id, table.schema_id)
348                    .map(|schema| schema.name.clone())?;
349                Ok((Arc::new(table.clone()), schema_name))
350            } else {
351                Err(err.into())
352            }
353        }
354    }
355}
356
357pub fn reject_internal_table_dependency(
358    table: &TableCatalog,
359    statement_name: &str,
360) -> RwResult<()> {
361    if table.is_internal_table() {
362        return Err(RwError::from(ErrorCode::InvalidInputSyntax(format!(
363            "{statement_name} does not support internal table \"{}\"",
364            table.name()
365        ))));
366    }
367    Ok(())
368}
369
370pub fn reject_internal_table_dependencies<'a, I>(
371    session: &SessionImpl,
372    dependent_relations: I,
373    statement_name: &str,
374) -> RwResult<()>
375where
376    I: IntoIterator<Item = &'a ObjectId>,
377{
378    let catalog_reader = session.env().catalog_reader().read_guard();
379    for object_id in dependent_relations {
380        if let Ok(table) = catalog_reader.get_any_table_by_id(object_id.as_table_id()) {
381            reject_internal_table_dependency(table.as_ref(), statement_name)?;
382        }
383    }
384    Ok(())
385}
386
387/// Execute an async operation with a notification if it takes too long.
388/// This is useful for operations that might be delayed due to high barrier latency.
389///
390/// The notification timeout duration is controlled by the `slow_ddl_notification_secs` session variable.
391///
392/// # Arguments
393/// * `operation_fut` - The async operation to execute
394/// * `session` - The session to send notifications to
395/// * `operation_name` - The name of the operation for the notification message (e.g., "DROP TABLE")
396///
397/// # Example
398/// ```ignore
399/// execute_with_long_running_notification(
400///     catalog_writer.drop_table(source_id, table_id, cascade),
401///     &session,
402///     "DROP TABLE",
403/// ).await?;
404/// ```
405#[derive(Clone, Copy)]
406pub enum LongRunningNotificationAction {
407    SuggestRecover,
408    DiagnoseBarrierLatency,
409    MonitorBackfillJob,
410}
411
412impl LongRunningNotificationAction {
413    fn build_message(self, operation_name: &str, notify_timeout_secs: u32) -> String {
414        match self {
415            LongRunningNotificationAction::SuggestRecover => format!(
416                "{} has taken more than {} secs, likely due to high barrier latency.\n\
417                You may trigger cluster recovery to let {} take effect immediately.\n\
418                Run RECOVER in a separate session to trigger recovery.\n\
419                See: https://docs.risingwave.com/sql/commands/sql-recover#recover",
420                operation_name, notify_timeout_secs, operation_name
421            ),
422            LongRunningNotificationAction::DiagnoseBarrierLatency => format!(
423                "{} has taken more than {} secs, likely due to high barrier latency.\n\
424                See: https://docs.risingwave.com/performance/metrics#barrier-monitoring for steps to diagnose high barrier latency.",
425                operation_name, notify_timeout_secs
426            ),
427            LongRunningNotificationAction::MonitorBackfillJob => format!(
428                "{} has taken more than {} secs, barrier latency might be high. Please check barrier latency metrics to confirm.\n\
429                You can also run SHOW JOBS to track the progress of the job.\n\
430                See: https://docs.risingwave.com/performance/metrics#barrier-monitoring and https://docs.risingwave.com/sql/commands/sql-show-jobs",
431                operation_name, notify_timeout_secs
432            ),
433        }
434    }
435}
436
437pub async fn execute_with_long_running_notification<F, T>(
438    operation_fut: F,
439    session: &SessionImpl,
440    operation_name: &str,
441    action: LongRunningNotificationAction,
442) -> RwResult<T>
443where
444    F: std::future::Future<Output = RwResult<T>>,
445{
446    let notify_timeout_secs = session.config().slow_ddl_notification_secs();
447
448    // If timeout is 0, disable notifications and just execute the operation
449    if notify_timeout_secs == 0 {
450        return operation_fut.await;
451    }
452
453    let notify_fut = sleep(Duration::from_secs(notify_timeout_secs as u64));
454    tokio::pin!(operation_fut);
455
456    select! {
457        _ = notify_fut => {
458            session.notice_to_user(action.build_message(operation_name, notify_timeout_secs));
459            operation_fut.await
460        }
461        result = &mut operation_fut => {
462            result
463        }
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use postgres_types::{ToSql, Type};
470    use risingwave_common::array::*;
471
472    use super::*;
473
474    #[test]
475    fn test_to_pg_field() {
476        let field = Field::with_name(DataType::Int32, "v1");
477        let pg_field = to_pg_field(&field);
478        assert_eq!(pg_field.get_name(), "v1");
479        assert_eq!(pg_field.get_type_oid(), DataType::Int32.to_oid());
480    }
481
482    #[test]
483    fn test_to_pg_rows() {
484        let chunk = DataChunk::from_pretty(
485            "i I f    T
486             1 6 6.01 aaa
487             2 . .    .
488             3 7 7.01 vvv
489             4 . .    .  ",
490        );
491        let static_session = StaticSessionData {
492            timezone: "UTC".into(),
493        };
494        let rows = to_pg_rows(
495            &[
496                DataType::Int32,
497                DataType::Int64,
498                DataType::Float32,
499                DataType::Varchar,
500            ],
501            chunk,
502            &[],
503            &static_session,
504        );
505        let expected: Vec<Vec<Option<Bytes>>> = vec![
506            vec![
507                Some("1".into()),
508                Some("6".into()),
509                Some("6.01".into()),
510                Some("aaa".into()),
511            ],
512            vec![Some("2".into()), None, None, None],
513            vec![
514                Some("3".into()),
515                Some("7".into()),
516                Some("7.01".into()),
517                Some("vvv".into()),
518            ],
519            vec![Some("4".into()), None, None, None],
520        ];
521        let vec = rows
522            .unwrap()
523            .into_iter()
524            .map(|r| r.values().iter().cloned().collect_vec())
525            .collect_vec();
526
527        assert_eq!(vec, expected);
528    }
529
530    #[test]
531    fn test_to_pg_rows_mix_format() {
532        let chunk = DataChunk::from_pretty(
533            "i I f    T
534             1 6 6.01 aaa
535            ",
536        );
537        let static_session = StaticSessionData {
538            timezone: "UTC".into(),
539        };
540        let rows = to_pg_rows(
541            &[
542                DataType::Int32,
543                DataType::Int64,
544                DataType::Float32,
545                DataType::Varchar,
546            ],
547            chunk,
548            &[Format::Binary, Format::Binary, Format::Binary, Format::Text],
549            &static_session,
550        );
551        let mut raw_params = vec![BytesMut::new(); 3];
552        1_i32.to_sql(&Type::ANY, &mut raw_params[0]).unwrap();
553        6_i64.to_sql(&Type::ANY, &mut raw_params[1]).unwrap();
554        6.01_f32.to_sql(&Type::ANY, &mut raw_params[2]).unwrap();
555        let raw_params = raw_params
556            .into_iter()
557            .map(|b| b.freeze())
558            .collect::<Vec<_>>();
559        let expected: Vec<Vec<Option<Bytes>>> = vec![vec![
560            Some(raw_params[0].clone()),
561            Some(raw_params[1].clone()),
562            Some(raw_params[2].clone()),
563            Some("aaa".into()),
564        ]];
565        let vec = rows
566            .unwrap()
567            .into_iter()
568            .map(|r| r.values().iter().cloned().collect_vec())
569            .collect_vec();
570
571        assert_eq!(vec, expected);
572    }
573
574    #[test]
575    fn test_value_format() {
576        use DataType as T;
577        use ScalarRefImpl as S;
578        let static_session = StaticSessionData {
579            timezone: "UTC".into(),
580        };
581
582        let f = |t, d, f| pg_value_format(t, d, f, &static_session).unwrap();
583        assert_eq!(&f(&T::Float32, S::Float32(1_f32.into()), Format::Text), "1");
584        assert_eq!(
585            &f(&T::Float32, S::Float32(f32::NAN.into()), Format::Text),
586            "NaN"
587        );
588        assert_eq!(
589            &f(&T::Float64, S::Float64(f64::NAN.into()), Format::Text),
590            "NaN"
591        );
592        assert_eq!(
593            &f(&T::Float32, S::Float32(f32::INFINITY.into()), Format::Text),
594            "Infinity"
595        );
596        assert_eq!(
597            &f(
598                &T::Float32,
599                S::Float32(f32::NEG_INFINITY.into()),
600                Format::Text
601            ),
602            "-Infinity"
603        );
604        assert_eq!(
605            &f(&T::Float64, S::Float64(f64::INFINITY.into()), Format::Text),
606            "Infinity"
607        );
608        assert_eq!(
609            &f(
610                &T::Float64,
611                S::Float64(f64::NEG_INFINITY.into()),
612                Format::Text
613            ),
614            "-Infinity"
615        );
616        assert_eq!(&f(&T::Boolean, S::Bool(true), Format::Text), "t");
617        assert_eq!(&f(&T::Boolean, S::Bool(false), Format::Text), "f");
618        assert_eq!(
619            &f(
620                &T::Timestamptz,
621                S::Timestamptz(Timestamptz::from_micros(-1).unwrap()),
622                Format::Text
623            ),
624            "1969-12-31 23:59:59.999999+00:00"
625        );
626    }
627
628    #[test]
629    fn test_local_fs_connector_config_gate() {
630        let frontend_config = FrontendConfig::default();
631        let default_enabled = cfg!(debug_assertions);
632        assert_eq!(
633            is_local_fs_connector_enabled(&frontend_config, POSIX_FS_CONNECTOR),
634            default_enabled
635        );
636        assert_eq!(
637            is_local_fs_connector_enabled(&frontend_config, BATCH_POSIX_FS_CONNECTOR),
638            default_enabled
639        );
640        assert_eq!(
641            is_local_fs_connector_enabled(&frontend_config, FS_SINK),
642            default_enabled
643        );
644        assert!(is_local_fs_connector_enabled(
645            &frontend_config,
646            KAFKA_CONNECTOR
647        ));
648
649        let disabled_config = FrontendConfig {
650            unsafe_enable_local_fs_connector: false,
651            ..Default::default()
652        };
653        assert!(!is_local_fs_connector_enabled(
654            &disabled_config,
655            POSIX_FS_CONNECTOR
656        ));
657        assert!(!is_local_fs_connector_enabled(
658            &disabled_config,
659            BATCH_POSIX_FS_CONNECTOR
660        ));
661        assert!(!is_local_fs_connector_enabled(&disabled_config, FS_SINK));
662
663        let enabled_config = FrontendConfig {
664            unsafe_enable_local_fs_connector: true,
665            ..Default::default()
666        };
667        assert!(is_local_fs_connector_enabled(
668            &enabled_config,
669            POSIX_FS_CONNECTOR
670        ));
671        assert!(is_local_fs_connector_enabled(
672            &enabled_config,
673            BATCH_POSIX_FS_CONNECTOR
674        ));
675        assert!(is_local_fs_connector_enabled(&enabled_config, FS_SINK));
676    }
677}
risingwave_frontend/handler/util.rs

risingwave_frontend/handler/
util.rs
