risingwave_frontend/handler/
alter_streaming_rate_limit.rs

1// Copyright 2024 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 pgwire::pg_response::{PgResponse, StatementType};
16use risingwave_common::bail;
17use risingwave_pb::common::ThrottleType as PbThrottleType;
18use risingwave_pb::meta::ThrottleTarget as PbThrottleTarget;
19use risingwave_sqlparser::ast::ObjectName;
20
21use super::{HandlerArgs, RwPgResponse};
22use crate::catalog::root_catalog::{Catalog, SchemaPath};
23use crate::catalog::table_catalog::TableType;
24use crate::error::{ErrorCode, Result};
25use crate::handler::util::{LongRunningNotificationAction, execute_with_long_running_notification};
26use crate::session::SessionImpl;
27use crate::{Binder, TableCatalog};
28
29fn check_table_mismatch<'a>(
30    reader: &'a Catalog,
31    session: &SessionImpl,
32    db_name: &str,
33    schema_path: SchemaPath<'_>,
34    table_name: &str,
35) -> Result<&'a TableCatalog> {
36    let (table, schema_name) = reader.get_table_by_name(db_name, schema_path, table_name, true)?;
37    if table.table_type != TableType::Table {
38        return Err(
39            ErrorCode::InvalidInputSyntax(format!("\"{table_name}\" is not a TABLE",)).into(),
40        );
41    }
42    session.check_privilege_for_drop_alter(schema_name, &**table)?;
43    Ok(table)
44}
45
46pub async fn handle_alter_streaming_rate_limit(
47    handler_args: HandlerArgs,
48    throttle_target: PbThrottleTarget,
49    throttle_type: PbThrottleType,
50    table_name: ObjectName,
51    rate_limit: i32,
52) -> Result<RwPgResponse> {
53    let session = handler_args.clone().session;
54    let db_name = &session.database();
55    let (schema_name, real_table_name) =
56        Binder::resolve_schema_qualified_name(db_name, &table_name)?;
57    let search_path = session.config().search_path();
58    let user_name = &session.user_name();
59
60    let schema_path = SchemaPath::new(schema_name.as_deref(), &search_path, user_name);
61
62    let (stmt_type, id) = match (throttle_target, throttle_type) {
63        (PbThrottleTarget::Mv, PbThrottleType::Backfill) => {
64            let reader = session.env().catalog_reader().read_guard();
65            let (table, schema_name) =
66                reader.get_any_table_by_name(db_name, schema_path, &real_table_name)?;
67            if table.table_type != TableType::MaterializedView {
68                return Err(ErrorCode::InvalidInputSyntax(format!(
69                    "\"{table_name}\" is not a materialized view",
70                ))
71                .into());
72            }
73            session.check_privilege_for_drop_alter(schema_name, &**table)?;
74            (StatementType::ALTER_MATERIALIZED_VIEW, table.id.as_raw_id())
75        }
76        (PbThrottleTarget::Source, PbThrottleType::Source) => {
77            let reader = session.env().catalog_reader().read_guard();
78            let (source, schema_name) =
79                reader.get_source_by_name(db_name, schema_path, &real_table_name)?;
80            session.check_privilege_for_drop_alter(schema_name, &**source)?;
81            if source.is_cdc_table_source() {
82                return Err(ErrorCode::NotSupported(
83                    "SOURCE_RATE_LIMIT is not supported for CDC table sources".to_owned(),
84                    "Alter the upstream shared CDC source instead".to_owned(),
85                )
86                .into());
87            }
88            (StatementType::ALTER_SOURCE, source.id.as_raw_id())
89        }
90        (PbThrottleTarget::Table, PbThrottleType::Dml) => {
91            let reader = session.env().catalog_reader().read_guard();
92            let table =
93                check_table_mismatch(&reader, &session, db_name, schema_path, &real_table_name)?;
94            (StatementType::ALTER_TABLE, table.id.as_raw_id())
95        }
96        (PbThrottleTarget::Table, PbThrottleType::Source) => {
97            let reader = session.env().catalog_reader().read_guard();
98            let table =
99                check_table_mismatch(&reader, &session, db_name, schema_path, &real_table_name)?;
100            let source_id = if let Some(id) = table.associated_source_id {
101                id.as_raw_id()
102            } else {
103                bail!("ALTER SOURCE_RATE_LIMIT is not for table without source")
104            };
105            (StatementType::ALTER_TABLE, source_id)
106        }
107        (PbThrottleTarget::Table, PbThrottleType::Backfill) => {
108            let reader = session.env().catalog_reader().read_guard();
109            let table =
110                check_table_mismatch(&reader, &session, db_name, schema_path, &real_table_name)?;
111            if table.cdc_table_type.is_none() {
112                return Err(ErrorCode::InvalidInputSyntax(format!(
113                    "\"{table_name}\" is not a CDC table",
114                ))
115                .into());
116            }
117            (StatementType::ALTER_TABLE, table.id.as_raw_id())
118        }
119        (PbThrottleTarget::Sink, PbThrottleType::Sink) => {
120            let reader = session.env().catalog_reader().read_guard();
121            let (sink, schema_name) =
122                reader.get_any_sink_by_name(db_name, schema_path, &real_table_name)?;
123            if sink.target_table.is_some() {
124                bail!("ALTER SINK_RATE_LIMIT is not for sink into table")
125            }
126            session.check_privilege_for_drop_alter(schema_name, &**sink)?;
127            (StatementType::ALTER_SINK, sink.id.as_raw_id())
128        }
129        (PbThrottleTarget::Sink, PbThrottleType::Backfill) => {
130            let reader = session.env().catalog_reader().read_guard();
131            let (sink, schema_name) =
132                reader.get_any_sink_by_name(db_name, schema_path, &real_table_name)?;
133            session.check_privilege_for_drop_alter(schema_name, &**sink)?;
134            (StatementType::ALTER_SINK, sink.id.as_raw_id())
135        }
136        _ => bail!(
137            "Unsupported throttle target: {:?} and throttle type: {:?}",
138            throttle_target,
139            throttle_type
140        ),
141    };
142    execute_with_long_running_notification(
143        handle_alter_streaming_rate_limit_by_id(
144            &session,
145            throttle_target,
146            throttle_type,
147            id,
148            rate_limit,
149            stmt_type,
150        ),
151        &session,
152        "ALTER STREAMING RATE LIMIT",
153        LongRunningNotificationAction::SuggestRecover,
154    )
155    .await
156}
157
158pub async fn handle_alter_streaming_rate_limit_by_id(
159    session: &SessionImpl,
160    throttle_target: PbThrottleTarget,
161    throttle_type: PbThrottleType,
162    id: u32,
163    rate_limit: i32,
164    stmt_type: StatementType,
165) -> Result<RwPgResponse> {
166    let meta_client = session.env().meta_client();
167
168    let rate_limit = if rate_limit < 0 {
169        None
170    } else {
171        Some(rate_limit as u32)
172    };
173
174    meta_client
175        .apply_throttle(throttle_target, throttle_type, id, rate_limit)
176        .await?;
177
178    Ok(PgResponse::empty_result(stmt_type))
179}
risingwave_frontend/handler/alter_streaming_rate_limit.rs

risingwave_frontend/handler/
alter_streaming_rate_limit.rs
