Skip to main content

risingwave_meta/controller/
utils.rs

1// Copyright 2023 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 std::collections::{BTreeSet, HashMap, HashSet};
16use std::sync::Arc;
17
18use anyhow::{Context, anyhow};
19use itertools::Itertools;
20use risingwave_common::bitmap::Bitmap;
21use risingwave_common::catalog::{
22    FragmentTypeFlag, FragmentTypeMask, ICEBERG_SINK_PREFIX, ICEBERG_SOURCE_PREFIX,
23};
24use risingwave_common::hash::{ActorMapping, VnodeBitmapExt, WorkerSlotId, WorkerSlotMapping};
25use risingwave_common::id::{JobId, SubscriptionId};
26use risingwave_common::types::{DataType, Datum};
27use risingwave_common::util::value_encoding::DatumToProtoExt;
28use risingwave_common::util::worker_util::DEFAULT_RESOURCE_GROUP;
29use risingwave_common::{bail, hash};
30use risingwave_connector::WithOptionsSecResolved;
31use risingwave_meta_model::fragment::DistributionType;
32use risingwave_meta_model::object::ObjectType;
33use risingwave_meta_model::prelude::*;
34use risingwave_meta_model::streaming_job::BackfillOrders;
35use risingwave_meta_model::table::TableType;
36use risingwave_meta_model::user_privilege::Action;
37use risingwave_meta_model::{
38    ActorId, ColumnCatalogArray, CreateType, DataTypeArray, DatabaseId, DispatcherType, FragmentId,
39    JobStatus, ObjectId, PrivilegeId, SchemaId, SinkId, SourceId, StreamNode, StreamSourceInfo,
40    TableId, TableIdArray, UserId, WorkerId, connection, database, fragment, fragment_relation,
41    function, index, object, object_dependency, schema, secret, sink, source, streaming_job,
42    subscription, table, user, user_default_privilege, user_privilege, view,
43};
44use risingwave_meta_model_migration::WithQuery;
45use risingwave_pb::catalog::{
46    PbConnection, PbDatabase, PbFunction, PbIndex, PbSchema, PbSecret, PbSink, PbSource,
47    PbSubscription, PbTable, PbView,
48};
49use risingwave_pb::common::{PbObjectType, WorkerNode};
50use risingwave_pb::expr::{PbExprNode, expr_node};
51use risingwave_pb::meta::object::PbObjectInfo;
52use risingwave_pb::meta::subscribe_response::Info as NotificationInfo;
53use risingwave_pb::meta::{
54    ObjectDependency as PbObjectDependency, PbFragmentWorkerSlotMapping, PbObject, PbObjectGroup,
55};
56use risingwave_pb::plan_common::column_desc::GeneratedOrDefaultColumn;
57use risingwave_pb::plan_common::{ColumnCatalog, DefaultColumnDesc};
58use risingwave_pb::secret::PbSecretRef;
59use risingwave_pb::secret::secret_ref::PbRefAsType;
60use risingwave_pb::stream_plan::{PbDispatchOutputMapping, PbDispatcher, PbDispatcherType};
61use risingwave_pb::user::grant_privilege::{PbActionWithGrantOption, PbObject as PbGrantObject};
62use risingwave_pb::user::{PbAction, PbGrantPrivilege, PbUserInfo};
63use risingwave_sqlparser::ast::{
64    Ident, ObjectName, SecretRefAsType, SecretRefValue, SqlOption, Statement as SqlStatement,
65};
66use risingwave_sqlparser::parser::Parser;
67use sea_orm::sea_query::{
68    Alias, CommonTableExpression, Expr, OnConflict, Query, QueryStatementBuilder, SelectStatement,
69    UnionType, WithClause,
70};
71use sea_orm::{
72    ColumnTrait, ConnectionTrait, DatabaseTransaction, DerivePartialModel, EntityTrait,
73    FromQueryResult, IntoActiveModel, JoinType, Order, PaginatorTrait, QueryFilter, QuerySelect,
74    RelationTrait, Set, Statement,
75};
76use thiserror_ext::AsReport;
77use tracing::warn;
78
79use crate::barrier::SharedFragmentInfo;
80use crate::controller::ObjectModel;
81use crate::controller::fragment::FragmentTypeMaskExt;
82use crate::controller::scale::resolve_streaming_job_definition;
83use crate::model::{FragmentDownstreamRelation, StreamContext};
84use crate::{MetaError, MetaResult};
85
86/// This function will construct a query using recursive cte to find all objects[(id, `obj_type`)] that are used by the given object.
87///
88/// # Examples
89///
90/// ```
91/// use risingwave_meta::controller::utils::construct_obj_dependency_query;
92/// use sea_orm::sea_query::*;
93/// use sea_orm::*;
94///
95/// let query = construct_obj_dependency_query(1.into());
96///
97/// assert_eq!(
98///     query.to_string(MysqlQueryBuilder),
99///     r#"WITH RECURSIVE `used_by_object_ids` (`used_by`) AS (SELECT `used_by` FROM `object_dependency` WHERE `object_dependency`.`oid` = 1 UNION ALL (SELECT `oid` FROM `object` WHERE `object`.`database_id` = 1 OR `object`.`schema_id` = 1) UNION ALL (SELECT `object_dependency`.`used_by` FROM `object_dependency` INNER JOIN `used_by_object_ids` ON `used_by_object_ids`.`used_by` = `oid`)) SELECT DISTINCT `oid`, `obj_type`, `schema_id`, `database_id` FROM `used_by_object_ids` INNER JOIN `object` ON `used_by_object_ids`.`used_by` = `oid` ORDER BY `oid` DESC"#
100/// );
101/// assert_eq!(
102///     query.to_string(PostgresQueryBuilder),
103///     r#"WITH RECURSIVE "used_by_object_ids" ("used_by") AS (SELECT "used_by" FROM "object_dependency" WHERE "object_dependency"."oid" = 1 UNION ALL (SELECT "oid" FROM "object" WHERE "object"."database_id" = 1 OR "object"."schema_id" = 1) UNION ALL (SELECT "object_dependency"."used_by" FROM "object_dependency" INNER JOIN "used_by_object_ids" ON "used_by_object_ids"."used_by" = "oid")) SELECT DISTINCT "oid", "obj_type", "schema_id", "database_id" FROM "used_by_object_ids" INNER JOIN "object" ON "used_by_object_ids"."used_by" = "oid" ORDER BY "oid" DESC"#
104/// );
105/// assert_eq!(
106///     query.to_string(SqliteQueryBuilder),
107///     r#"WITH RECURSIVE "used_by_object_ids" ("used_by") AS (SELECT "used_by" FROM "object_dependency" WHERE "object_dependency"."oid" = 1 UNION ALL SELECT "oid" FROM "object" WHERE "object"."database_id" = 1 OR "object"."schema_id" = 1 UNION ALL SELECT "object_dependency"."used_by" FROM "object_dependency" INNER JOIN "used_by_object_ids" ON "used_by_object_ids"."used_by" = "oid") SELECT DISTINCT "oid", "obj_type", "schema_id", "database_id" FROM "used_by_object_ids" INNER JOIN "object" ON "used_by_object_ids"."used_by" = "oid" ORDER BY "oid" DESC"#
108/// );
109/// ```
110pub fn construct_obj_dependency_query(obj_id: ObjectId) -> WithQuery {
111    let cte_alias = Alias::new("used_by_object_ids");
112    let cte_return_alias = Alias::new("used_by");
113
114    let mut base_query = SelectStatement::new()
115        .column(object_dependency::Column::UsedBy)
116        .from(ObjectDependency)
117        .and_where(object_dependency::Column::Oid.eq(obj_id))
118        .to_owned();
119
120    let belonged_obj_query = SelectStatement::new()
121        .column(object::Column::Oid)
122        .from(Object)
123        .and_where(
124            object::Column::DatabaseId
125                .eq(obj_id)
126                .or(object::Column::SchemaId.eq(obj_id)),
127        )
128        .to_owned();
129
130    let cte_referencing = Query::select()
131        .column((ObjectDependency, object_dependency::Column::UsedBy))
132        .from(ObjectDependency)
133        .inner_join(
134            cte_alias.clone(),
135            Expr::col((cte_alias.clone(), cte_return_alias.clone()))
136                .equals(object_dependency::Column::Oid),
137        )
138        .to_owned();
139
140    let mut common_table_expr = CommonTableExpression::new();
141    common_table_expr
142        .query(
143            base_query
144                .union(UnionType::All, belonged_obj_query)
145                .union(UnionType::All, cte_referencing)
146                .to_owned(),
147        )
148        .column(cte_return_alias.clone())
149        .table_name(cte_alias.clone());
150
151    SelectStatement::new()
152        .distinct()
153        .columns([
154            object::Column::Oid,
155            object::Column::ObjType,
156            object::Column::SchemaId,
157            object::Column::DatabaseId,
158        ])
159        .from(cte_alias.clone())
160        .inner_join(
161            Object,
162            Expr::col((cte_alias, cte_return_alias)).equals(object::Column::Oid),
163        )
164        .order_by(object::Column::Oid, Order::Desc)
165        .to_owned()
166        .with(
167            WithClause::new()
168                .recursive(true)
169                .cte(common_table_expr)
170                .to_owned(),
171        )
172}
173
174fn to_pb_object_type(obj_type: ObjectType) -> PbObjectType {
175    match obj_type {
176        ObjectType::Database => PbObjectType::Database,
177        ObjectType::Schema => PbObjectType::Schema,
178        ObjectType::Table => PbObjectType::Table,
179        ObjectType::Source => PbObjectType::Source,
180        ObjectType::Sink => PbObjectType::Sink,
181        ObjectType::View => PbObjectType::View,
182        ObjectType::Index => PbObjectType::Index,
183        ObjectType::Function => PbObjectType::Function,
184        ObjectType::Connection => PbObjectType::Connection,
185        ObjectType::Subscription => PbObjectType::Subscription,
186        ObjectType::Secret => PbObjectType::Secret,
187    }
188}
189
190/// Collect object dependencies with optional object id filtering, and optionally exclude non-created streaming jobs.
191async fn list_object_dependencies_impl(
192    txn: &DatabaseTransaction,
193    object_id: Option<ObjectId>,
194    include_creating: bool,
195) -> MetaResult<Vec<PbObjectDependency>> {
196    let referenced_alias = Alias::new("referenced_obj");
197    let mut query = ObjectDependency::find()
198        .select_only()
199        .columns([
200            object_dependency::Column::Oid,
201            object_dependency::Column::UsedBy,
202        ])
203        .column_as(
204            Expr::col((referenced_alias.clone(), object::Column::ObjType)),
205            "referenced_obj_type",
206        )
207        .join(
208            JoinType::InnerJoin,
209            object_dependency::Relation::Object1.def(),
210        )
211        .join_as(
212            JoinType::InnerJoin,
213            object_dependency::Relation::Object2.def(),
214            referenced_alias.clone(),
215        );
216    if let Some(object_id) = object_id {
217        query = query.filter(object_dependency::Column::UsedBy.eq(object_id));
218    }
219    let mut obj_dependencies: Vec<PbObjectDependency> = query
220        .into_tuple()
221        .all(txn)
222        .await?
223        .into_iter()
224        .map(|(oid, used_by, referenced_type)| PbObjectDependency {
225            object_id: used_by,
226            referenced_object_id: oid,
227            referenced_object_type: to_pb_object_type(referenced_type) as i32,
228        })
229        .collect();
230
231    let mut sink_query = Sink::find()
232        .select_only()
233        .columns([sink::Column::SinkId, sink::Column::TargetTable])
234        .filter(sink::Column::TargetTable.is_not_null());
235    if let Some(object_id) = object_id {
236        sink_query = sink_query.filter(
237            sink::Column::SinkId
238                .eq(object_id)
239                .or(sink::Column::TargetTable.eq(object_id)),
240        );
241    }
242    let sink_dependencies: Vec<(SinkId, TableId)> = sink_query.into_tuple().all(txn).await?;
243    obj_dependencies.extend(sink_dependencies.into_iter().map(|(sink_id, table_id)| {
244        PbObjectDependency {
245            object_id: table_id.into(),
246            referenced_object_id: sink_id.into(),
247            referenced_object_type: PbObjectType::Sink as i32,
248        }
249    }));
250
251    if !include_creating {
252        let mut streaming_job_ids = obj_dependencies
253            .iter()
254            .map(|dependency| dependency.object_id)
255            .collect_vec();
256        streaming_job_ids.sort_unstable();
257        streaming_job_ids.dedup();
258
259        if !streaming_job_ids.is_empty() {
260            let non_created_jobs: HashSet<JobId> = StreamingJob::find()
261                .select_only()
262                .columns([streaming_job::Column::JobId])
263                .filter(
264                    streaming_job::Column::JobId
265                        .is_in(streaming_job_ids)
266                        .and(streaming_job::Column::JobStatus.ne(JobStatus::Created)),
267                )
268                .into_tuple()
269                .all(txn)
270                .await?
271                .into_iter()
272                .collect();
273
274            if !non_created_jobs.is_empty() {
275                obj_dependencies.retain(|dependency| {
276                    !non_created_jobs.contains(&dependency.object_id.as_job_id())
277                });
278            }
279        }
280    }
281
282    Ok(obj_dependencies)
283}
284
285/// List object dependencies with optional filtering of non-created streaming jobs.
286pub async fn list_object_dependencies(
287    txn: &DatabaseTransaction,
288    include_creating: bool,
289) -> MetaResult<Vec<PbObjectDependency>> {
290    list_object_dependencies_impl(txn, None, include_creating).await
291}
292
293/// List object dependencies for the specified object id without filtering by job status.
294pub async fn list_object_dependencies_by_object_id(
295    txn: &DatabaseTransaction,
296    object_id: ObjectId,
297) -> MetaResult<Vec<PbObjectDependency>> {
298    list_object_dependencies_impl(txn, Some(object_id), true).await
299}
300
301/// This function will construct a query using recursive cte to find if dependent objects are already relying on the target table.
302///
303/// # Examples
304///
305/// ```
306/// use risingwave_meta::controller::utils::construct_sink_cycle_check_query;
307/// use sea_orm::sea_query::*;
308/// use sea_orm::*;
309///
310/// let query = construct_sink_cycle_check_query(1.into(), vec![2.into(), 3.into()]);
311///
312/// assert_eq!(
313///     query.to_string(MysqlQueryBuilder),
314///     r#"WITH RECURSIVE `used_by_object_ids_with_sink` (`oid`, `used_by`) AS (SELECT `oid`, `used_by` FROM `object_dependency` WHERE `object_dependency`.`oid` = 1 UNION ALL (SELECT `obj_dependency_with_sink`.`oid`, `obj_dependency_with_sink`.`used_by` FROM (SELECT `oid`, `used_by` FROM `object_dependency` UNION ALL (SELECT `sink_id`, `target_table` FROM `sink` WHERE `sink`.`target_table` IS NOT NULL)) AS `obj_dependency_with_sink` INNER JOIN `used_by_object_ids_with_sink` ON `used_by_object_ids_with_sink`.`used_by` = `obj_dependency_with_sink`.`oid` WHERE `used_by_object_ids_with_sink`.`used_by` <> `used_by_object_ids_with_sink`.`oid`)) SELECT COUNT(`used_by_object_ids_with_sink`.`used_by`) FROM `used_by_object_ids_with_sink` WHERE `used_by_object_ids_with_sink`.`used_by` IN (2, 3)"#
315/// );
316/// assert_eq!(
317///     query.to_string(PostgresQueryBuilder),
318///     r#"WITH RECURSIVE "used_by_object_ids_with_sink" ("oid", "used_by") AS (SELECT "oid", "used_by" FROM "object_dependency" WHERE "object_dependency"."oid" = 1 UNION ALL (SELECT "obj_dependency_with_sink"."oid", "obj_dependency_with_sink"."used_by" FROM (SELECT "oid", "used_by" FROM "object_dependency" UNION ALL (SELECT "sink_id", "target_table" FROM "sink" WHERE "sink"."target_table" IS NOT NULL)) AS "obj_dependency_with_sink" INNER JOIN "used_by_object_ids_with_sink" ON "used_by_object_ids_with_sink"."used_by" = "obj_dependency_with_sink"."oid" WHERE "used_by_object_ids_with_sink"."used_by" <> "used_by_object_ids_with_sink"."oid")) SELECT COUNT("used_by_object_ids_with_sink"."used_by") FROM "used_by_object_ids_with_sink" WHERE "used_by_object_ids_with_sink"."used_by" IN (2, 3)"#
319/// );
320/// assert_eq!(
321///     query.to_string(SqliteQueryBuilder),
322///     r#"WITH RECURSIVE "used_by_object_ids_with_sink" ("oid", "used_by") AS (SELECT "oid", "used_by" FROM "object_dependency" WHERE "object_dependency"."oid" = 1 UNION ALL SELECT "obj_dependency_with_sink"."oid", "obj_dependency_with_sink"."used_by" FROM (SELECT "oid", "used_by" FROM "object_dependency" UNION ALL SELECT "sink_id", "target_table" FROM "sink" WHERE "sink"."target_table" IS NOT NULL) AS "obj_dependency_with_sink" INNER JOIN "used_by_object_ids_with_sink" ON "used_by_object_ids_with_sink"."used_by" = "obj_dependency_with_sink"."oid" WHERE "used_by_object_ids_with_sink"."used_by" <> "used_by_object_ids_with_sink"."oid") SELECT COUNT("used_by_object_ids_with_sink"."used_by") FROM "used_by_object_ids_with_sink" WHERE "used_by_object_ids_with_sink"."used_by" IN (2, 3)"#
323/// );
324/// ```
325pub fn construct_sink_cycle_check_query(
326    target_table: ObjectId,
327    dependent_objects: Vec<ObjectId>,
328) -> WithQuery {
329    let cte_alias = Alias::new("used_by_object_ids_with_sink");
330    let depend_alias = Alias::new("obj_dependency_with_sink");
331
332    let mut base_query = SelectStatement::new()
333        .columns([
334            object_dependency::Column::Oid,
335            object_dependency::Column::UsedBy,
336        ])
337        .from(ObjectDependency)
338        .and_where(object_dependency::Column::Oid.eq(target_table))
339        .to_owned();
340
341    let query_sink_deps = SelectStatement::new()
342        .columns([sink::Column::SinkId, sink::Column::TargetTable])
343        .from(Sink)
344        .and_where(sink::Column::TargetTable.is_not_null())
345        .to_owned();
346
347    let cte_referencing = Query::select()
348        .column((depend_alias.clone(), object_dependency::Column::Oid))
349        .column((depend_alias.clone(), object_dependency::Column::UsedBy))
350        .from_subquery(
351            SelectStatement::new()
352                .columns([
353                    object_dependency::Column::Oid,
354                    object_dependency::Column::UsedBy,
355                ])
356                .from(ObjectDependency)
357                .union(UnionType::All, query_sink_deps)
358                .to_owned(),
359            depend_alias.clone(),
360        )
361        .inner_join(
362            cte_alias.clone(),
363            Expr::col((cte_alias.clone(), object_dependency::Column::UsedBy))
364                .eq(Expr::col((depend_alias, object_dependency::Column::Oid))),
365        )
366        .and_where(
367            Expr::col((cte_alias.clone(), object_dependency::Column::UsedBy)).ne(Expr::col((
368                cte_alias.clone(),
369                object_dependency::Column::Oid,
370            ))),
371        )
372        .to_owned();
373
374    let mut common_table_expr = CommonTableExpression::new();
375    common_table_expr
376        .query(base_query.union(UnionType::All, cte_referencing).to_owned())
377        .columns([
378            object_dependency::Column::Oid,
379            object_dependency::Column::UsedBy,
380        ])
381        .table_name(cte_alias.clone());
382
383    SelectStatement::new()
384        .expr(Expr::col((cte_alias.clone(), object_dependency::Column::UsedBy)).count())
385        .from(cte_alias.clone())
386        .and_where(
387            Expr::col((cte_alias, object_dependency::Column::UsedBy)).is_in(dependent_objects),
388        )
389        .to_owned()
390        .with(
391            WithClause::new()
392                .recursive(true)
393                .cte(common_table_expr)
394                .to_owned(),
395        )
396}
397
398#[derive(Clone, DerivePartialModel, FromQueryResult, Debug)]
399#[sea_orm(entity = "Object")]
400pub struct PartialObject {
401    pub oid: ObjectId,
402    pub obj_type: ObjectType,
403    pub schema_id: Option<SchemaId>,
404    pub database_id: Option<DatabaseId>,
405}
406
407impl From<object::Model> for PartialObject {
408    fn from(object: object::Model) -> Self {
409        Self {
410            oid: object.oid,
411            obj_type: object.obj_type,
412            schema_id: object.schema_id,
413            database_id: object.database_id,
414        }
415    }
416}
417
418pub async fn get_belong_objects<C>(db: &C, object_id: ObjectId) -> MetaResult<Vec<object::Model>>
419where
420    C: ConnectionTrait,
421{
422    get_belong_objects_by_ids(db, [object_id]).await
423}
424
425/// Returns all objects that transitively belong to `object_ids`, excluding `object_ids`.
426pub async fn get_belong_objects_by_ids<C>(
427    db: &C,
428    object_ids: impl IntoIterator<Item = ObjectId>,
429) -> MetaResult<Vec<object::Model>>
430where
431    C: ConnectionTrait,
432{
433    let mut parents = object_ids.into_iter().collect_vec();
434    let mut visited = parents.iter().copied().collect::<HashSet<_>>();
435    let mut objects = vec![];
436    loop {
437        let children = Object::find()
438            .filter(object::Column::BelongToOid.is_in(parents))
439            .all(db)
440            .await?
441            .into_iter()
442            .filter(|object| visited.insert(object.oid))
443            .collect_vec();
444        if children.is_empty() {
445            break;
446        }
447
448        parents = children.iter().map(|object| object.oid).collect();
449        objects.extend(children);
450    }
451
452    Ok(objects)
453}
454
455#[derive(Clone, DerivePartialModel, FromQueryResult)]
456#[sea_orm(entity = "Fragment")]
457pub struct PartialFragmentStateTables {
458    pub fragment_id: FragmentId,
459    pub job_id: ObjectId,
460    pub state_table_ids: TableIdArray,
461}
462
463#[derive(Clone, Eq, PartialEq, Debug)]
464pub struct PartialActorLocation {
465    pub actor_id: ActorId,
466    pub fragment_id: FragmentId,
467    pub worker_id: WorkerId,
468}
469
470#[derive(FromQueryResult, Debug, Eq, PartialEq, Clone)]
471pub struct FragmentDesc {
472    pub fragment_id: FragmentId,
473    pub job_id: JobId,
474    pub fragment_type_mask: i32,
475    pub distribution_type: DistributionType,
476    pub state_table_ids: TableIdArray,
477    pub parallelism: i64,
478    pub vnode_count: i32,
479    pub stream_node: StreamNode,
480    pub parallelism_policy: String,
481}
482
483/// List all objects that are using the given one in a cascade way. It runs a recursive CTE to find all the dependencies.
484pub async fn get_referring_objects_cascade<C>(
485    obj_id: ObjectId,
486    object_type: ObjectType,
487    db: &C,
488) -> MetaResult<Vec<PartialObject>>
489where
490    C: ConnectionTrait,
491{
492    let query = construct_obj_dependency_query(obj_id);
493    let (sql, values) = query.build_any(&*db.get_database_backend().get_query_builder());
494    let mut objects = PartialObject::find_by_statement(Statement::from_sql_and_values(
495        db.get_database_backend(),
496        sql,
497        values,
498    ))
499    .all(db)
500    .await?;
501
502    let target_objects = std::iter::once((obj_id, object_type))
503        .chain(objects.iter().map(|object| (object.oid, object.obj_type)))
504        .collect_vec();
505    let mut existing_object_ids: HashSet<ObjectId> =
506        objects.iter().map(|object| object.oid).collect();
507    for (target_id, target_type) in target_objects {
508        let incoming_sink_objects =
509            get_incoming_sink_objects_for_target(target_id, target_type, db).await?;
510        for incoming_sink_object in incoming_sink_objects {
511            if existing_object_ids.insert(incoming_sink_object.oid) {
512                objects.push(incoming_sink_object);
513            }
514        }
515    }
516    Ok(objects)
517}
518
519/// Check if create a sink with given dependent objects into the target table will cause a cycle, return true if it will.
520pub async fn check_sink_into_table_cycle<C>(
521    target_table: ObjectId,
522    dependent_objs: Vec<ObjectId>,
523    db: &C,
524) -> MetaResult<bool>
525where
526    C: ConnectionTrait,
527{
528    if dependent_objs.is_empty() {
529        return Ok(false);
530    }
531
532    // special check for self referencing
533    if dependent_objs.contains(&target_table) {
534        return Ok(true);
535    }
536
537    let query = construct_sink_cycle_check_query(target_table, dependent_objs);
538    let (sql, values) = query.build_any(&*db.get_database_backend().get_query_builder());
539
540    let res = db
541        .query_one(Statement::from_sql_and_values(
542            db.get_database_backend(),
543            sql,
544            values,
545        ))
546        .await?
547        .unwrap();
548
549    let cnt: i64 = res.try_get_by(0)?;
550
551    Ok(cnt != 0)
552}
553
554/// Formats SQL options with secret values properly resolved.
555///
556/// This function processes configuration options that may contain sensitive data:
557/// - Plaintext options are directly converted to `SqlOption`.
558/// - Secret options are retrieved from the database and formatted as `SECRET {name}` without
559///   exposing the actual secret value.
560pub async fn format_with_option_secret_resolved(
561    txn: &DatabaseTransaction,
562    options_with_secret: &WithOptionsSecResolved,
563) -> MetaResult<Vec<SqlOption>> {
564    let mut options = Vec::new();
565    for (k, v) in options_with_secret.as_plaintext() {
566        let sql_option = SqlOption::try_from((k, v))
567            .map_err(|e| MetaError::invalid_parameter(e.to_report_string()))?;
568        options.push(sql_option);
569    }
570    for (name, secret_ref) in options_with_secret.as_secret() {
571        let secret_ref_value = resolve_secret_ref_value(txn, secret_ref).await?;
572        options.push(SqlOption::from_secret_ref(name, secret_ref_value));
573    }
574    Ok(options)
575}
576
577async fn resolve_secret_ref_value(
578    txn: &DatabaseTransaction,
579    secret_ref: &PbSecretRef,
580) -> MetaResult<SecretRefValue> {
581    let (secret, object) = Secret::find_by_id(secret_ref.secret_id)
582        .find_also_related(Object)
583        .one(txn)
584        .await?
585        .ok_or_else(|| MetaError::catalog_id_not_found("secret", secret_ref.secret_id))?;
586    let object =
587        object.ok_or_else(|| MetaError::catalog_id_not_found("object", secret_ref.secret_id))?;
588    let schema_id = object.schema_id.ok_or_else(|| {
589        MetaError::invalid_parameter(format!(
590            "secret {} does not belong to a schema",
591            secret_ref.secret_id
592        ))
593    })?;
594    let schema = Schema::find_by_id(schema_id)
595        .one(txn)
596        .await?
597        .ok_or_else(|| MetaError::catalog_id_not_found("schema", schema_id))?;
598
599    let ref_as = match PbRefAsType::try_from(secret_ref.ref_as) {
600        Ok(PbRefAsType::File) => SecretRefAsType::File,
601        Ok(PbRefAsType::Text | PbRefAsType::Unspecified) => SecretRefAsType::Text,
602        Err(_) => {
603            return Err(MetaError::invalid_parameter(format!(
604                "invalid reference type {} for secret {}",
605                secret_ref.ref_as, secret_ref.secret_id
606            )));
607        }
608    };
609
610    Ok(SecretRefValue {
611        secret_name: ObjectName(vec![
612            Ident::from_real_value(&schema.name),
613            Ident::from_real_value(&secret.name),
614        ]),
615        ref_as,
616    })
617}
618
619/// `ensure_object_id` ensures the existence of target object in the cluster.
620pub async fn ensure_object_id<C>(
621    object_type: ObjectType,
622    obj_id: impl Into<ObjectId>,
623    db: &C,
624) -> MetaResult<()>
625where
626    C: ConnectionTrait,
627{
628    let obj_id = obj_id.into();
629    let count = Object::find_by_id(obj_id).count(db).await?;
630    if count == 0 {
631        return Err(MetaError::catalog_id_not_found(
632            object_type.as_str(),
633            obj_id,
634        ));
635    }
636    Ok(())
637}
638
639pub async fn ensure_job_not_canceled<C>(job_id: JobId, db: &C) -> MetaResult<()>
640where
641    C: ConnectionTrait,
642{
643    let count = Object::find_by_id(job_id).count(db).await?;
644    if count == 0 {
645        return Err(MetaError::cancelled(format!(
646            "job {} might be cancelled manually or by recovery",
647            job_id
648        )));
649    }
650    Ok(())
651}
652
653/// `ensure_user_id` ensures the existence of target user in the cluster.
654pub async fn ensure_user_id<C>(user_id: UserId, db: &C) -> MetaResult<()>
655where
656    C: ConnectionTrait,
657{
658    let count = User::find_by_id(user_id).count(db).await?;
659    if count == 0 {
660        return Err(anyhow!("user {} was concurrently dropped", user_id).into());
661    }
662    Ok(())
663}
664
665/// `check_database_name_duplicate` checks whether the database name is already used in the cluster.
666pub async fn check_database_name_duplicate<C>(name: &str, db: &C) -> MetaResult<()>
667where
668    C: ConnectionTrait,
669{
670    let count = Database::find()
671        .filter(database::Column::Name.eq(name))
672        .count(db)
673        .await?;
674    if count > 0 {
675        assert_eq!(count, 1);
676        return Err(MetaError::catalog_duplicated("database", name));
677    }
678    Ok(())
679}
680
681/// `check_function_signature_duplicate` checks whether the function name and its signature is already used in the target namespace.
682pub async fn check_function_signature_duplicate<C>(
683    pb_function: &PbFunction,
684    db: &C,
685) -> MetaResult<()>
686where
687    C: ConnectionTrait,
688{
689    let count = Function::find()
690        .inner_join(Object)
691        .filter(
692            object::Column::DatabaseId
693                .eq(pb_function.database_id)
694                .and(object::Column::SchemaId.eq(pb_function.schema_id))
695                .and(function::Column::Name.eq(&pb_function.name))
696                .and(
697                    function::Column::ArgTypes
698                        .eq(DataTypeArray::from(pb_function.arg_types.clone())),
699                ),
700        )
701        .count(db)
702        .await?;
703    if count > 0 {
704        assert_eq!(count, 1);
705        return Err(MetaError::catalog_duplicated("function", &pb_function.name));
706    }
707    Ok(())
708}
709
710/// `check_connection_name_duplicate` checks whether the connection name is already used in the target namespace.
711pub async fn check_connection_name_duplicate<C>(
712    pb_connection: &PbConnection,
713    db: &C,
714) -> MetaResult<()>
715where
716    C: ConnectionTrait,
717{
718    let count = Connection::find()
719        .inner_join(Object)
720        .filter(
721            object::Column::DatabaseId
722                .eq(pb_connection.database_id)
723                .and(object::Column::SchemaId.eq(pb_connection.schema_id))
724                .and(connection::Column::Name.eq(&pb_connection.name)),
725        )
726        .count(db)
727        .await?;
728    if count > 0 {
729        assert_eq!(count, 1);
730        return Err(MetaError::catalog_duplicated(
731            "connection",
732            &pb_connection.name,
733        ));
734    }
735    Ok(())
736}
737
738pub async fn check_secret_name_duplicate<C>(pb_secret: &PbSecret, db: &C) -> MetaResult<()>
739where
740    C: ConnectionTrait,
741{
742    let count = Secret::find()
743        .inner_join(Object)
744        .filter(
745            object::Column::DatabaseId
746                .eq(pb_secret.database_id)
747                .and(object::Column::SchemaId.eq(pb_secret.schema_id))
748                .and(secret::Column::Name.eq(&pb_secret.name)),
749        )
750        .count(db)
751        .await?;
752    if count > 0 {
753        assert_eq!(count, 1);
754        return Err(MetaError::catalog_duplicated("secret", &pb_secret.name));
755    }
756    Ok(())
757}
758
759pub async fn check_subscription_name_duplicate<C>(
760    pb_subscription: &PbSubscription,
761    db: &C,
762) -> MetaResult<()>
763where
764    C: ConnectionTrait,
765{
766    let count = Subscription::find()
767        .inner_join(Object)
768        .filter(
769            object::Column::DatabaseId
770                .eq(pb_subscription.database_id)
771                .and(object::Column::SchemaId.eq(pb_subscription.schema_id))
772                .and(subscription::Column::Name.eq(&pb_subscription.name)),
773        )
774        .count(db)
775        .await?;
776    if count > 0 {
777        assert_eq!(count, 1);
778        return Err(MetaError::catalog_duplicated(
779            "subscription",
780            &pb_subscription.name,
781        ));
782    }
783    Ok(())
784}
785
786/// `check_user_name_duplicate` checks whether the user is already existed in the cluster.
787pub async fn check_user_name_duplicate<C>(name: &str, db: &C) -> MetaResult<()>
788where
789    C: ConnectionTrait,
790{
791    let count = User::find()
792        .filter(user::Column::Name.eq(name))
793        .count(db)
794        .await?;
795    if count > 0 {
796        assert_eq!(count, 1);
797        return Err(MetaError::catalog_duplicated("user", name));
798    }
799    Ok(())
800}
801
802/// `check_relation_name_duplicate` checks whether the relation name is already used in the target namespace.
803pub async fn check_relation_name_duplicate<C>(
804    name: &str,
805    database_id: DatabaseId,
806    schema_id: SchemaId,
807    db: &C,
808) -> MetaResult<()>
809where
810    C: ConnectionTrait,
811{
812    macro_rules! check_duplicated {
813        ($obj_type:expr, $entity:ident, $table:ident) => {
814            let object_id = Object::find()
815                .select_only()
816                .column(object::Column::Oid)
817                .inner_join($entity)
818                .filter(
819                    object::Column::DatabaseId
820                        .eq(Some(database_id))
821                        .and(object::Column::SchemaId.eq(Some(schema_id)))
822                        .and($table::Column::Name.eq(name)),
823                )
824                .into_tuple::<ObjectId>()
825                .one(db)
826                .await?;
827            if let Some(oid) = object_id {
828                let check_creation = if $obj_type == ObjectType::View {
829                    false
830                } else if $obj_type == ObjectType::Source {
831                    let source_info = Source::find_by_id(oid.as_source_id())
832                        .select_only()
833                        .column(source::Column::SourceInfo)
834                        .into_tuple::<Option<StreamSourceInfo>>()
835                        .one(db)
836                        .await?
837                        .unwrap();
838                    source_info.map_or(false, |info| info.to_protobuf().is_shared())
839                } else {
840                    true
841                };
842                let job_id = oid.as_job_id();
843                return if check_creation
844                    && !matches!(
845                        StreamingJob::find_by_id(job_id)
846                            .select_only()
847                            .column(streaming_job::Column::JobStatus)
848                            .into_tuple::<JobStatus>()
849                            .one(db)
850                            .await?,
851                        Some(JobStatus::Created)
852                    ) {
853                    Err(MetaError::catalog_under_creation(
854                        $obj_type.as_str(),
855                        name,
856                        job_id,
857                    ))
858                } else {
859                    Err(MetaError::catalog_duplicated($obj_type.as_str(), name))
860                };
861            }
862        };
863    }
864    check_duplicated!(ObjectType::Table, Table, table);
865    check_duplicated!(ObjectType::Source, Source, source);
866    check_duplicated!(ObjectType::Sink, Sink, sink);
867    check_duplicated!(ObjectType::Index, Index, index);
868    check_duplicated!(ObjectType::View, View, view);
869
870    Ok(())
871}
872
873/// `check_schema_name_duplicate` checks whether the schema name is already used in the target database.
874pub async fn check_schema_name_duplicate<C>(
875    name: &str,
876    database_id: DatabaseId,
877    db: &C,
878) -> MetaResult<()>
879where
880    C: ConnectionTrait,
881{
882    let count = Object::find()
883        .inner_join(Schema)
884        .filter(
885            object::Column::ObjType
886                .eq(ObjectType::Schema)
887                .and(object::Column::DatabaseId.eq(Some(database_id)))
888                .and(schema::Column::Name.eq(name)),
889        )
890        .count(db)
891        .await?;
892    if count != 0 {
893        return Err(MetaError::catalog_duplicated("schema", name));
894    }
895
896    Ok(())
897}
898
899/// Validates a RESTRICT drop and returns referring objects that should be dropped together.
900///
901/// Objects related through `belong_to_oid` are catalog-owned by the target. Indexes are also
902/// logically owned by their primary table for dropping, although they remain schema-owned catalog
903/// objects. These objects are dropped together with the target; all other referrers prevent the
904/// drop.
905pub async fn validate_restrict_drop_and_collect_owned_objects<C>(
906    object_type: ObjectType,
907    object_id: ObjectId,
908    db: &C,
909) -> MetaResult<Vec<PartialObject>>
910where
911    C: ConnectionTrait,
912{
913    let referring_objects = get_referring_objects(object_id, object_type, db).await?;
914    let owned_object_ids = get_belong_objects(db, object_id)
915        .await?
916        .into_iter()
917        .map(|object| object.oid)
918        .collect::<HashSet<_>>();
919    let mut non_owned_objects = referring_objects.clone();
920    non_owned_objects.retain(|object| !owned_object_ids.contains(&object.oid));
921    if object_type == ObjectType::Table {
922        non_owned_objects.retain(|object| object.obj_type != ObjectType::Index);
923    }
924
925    if !non_owned_objects.is_empty() {
926        let referring_objs_map = non_owned_objects
927            .into_iter()
928            .into_group_map_by(|o| o.obj_type);
929        let mut details = vec![];
930        for (obj_type, objs) in referring_objs_map {
931            match obj_type {
932                ObjectType::Table => {
933                    let tables: Vec<(String, String)> = Object::find()
934                        .join(JoinType::InnerJoin, object::Relation::Table.def())
935                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
936                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
937                        .select_only()
938                        .column(schema::Column::Name)
939                        .column(table::Column::Name)
940                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
941                        .into_tuple()
942                        .all(db)
943                        .await?;
944                    details.extend(tables.into_iter().map(|(schema_name, table_name)| {
945                        format!(
946                            "materialized view {}.{} depends on it",
947                            schema_name, table_name
948                        )
949                    }));
950                }
951                ObjectType::Sink => {
952                    let sinks: Vec<(String, String)> = Object::find()
953                        .join(JoinType::InnerJoin, object::Relation::Sink.def())
954                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
955                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
956                        .select_only()
957                        .column(schema::Column::Name)
958                        .column(sink::Column::Name)
959                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
960                        .into_tuple()
961                        .all(db)
962                        .await?;
963                    details.extend(sinks.into_iter().map(|(schema_name, sink_name)| {
964                        format!("sink {}.{} depends on it", schema_name, sink_name)
965                    }));
966                }
967                ObjectType::View => {
968                    let views: Vec<(String, String)> = Object::find()
969                        .join(JoinType::InnerJoin, object::Relation::View.def())
970                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
971                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
972                        .select_only()
973                        .column(schema::Column::Name)
974                        .column(view::Column::Name)
975                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
976                        .into_tuple()
977                        .all(db)
978                        .await?;
979                    details.extend(views.into_iter().map(|(schema_name, view_name)| {
980                        format!("view {}.{} depends on it", schema_name, view_name)
981                    }));
982                }
983                ObjectType::Subscription => {
984                    let subscriptions: Vec<(String, String)> = Object::find()
985                        .join(JoinType::InnerJoin, object::Relation::Subscription.def())
986                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
987                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
988                        .select_only()
989                        .column(schema::Column::Name)
990                        .column(subscription::Column::Name)
991                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
992                        .into_tuple()
993                        .all(db)
994                        .await?;
995                    details.extend(subscriptions.into_iter().map(
996                        |(schema_name, subscription_name)| {
997                            format!(
998                                "subscription {}.{} depends on it",
999                                schema_name, subscription_name
1000                            )
1001                        },
1002                    ));
1003                }
1004                ObjectType::Source => {
1005                    let sources: Vec<(String, String)> = Object::find()
1006                        .join(JoinType::InnerJoin, object::Relation::Source.def())
1007                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
1008                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
1009                        .select_only()
1010                        .column(schema::Column::Name)
1011                        .column(source::Column::Name)
1012                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
1013                        .into_tuple()
1014                        .all(db)
1015                        .await?;
1016                    details.extend(sources.into_iter().map(|(schema_name, view_name)| {
1017                        format!("source {}.{} depends on it", schema_name, view_name)
1018                    }));
1019                }
1020                ObjectType::Connection => {
1021                    let connections: Vec<(String, String)> = Object::find()
1022                        .join(JoinType::InnerJoin, object::Relation::Connection.def())
1023                        .join(JoinType::InnerJoin, object::Relation::Database2.def())
1024                        .join(JoinType::InnerJoin, object::Relation::Schema2.def())
1025                        .select_only()
1026                        .column(schema::Column::Name)
1027                        .column(connection::Column::Name)
1028                        .filter(object::Column::Oid.is_in(objs.iter().map(|o| o.oid)))
1029                        .into_tuple()
1030                        .all(db)
1031                        .await?;
1032                    details.extend(connections.into_iter().map(|(schema_name, view_name)| {
1033                        format!("connection {}.{} depends on it", schema_name, view_name)
1034                    }));
1035                }
1036                // only the table, source, sink, subscription, view, connection and index will depend on other objects.
1037                _ => bail!("unexpected referring object type: {}", obj_type.as_str()),
1038            }
1039        }
1040        if details.is_empty() {
1041            return Ok(referring_objects);
1042        }
1043
1044        return Err(MetaError::permission_denied(format!(
1045            "{} used by {} other objects. \nDETAIL: {}\n\
1046            {}",
1047            object_type.as_str(),
1048            details.len(),
1049            details.join("\n"),
1050            match object_type {
1051                ObjectType::Function | ObjectType::Connection | ObjectType::Secret =>
1052                    "HINT: DROP the dependent objects first.",
1053                ObjectType::Database | ObjectType::Schema => unreachable!(),
1054                _ => "HINT:  Use DROP ... CASCADE to drop the dependent objects too.",
1055            }
1056        )));
1057    }
1058    Ok(referring_objects)
1059}
1060
1061async fn get_incoming_sink_objects_for_target<C>(
1062    target_id: ObjectId,
1063    target_type: ObjectType,
1064    db: &C,
1065) -> MetaResult<Vec<PartialObject>>
1066where
1067    C: ConnectionTrait,
1068{
1069    match target_type {
1070        ObjectType::Table => {
1071            let incoming_sink_ids = Sink::find()
1072                .select_only()
1073                .column(sink::Column::SinkId)
1074                .filter(sink::Column::TargetTable.eq(target_id.as_table_id()))
1075                .into_tuple::<SinkId>()
1076                .all(db)
1077                .await?;
1078            Object::find()
1079                .filter(object::Column::Oid.is_in(incoming_sink_ids))
1080                .into_partial_model()
1081                .all(db)
1082                .await
1083                .map_err(Into::into)
1084        }
1085        _ => Ok(vec![]),
1086    }
1087}
1088
1089/// List all objects that are using the given one.
1090pub async fn get_referring_objects<C>(
1091    object_id: ObjectId,
1092    object_type: ObjectType,
1093    db: &C,
1094) -> MetaResult<Vec<PartialObject>>
1095where
1096    C: ConnectionTrait,
1097{
1098    let mut objects: Vec<PartialObject> = ObjectDependency::find()
1099        .filter(object_dependency::Column::Oid.eq(object_id))
1100        .join(
1101            JoinType::InnerJoin,
1102            object_dependency::Relation::Object1.def(),
1103        )
1104        .into_partial_model()
1105        .all(db)
1106        .await?;
1107
1108    let incoming_sink_objects =
1109        get_incoming_sink_objects_for_target(object_id, object_type, db).await?;
1110    let mut existing_object_ids: HashSet<ObjectId> =
1111        objects.iter().map(|object| object.oid).collect();
1112    objects.extend(
1113        incoming_sink_objects
1114            .into_iter()
1115            .filter(|object| existing_object_ids.insert(object.oid)),
1116    );
1117
1118    Ok(objects)
1119}
1120
1121/// `ensure_schema_empty` ensures that the schema is empty, used by `DROP SCHEMA`.
1122pub async fn ensure_schema_empty<C>(schema_id: SchemaId, db: &C) -> MetaResult<()>
1123where
1124    C: ConnectionTrait,
1125{
1126    let count = Object::find()
1127        .filter(object::Column::SchemaId.eq(Some(schema_id)))
1128        .count(db)
1129        .await?;
1130    if count != 0 {
1131        return Err(MetaError::permission_denied("schema is not empty"));
1132    }
1133
1134    Ok(())
1135}
1136
1137/// `list_user_info_by_ids` lists all users' info by their ids.
1138pub async fn list_user_info_by_ids<C>(
1139    user_ids: impl IntoIterator<Item = UserId>,
1140    db: &C,
1141) -> MetaResult<Vec<PbUserInfo>>
1142where
1143    C: ConnectionTrait,
1144{
1145    let mut user_infos = vec![];
1146    for user_id in user_ids {
1147        let user = User::find_by_id(user_id)
1148            .one(db)
1149            .await?
1150            .ok_or_else(|| MetaError::catalog_id_not_found("user", user_id))?;
1151        let mut user_info: PbUserInfo = user.into();
1152        user_info.grant_privileges = get_user_privilege(user_id, db).await?;
1153        user_infos.push(user_info);
1154    }
1155    Ok(user_infos)
1156}
1157
1158/// `get_object_owner` returns the owner of the given object.
1159pub async fn get_object_owner<C>(object_id: ObjectId, db: &C) -> MetaResult<UserId>
1160where
1161    C: ConnectionTrait,
1162{
1163    let obj_owner: UserId = Object::find_by_id(object_id)
1164        .select_only()
1165        .column(object::Column::OwnerId)
1166        .into_tuple()
1167        .one(db)
1168        .await?
1169        .ok_or_else(|| MetaError::catalog_id_not_found("object", object_id))?;
1170    Ok(obj_owner)
1171}
1172
1173/// `construct_privilege_dependency_query` constructs a query to find all privileges that are dependent on the given one.
1174///
1175/// # Examples
1176///
1177/// ```
1178/// use risingwave_meta::controller::utils::construct_privilege_dependency_query;
1179/// use sea_orm::sea_query::*;
1180/// use sea_orm::*;
1181///
1182/// let query = construct_privilege_dependency_query(vec![1, 2, 3]);
1183///
1184/// assert_eq!(
1185///    query.to_string(MysqlQueryBuilder),
1186///   r#"WITH RECURSIVE `granted_privilege_ids` (`id`, `user_id`) AS (SELECT `id`, `user_id` FROM `user_privilege` WHERE `user_privilege`.`id` IN (1, 2, 3) UNION ALL (SELECT `user_privilege`.`id`, `user_privilege`.`user_id` FROM `user_privilege` INNER JOIN `granted_privilege_ids` ON `granted_privilege_ids`.`id` = `dependent_id`)) SELECT `id`, `user_id` FROM `granted_privilege_ids`"#
1187/// );
1188/// assert_eq!(
1189///   query.to_string(PostgresQueryBuilder),
1190///  r#"WITH RECURSIVE "granted_privilege_ids" ("id", "user_id") AS (SELECT "id", "user_id" FROM "user_privilege" WHERE "user_privilege"."id" IN (1, 2, 3) UNION ALL (SELECT "user_privilege"."id", "user_privilege"."user_id" FROM "user_privilege" INNER JOIN "granted_privilege_ids" ON "granted_privilege_ids"."id" = "dependent_id")) SELECT "id", "user_id" FROM "granted_privilege_ids""#
1191/// );
1192/// assert_eq!(
1193///  query.to_string(SqliteQueryBuilder),
1194///  r#"WITH RECURSIVE "granted_privilege_ids" ("id", "user_id") AS (SELECT "id", "user_id" FROM "user_privilege" WHERE "user_privilege"."id" IN (1, 2, 3) UNION ALL SELECT "user_privilege"."id", "user_privilege"."user_id" FROM "user_privilege" INNER JOIN "granted_privilege_ids" ON "granted_privilege_ids"."id" = "dependent_id") SELECT "id", "user_id" FROM "granted_privilege_ids""#
1195/// );
1196/// ```
1197pub fn construct_privilege_dependency_query(ids: Vec<PrivilegeId>) -> WithQuery {
1198    let cte_alias = Alias::new("granted_privilege_ids");
1199    let cte_return_privilege_alias = Alias::new("id");
1200    let cte_return_user_alias = Alias::new("user_id");
1201
1202    let mut base_query = SelectStatement::new()
1203        .columns([user_privilege::Column::Id, user_privilege::Column::UserId])
1204        .from(UserPrivilege)
1205        .and_where(user_privilege::Column::Id.is_in(ids))
1206        .to_owned();
1207
1208    let cte_referencing = Query::select()
1209        .columns([
1210            (UserPrivilege, user_privilege::Column::Id),
1211            (UserPrivilege, user_privilege::Column::UserId),
1212        ])
1213        .from(UserPrivilege)
1214        .inner_join(
1215            cte_alias.clone(),
1216            Expr::col((cte_alias.clone(), cte_return_privilege_alias.clone()))
1217                .equals(user_privilege::Column::DependentId),
1218        )
1219        .to_owned();
1220
1221    let mut common_table_expr = CommonTableExpression::new();
1222    common_table_expr
1223        .query(base_query.union(UnionType::All, cte_referencing).to_owned())
1224        .columns([
1225            cte_return_privilege_alias.clone(),
1226            cte_return_user_alias.clone(),
1227        ])
1228        .table_name(cte_alias.clone());
1229
1230    SelectStatement::new()
1231        .columns([cte_return_privilege_alias, cte_return_user_alias])
1232        .from(cte_alias)
1233        .to_owned()
1234        .with(
1235            WithClause::new()
1236                .recursive(true)
1237                .cte(common_table_expr)
1238                .to_owned(),
1239        )
1240}
1241
1242pub async fn get_internal_tables_by_id<C>(job_id: JobId, db: &C) -> MetaResult<Vec<TableId>>
1243where
1244    C: ConnectionTrait,
1245{
1246    let table_ids: Vec<TableId> = Table::find()
1247        .select_only()
1248        .column(table::Column::TableId)
1249        .filter(
1250            table::Column::TableType
1251                .eq(TableType::Internal)
1252                .and(table::Column::BelongsToJobId.eq(job_id)),
1253        )
1254        .into_tuple()
1255        .all(db)
1256        .await?;
1257    Ok(table_ids)
1258}
1259
1260pub async fn get_index_state_tables_by_table_id<C>(
1261    table_id: TableId,
1262    db: &C,
1263) -> MetaResult<Vec<TableId>>
1264where
1265    C: ConnectionTrait,
1266{
1267    let mut index_table_ids: Vec<TableId> = Index::find()
1268        .select_only()
1269        .column(index::Column::IndexTableId)
1270        .filter(index::Column::PrimaryTableId.eq(table_id))
1271        .into_tuple()
1272        .all(db)
1273        .await?;
1274
1275    if !index_table_ids.is_empty() {
1276        let internal_table_ids: Vec<TableId> = Table::find()
1277            .select_only()
1278            .column(table::Column::TableId)
1279            .filter(
1280                table::Column::TableType
1281                    .eq(TableType::Internal)
1282                    .and(table::Column::BelongsToJobId.is_in(index_table_ids.clone())),
1283            )
1284            .into_tuple()
1285            .all(db)
1286            .await?;
1287
1288        index_table_ids.extend(internal_table_ids);
1289    }
1290
1291    Ok(index_table_ids)
1292}
1293
1294/// `get_iceberg_related_object_ids` returns the related object ids of the iceberg source, sink and internal tables.
1295pub async fn get_iceberg_related_object_ids<C>(
1296    object_id: ObjectId,
1297    db: &C,
1298) -> MetaResult<Vec<ObjectId>>
1299where
1300    C: ConnectionTrait,
1301{
1302    let object = Object::find_by_id(object_id)
1303        .one(db)
1304        .await?
1305        .ok_or_else(|| MetaError::catalog_id_not_found("object", object_id))?;
1306    if object.obj_type != ObjectType::Table {
1307        return Ok(vec![]);
1308    }
1309
1310    let table = Table::find_by_id(object_id.as_table_id())
1311        .one(db)
1312        .await?
1313        .ok_or_else(|| MetaError::catalog_id_not_found("table", object_id))?;
1314    if !matches!(table.engine, Some(table::Engine::Iceberg)) {
1315        return Ok(vec![]);
1316    }
1317
1318    let database_id = object.database_id.unwrap();
1319    let schema_id = object.schema_id.unwrap();
1320
1321    let mut related_objects = vec![];
1322
1323    let iceberg_sink_name = format!("{}{}", ICEBERG_SINK_PREFIX, table.name);
1324    let iceberg_sink_id = Sink::find()
1325        .inner_join(Object)
1326        .select_only()
1327        .column(sink::Column::SinkId)
1328        .filter(
1329            object::Column::DatabaseId
1330                .eq(database_id)
1331                .and(object::Column::SchemaId.eq(schema_id))
1332                .and(sink::Column::Name.eq(&iceberg_sink_name)),
1333        )
1334        .into_tuple::<SinkId>()
1335        .one(db)
1336        .await?;
1337    if let Some(sink_id) = iceberg_sink_id {
1338        related_objects.push(sink_id.as_object_id());
1339        let sink_internal_tables = get_internal_tables_by_id(sink_id.as_job_id(), db).await?;
1340        related_objects.extend(
1341            sink_internal_tables
1342                .into_iter()
1343                .map(|tid| tid.as_object_id()),
1344        );
1345    } else {
1346        warn!(
1347            "iceberg table {} missing sink {}",
1348            table.name, iceberg_sink_name
1349        );
1350    }
1351
1352    let iceberg_source_name = format!("{}{}", ICEBERG_SOURCE_PREFIX, table.name);
1353    let iceberg_source_id = Source::find()
1354        .inner_join(Object)
1355        .select_only()
1356        .column(source::Column::SourceId)
1357        .filter(
1358            object::Column::DatabaseId
1359                .eq(database_id)
1360                .and(object::Column::SchemaId.eq(schema_id))
1361                .and(source::Column::Name.eq(&iceberg_source_name)),
1362        )
1363        .into_tuple::<SourceId>()
1364        .one(db)
1365        .await?;
1366    if let Some(source_id) = iceberg_source_id {
1367        related_objects.push(source_id.as_object_id());
1368    } else {
1369        warn!(
1370            "iceberg table {} missing source {}",
1371            table.name, iceberg_source_name
1372        );
1373    }
1374
1375    Ok(related_objects)
1376}
1377
1378/// Load streaming jobs by job ids.
1379pub(crate) async fn load_streaming_jobs_by_ids<C>(
1380    txn: &C,
1381    job_ids: impl IntoIterator<Item = JobId>,
1382) -> MetaResult<HashMap<JobId, streaming_job::Model>>
1383where
1384    C: ConnectionTrait,
1385{
1386    let job_ids: HashSet<JobId> = job_ids.into_iter().collect();
1387    if job_ids.is_empty() {
1388        return Ok(HashMap::new());
1389    }
1390    let jobs = streaming_job::Entity::find()
1391        .filter(streaming_job::Column::JobId.is_in(job_ids.clone()))
1392        .all(txn)
1393        .await?;
1394    Ok(jobs.into_iter().map(|job| (job.job_id, job)).collect())
1395}
1396
1397#[derive(Clone, DerivePartialModel, FromQueryResult)]
1398#[sea_orm(entity = "UserPrivilege")]
1399pub struct PartialUserPrivilege {
1400    pub id: PrivilegeId,
1401    pub user_id: UserId,
1402}
1403
1404pub async fn get_referring_privileges_cascade<C>(
1405    ids: Vec<PrivilegeId>,
1406    db: &C,
1407) -> MetaResult<Vec<PartialUserPrivilege>>
1408where
1409    C: ConnectionTrait,
1410{
1411    let query = construct_privilege_dependency_query(ids);
1412    let (sql, values) = query.build_any(&*db.get_database_backend().get_query_builder());
1413    let privileges = PartialUserPrivilege::find_by_statement(Statement::from_sql_and_values(
1414        db.get_database_backend(),
1415        sql,
1416        values,
1417    ))
1418    .all(db)
1419    .await?;
1420
1421    Ok(privileges)
1422}
1423
1424/// `ensure_privileges_not_referred` ensures that the privileges are not granted to any other users.
1425pub async fn ensure_privileges_not_referred<C>(ids: Vec<PrivilegeId>, db: &C) -> MetaResult<()>
1426where
1427    C: ConnectionTrait,
1428{
1429    let count = UserPrivilege::find()
1430        .filter(user_privilege::Column::DependentId.is_in(ids))
1431        .count(db)
1432        .await?;
1433    if count != 0 {
1434        return Err(MetaError::permission_denied(format!(
1435            "privileges granted to {} other ones.",
1436            count
1437        )));
1438    }
1439    Ok(())
1440}
1441
1442/// `get_user_privilege` returns the privileges of the given user.
1443pub async fn get_user_privilege<C>(user_id: UserId, db: &C) -> MetaResult<Vec<PbGrantPrivilege>>
1444where
1445    C: ConnectionTrait,
1446{
1447    let user_privileges = UserPrivilege::find()
1448        .find_also_related(Object)
1449        .filter(user_privilege::Column::UserId.eq(user_id))
1450        .all(db)
1451        .await?;
1452    Ok(user_privileges
1453        .into_iter()
1454        .map(|(privilege, object)| {
1455            let object = object.unwrap();
1456            let obj = match object.obj_type {
1457                ObjectType::Database => PbGrantObject::DatabaseId(object.oid.as_database_id()),
1458                ObjectType::Schema => PbGrantObject::SchemaId(object.oid.as_schema_id()),
1459                ObjectType::Table | ObjectType::Index => {
1460                    PbGrantObject::TableId(object.oid.as_table_id())
1461                }
1462                ObjectType::Source => PbGrantObject::SourceId(object.oid.as_source_id()),
1463                ObjectType::Sink => PbGrantObject::SinkId(object.oid.as_sink_id()),
1464                ObjectType::View => PbGrantObject::ViewId(object.oid.as_view_id()),
1465                ObjectType::Function => PbGrantObject::FunctionId(object.oid.as_function_id()),
1466                ObjectType::Connection => {
1467                    PbGrantObject::ConnectionId(object.oid.as_connection_id())
1468                }
1469                ObjectType::Subscription => {
1470                    PbGrantObject::SubscriptionId(object.oid.as_subscription_id())
1471                }
1472                ObjectType::Secret => PbGrantObject::SecretId(object.oid.as_secret_id()),
1473            };
1474            PbGrantPrivilege {
1475                action_with_opts: vec![PbActionWithGrantOption {
1476                    action: PbAction::from(privilege.action) as _,
1477                    with_grant_option: privilege.with_grant_option,
1478                    granted_by: privilege.granted_by as _,
1479                }],
1480                object: Some(obj),
1481            }
1482        })
1483        .collect())
1484}
1485
1486/// Insert user privileges idempotently using the same key as `idx_user_privilege_item`.
1487/// An existing privilege is preserved, except that its grant option is upgraded when requested.
1488pub(crate) async fn upsert_user_privileges<C>(
1489    db: &C,
1490    privileges: impl IntoIterator<Item = user_privilege::ActiveModel>,
1491) -> MetaResult<()>
1492where
1493    C: ConnectionTrait,
1494{
1495    for privilege in privileges {
1496        let mut on_conflict = OnConflict::columns([
1497            user_privilege::Column::UserId,
1498            user_privilege::Column::Oid,
1499            user_privilege::Column::Action,
1500            user_privilege::Column::GrantedBy,
1501        ]);
1502        if *privilege.with_grant_option.as_ref() {
1503            on_conflict.update_column(user_privilege::Column::WithGrantOption);
1504        } else {
1505            // Workaround to support MYSQL for `DO NOTHING`.
1506            on_conflict.update_column(user_privilege::Column::UserId);
1507        }
1508
1509        UserPrivilege::insert(privilege)
1510            .on_conflict(on_conflict)
1511            .do_nothing()
1512            .exec(db)
1513            .await?;
1514    }
1515    Ok(())
1516}
1517
1518pub async fn get_table_columns(
1519    txn: &impl ConnectionTrait,
1520    id: TableId,
1521) -> MetaResult<ColumnCatalogArray> {
1522    let columns = Table::find_by_id(id)
1523        .select_only()
1524        .columns([table::Column::Columns])
1525        .into_tuple::<ColumnCatalogArray>()
1526        .one(txn)
1527        .await?
1528        .ok_or_else(|| MetaError::catalog_id_not_found("table", id))?;
1529    Ok(columns)
1530}
1531
1532/// `grant_default_privileges_automatically` grants default privileges automatically
1533/// for the given new object. It returns the list of user infos whose privileges are updated.
1534pub async fn grant_default_privileges_automatically<C>(
1535    db: &C,
1536    object_id: impl Into<ObjectId>,
1537) -> MetaResult<Vec<PbUserInfo>>
1538where
1539    C: ConnectionTrait,
1540{
1541    let object_id = object_id.into();
1542    let object = Object::find_by_id(object_id)
1543        .one(db)
1544        .await?
1545        .ok_or_else(|| MetaError::catalog_id_not_found("object", object_id))?;
1546    assert_ne!(object.obj_type, ObjectType::Database);
1547
1548    let for_mview_filter = if object.obj_type == ObjectType::Table {
1549        let table_type = Table::find_by_id(object_id.as_table_id())
1550            .select_only()
1551            .column(table::Column::TableType)
1552            .into_tuple::<TableType>()
1553            .one(db)
1554            .await?
1555            .ok_or_else(|| MetaError::catalog_id_not_found("table", object_id))?;
1556        user_default_privilege::Column::ForMaterializedView
1557            .eq(table_type == TableType::MaterializedView)
1558    } else {
1559        user_default_privilege::Column::ForMaterializedView.eq(false)
1560    };
1561    let schema_filter = if let Some(schema_id) = &object.schema_id {
1562        user_default_privilege::Column::SchemaId.eq(*schema_id)
1563    } else {
1564        user_default_privilege::Column::SchemaId.is_null()
1565    };
1566
1567    let default_privileges: Vec<(UserId, UserId, Action, bool)> = UserDefaultPrivilege::find()
1568        .select_only()
1569        .columns([
1570            user_default_privilege::Column::Grantee,
1571            user_default_privilege::Column::GrantedBy,
1572            user_default_privilege::Column::Action,
1573            user_default_privilege::Column::WithGrantOption,
1574        ])
1575        .filter(
1576            user_default_privilege::Column::DatabaseId
1577                .eq(object.database_id.unwrap())
1578                .and(schema_filter)
1579                .and(user_default_privilege::Column::UserId.eq(object.owner_id))
1580                .and(user_default_privilege::Column::ObjectType.eq(object.obj_type))
1581                .and(for_mview_filter),
1582        )
1583        .into_tuple()
1584        .all(db)
1585        .await?;
1586    if default_privileges.is_empty() {
1587        return Ok(vec![]);
1588    }
1589
1590    let updated_user_ids = default_privileges
1591        .iter()
1592        .map(|(grantee, _, _, _)| *grantee)
1593        .collect::<HashSet<_>>();
1594
1595    let mut new_privileges = vec![];
1596    for (grantee, granted_by, action, with_grant_option) in default_privileges {
1597        new_privileges.push(user_privilege::ActiveModel {
1598            user_id: Set(grantee),
1599            oid: Set(object_id),
1600            granted_by: Set(granted_by),
1601            action: Set(action),
1602            with_grant_option: Set(with_grant_option),
1603            ..Default::default()
1604        });
1605        if action == Action::Select {
1606            // Grant SELECT privilege for internal tables if the action is SELECT.
1607            let internal_table_ids = get_internal_tables_by_id(object_id.as_job_id(), db).await?;
1608            new_privileges.extend(internal_table_ids.into_iter().map(|internal_table_id| {
1609                user_privilege::ActiveModel {
1610                    user_id: Set(grantee),
1611                    oid: Set(internal_table_id.as_object_id()),
1612                    granted_by: Set(granted_by),
1613                    action: Set(Action::Select),
1614                    with_grant_option: Set(with_grant_option),
1615                    ..Default::default()
1616                }
1617            }));
1618
1619            // Additionally, grant SELECT privilege for iceberg related objects if the action is SELECT.
1620            let iceberg_privilege_object_ids =
1621                get_iceberg_related_object_ids(object_id, db).await?;
1622            new_privileges.extend(iceberg_privilege_object_ids.into_iter().map(
1623                |iceberg_object_id| user_privilege::ActiveModel {
1624                    user_id: Set(grantee),
1625                    oid: Set(iceberg_object_id),
1626                    granted_by: Set(granted_by),
1627                    action: Set(action),
1628                    with_grant_option: Set(with_grant_option),
1629                    ..Default::default()
1630                },
1631            ));
1632        }
1633    }
1634    upsert_user_privileges(db, new_privileges).await?;
1635
1636    let updated_user_infos = list_user_info_by_ids(updated_user_ids, db).await?;
1637    Ok(updated_user_infos)
1638}
1639
1640// todo: remove it after migrated to sql backend.
1641pub fn extract_grant_obj_id(object: &PbGrantObject) -> ObjectId {
1642    match object {
1643        PbGrantObject::DatabaseId(id) => (*id).into(),
1644        PbGrantObject::SchemaId(id) => (*id).into(),
1645        PbGrantObject::TableId(id) => (*id).into(),
1646        PbGrantObject::SourceId(id) => (*id).into(),
1647        PbGrantObject::SinkId(id) => (*id).into(),
1648        PbGrantObject::ViewId(id) => (*id).into(),
1649        PbGrantObject::FunctionId(id) => (*id).into(),
1650        PbGrantObject::SubscriptionId(id) => (*id).into(),
1651        PbGrantObject::ConnectionId(id) => (*id).into(),
1652        PbGrantObject::SecretId(id) => (*id).into(),
1653    }
1654}
1655
1656pub async fn insert_fragment_relations(
1657    db: &impl ConnectionTrait,
1658    downstream_fragment_relations: &FragmentDownstreamRelation,
1659) -> MetaResult<()> {
1660    let mut relations = vec![];
1661    for (upstream_fragment_id, downstreams) in downstream_fragment_relations {
1662        for downstream in downstreams {
1663            relations.push(
1664                fragment_relation::Model {
1665                    source_fragment_id: *upstream_fragment_id as _,
1666                    target_fragment_id: downstream.downstream_fragment_id as _,
1667                    dispatcher_type: downstream.dispatcher_type,
1668                    dist_key_indices: downstream
1669                        .dist_key_indices
1670                        .iter()
1671                        .map(|idx| *idx as i32)
1672                        .collect_vec()
1673                        .into(),
1674                    output_indices: downstream
1675                        .output_mapping
1676                        .indices
1677                        .iter()
1678                        .map(|idx| *idx as i32)
1679                        .collect_vec()
1680                        .into(),
1681                    output_type_mapping: Some(downstream.output_mapping.types.clone().into()),
1682                }
1683                .into_active_model(),
1684            );
1685        }
1686    }
1687    if !relations.is_empty() {
1688        FragmentRelation::insert_many(relations).exec(db).await?;
1689    }
1690    Ok(())
1691}
1692
1693pub fn compose_dispatchers(
1694    source_fragment_distribution: DistributionType,
1695    source_fragment_actors: &HashMap<crate::model::ActorId, Option<Bitmap>>,
1696    target_fragment_id: crate::model::FragmentId,
1697    target_fragment_distribution: DistributionType,
1698    target_fragment_actors: &HashMap<crate::model::ActorId, Option<Bitmap>>,
1699    dispatcher_type: DispatcherType,
1700    dist_key_indices: Vec<u32>,
1701    output_mapping: PbDispatchOutputMapping,
1702) -> (
1703    HashMap<crate::model::ActorId, PbDispatcher>,
1704    Option<HashMap<crate::model::ActorId, crate::model::ActorId>>,
1705) {
1706    match dispatcher_type {
1707        DispatcherType::Hash => {
1708            let dispatcher = PbDispatcher {
1709                r#type: PbDispatcherType::from(dispatcher_type) as _,
1710                dist_key_indices,
1711                output_mapping: output_mapping.into(),
1712                hash_mapping: Some(
1713                    ActorMapping::from_bitmaps(
1714                        &target_fragment_actors
1715                            .iter()
1716                            .map(|(actor_id, bitmap)| {
1717                                (
1718                                    *actor_id as _,
1719                                    bitmap
1720                                        .clone()
1721                                        .expect("downstream hash dispatch must have distribution"),
1722                                )
1723                            })
1724                            .collect(),
1725                    )
1726                    .to_protobuf(),
1727                ),
1728                dispatcher_id: target_fragment_id,
1729                downstream_actor_id: target_fragment_actors.keys().copied().collect(),
1730            };
1731            (
1732                source_fragment_actors
1733                    .keys()
1734                    .map(|source_actor_id| (*source_actor_id, dispatcher.clone()))
1735                    .collect(),
1736                None,
1737            )
1738        }
1739        DispatcherType::Broadcast | DispatcherType::Simple => {
1740            let dispatcher = PbDispatcher {
1741                r#type: PbDispatcherType::from(dispatcher_type) as _,
1742                dist_key_indices,
1743                output_mapping: output_mapping.into(),
1744                hash_mapping: None,
1745                dispatcher_id: target_fragment_id,
1746                downstream_actor_id: target_fragment_actors.keys().copied().collect(),
1747            };
1748            (
1749                source_fragment_actors
1750                    .keys()
1751                    .map(|source_actor_id| (*source_actor_id, dispatcher.clone()))
1752                    .collect(),
1753                None,
1754            )
1755        }
1756        DispatcherType::NoShuffle => {
1757            let no_shuffle_map = resolve_no_shuffle_actor_mapping(
1758                source_fragment_distribution,
1759                source_fragment_actors
1760                    .iter()
1761                    .map(|(&id, bitmap)| (id, bitmap)),
1762                target_fragment_distribution,
1763                target_fragment_actors
1764                    .iter()
1765                    .map(|(&id, bitmap)| (id, bitmap)),
1766            );
1767            let dispatchers = no_shuffle_map
1768                .iter()
1769                .map(|(&upstream_actor_id, &downstream_actor_id)| {
1770                    (
1771                        upstream_actor_id,
1772                        PbDispatcher {
1773                            r#type: PbDispatcherType::NoShuffle as _,
1774                            dist_key_indices: dist_key_indices.clone(),
1775                            output_mapping: output_mapping.clone().into(),
1776                            hash_mapping: None,
1777                            dispatcher_id: target_fragment_id,
1778                            downstream_actor_id: vec![downstream_actor_id],
1779                        },
1780                    )
1781                })
1782                .collect();
1783            (dispatchers, Some(no_shuffle_map))
1784        }
1785    }
1786}
1787
1788/// Resolve 1:1 actor mapping between two no-shuffle-connected fragments.
1789///
1790/// Returns `Vec<(upstream_actor_id, downstream_actor_id)>` pairs. The actor id
1791/// type is generic so the same logic can be reused both for real `ActorId` mapping
1792/// and for alignment checks with synthetic ids (e.g. `(WorkerId, actor_idx)`).
1793///
1794/// For `Single` distribution the single actors are paired directly.
1795/// For `Hash` distribution actors are matched by their bitmap's first vnode and
1796/// full bitmap equality is asserted.
1797pub fn resolve_no_shuffle_actor_mapping<
1798    'a,
1799    ActorId: Copy + Eq + std::hash::Hash + std::fmt::Debug,
1800>(
1801    source_fragment_distribution: DistributionType,
1802    source_fragment_actors: impl IntoIterator<Item = (ActorId, &'a Option<Bitmap>)>,
1803    target_fragment_distribution: DistributionType,
1804    target_fragment_actors: impl IntoIterator<Item = (ActorId, &'a Option<Bitmap>)>,
1805) -> HashMap<ActorId, ActorId> {
1806    assert_eq!(source_fragment_distribution, target_fragment_distribution);
1807
1808    match source_fragment_distribution {
1809        DistributionType::Single => {
1810            let assert_singleton = |bitmap: &Option<Bitmap>| {
1811                assert!(
1812                    bitmap.as_ref().map(|bitmap| bitmap.all()).unwrap_or(true),
1813                    "not singleton: {:?}",
1814                    bitmap
1815                );
1816            };
1817            let (source_actor_id, bitmap) =
1818                Itertools::exactly_one(source_fragment_actors.into_iter())
1819                    .ok()
1820                    .expect("Single distribution should have exactly one source actor");
1821            assert_singleton(bitmap);
1822            let (target_actor_id, bitmap) =
1823                Itertools::exactly_one(target_fragment_actors.into_iter())
1824                    .ok()
1825                    .expect("Single distribution should have exactly one target actor");
1826            assert_singleton(bitmap);
1827            HashMap::from([(source_actor_id, target_actor_id)])
1828        }
1829        DistributionType::Hash => {
1830            // Build target index keyed by first vnode (one pass over target).
1831            let mut target_by_vnode: HashMap<_, _> = target_fragment_actors
1832                .into_iter()
1833                .map(|(actor_id, bitmap)| {
1834                    let bitmap = bitmap
1835                        .as_ref()
1836                        .expect("hash distribution should have bitmap");
1837                    let first_vnode = bitmap.iter_vnodes().next().expect("non-empty bitmap");
1838                    (first_vnode, (actor_id, bitmap))
1839                })
1840                .collect();
1841
1842            let target_count = target_by_vnode.len();
1843
1844            // One pass over source, matching against the target index.
1845            let mapping: HashMap<_, _> = source_fragment_actors
1846                .into_iter()
1847                .map(|(source_actor_id, source_bitmap)| {
1848                    let source_bitmap = source_bitmap
1849                        .as_ref()
1850                        .expect("hash distribution should have bitmap");
1851                    let first_vnode = source_bitmap
1852                        .iter_vnodes()
1853                        .next()
1854                        .expect("non-empty bitmap");
1855                    let (target_actor_id, target_bitmap) =
1856                        target_by_vnode.remove(&first_vnode).unwrap_or_else(|| {
1857                            panic!(
1858                                "cannot find matched target actor: {:?} first_vnode {:?}",
1859                                source_actor_id, first_vnode,
1860                            )
1861                        });
1862                    assert_eq!(
1863                        source_bitmap, target_bitmap,
1864                        "bitmap mismatch for source {:?} target {:?} at first_vnode {:?}",
1865                        source_actor_id, target_actor_id, first_vnode,
1866                    );
1867                    (source_actor_id, target_actor_id)
1868                })
1869                .collect();
1870
1871            assert_eq!(
1872                mapping.len(),
1873                target_count,
1874                "no-shuffle should have equal upstream downstream actor count: {} vs {}",
1875                mapping.len(),
1876                target_count,
1877            );
1878
1879            mapping
1880        }
1881    }
1882}
1883
1884pub fn rebuild_fragment_mapping(fragment: &SharedFragmentInfo) -> PbFragmentWorkerSlotMapping {
1885    let fragment_worker_slot_mapping = match fragment.distribution_type {
1886        DistributionType::Single => {
1887            let actor = Itertools::exactly_one(fragment.actors.values()).unwrap();
1888            WorkerSlotMapping::new_single(WorkerSlotId::new(actor.worker_id as _, 0))
1889        }
1890        DistributionType::Hash => {
1891            let actor_bitmaps: HashMap<_, _> = fragment
1892                .actors
1893                .iter()
1894                .map(|(actor_id, actor_info)| {
1895                    let vnode_bitmap = actor_info
1896                        .vnode_bitmap
1897                        .as_ref()
1898                        .cloned()
1899                        .expect("actor bitmap shouldn't be none in hash fragment");
1900
1901                    (*actor_id as hash::ActorId, vnode_bitmap)
1902                })
1903                .collect();
1904
1905            let actor_mapping = ActorMapping::from_bitmaps(&actor_bitmaps);
1906
1907            let actor_locations = fragment
1908                .actors
1909                .iter()
1910                .map(|(actor_id, actor_info)| (*actor_id as hash::ActorId, actor_info.worker_id))
1911                .collect();
1912
1913            actor_mapping.to_worker_slot(&actor_locations)
1914        }
1915    };
1916
1917    PbFragmentWorkerSlotMapping {
1918        fragment_id: fragment.fragment_id,
1919        mapping: Some(fragment_worker_slot_mapping.to_protobuf()),
1920    }
1921}
1922
1923/// For the given streaming jobs, returns
1924/// - All source fragments
1925/// - All sink fragments
1926/// - All fragments
1927pub async fn get_fragments_for_jobs<C>(
1928    db: &C,
1929    streaming_jobs: Vec<JobId>,
1930) -> MetaResult<(
1931    HashMap<SourceId, BTreeSet<FragmentId>>,
1932    HashSet<FragmentId>,
1933    HashSet<FragmentId>,
1934)>
1935where
1936    C: ConnectionTrait,
1937{
1938    if streaming_jobs.is_empty() {
1939        return Ok((HashMap::default(), HashSet::default(), HashSet::default()));
1940    }
1941
1942    let fragments: Vec<(FragmentId, i32, StreamNode)> = Fragment::find()
1943        .select_only()
1944        .columns([
1945            fragment::Column::FragmentId,
1946            fragment::Column::FragmentTypeMask,
1947            fragment::Column::StreamNode,
1948        ])
1949        .filter(fragment::Column::JobId.is_in(streaming_jobs))
1950        .into_tuple()
1951        .all(db)
1952        .await?;
1953
1954    let fragment_ids: HashSet<_> = fragments
1955        .iter()
1956        .map(|(fragment_id, _, _)| *fragment_id)
1957        .collect();
1958
1959    let mut source_fragment_ids: HashMap<SourceId, BTreeSet<FragmentId>> = HashMap::new();
1960    let mut sink_fragment_ids: HashSet<FragmentId> = HashSet::new();
1961    for (fragment_id, mask, stream_node) in fragments {
1962        if FragmentTypeMask::from(mask).contains(FragmentTypeFlag::Source)
1963            && let Some(source_id) = stream_node.to_protobuf().find_stream_source()
1964        {
1965            source_fragment_ids
1966                .entry(source_id)
1967                .or_default()
1968                .insert(fragment_id);
1969        }
1970        if FragmentTypeMask::from(mask).contains(FragmentTypeFlag::Sink) {
1971            sink_fragment_ids.insert(fragment_id);
1972        }
1973    }
1974
1975    Ok((source_fragment_ids, sink_fragment_ids, fragment_ids))
1976}
1977
1978/// Build a object group for notifying the deletion of the given objects.
1979///
1980/// Note that only id fields are filled in the object info, as the arguments are partial objects.
1981/// As a result, the returned notification info should only be used for deletion.
1982pub(crate) fn build_object_group_for_delete(
1983    partial_objects: Vec<PartialObject>,
1984) -> NotificationInfo {
1985    let mut objects = vec![];
1986    for obj in partial_objects {
1987        match obj.obj_type {
1988            ObjectType::Database => objects.push(PbObject {
1989                object_info: Some(PbObjectInfo::Database(PbDatabase {
1990                    id: obj.oid.as_database_id(),
1991                    ..Default::default()
1992                })),
1993            }),
1994            ObjectType::Schema => objects.push(PbObject {
1995                object_info: Some(PbObjectInfo::Schema(PbSchema {
1996                    id: obj.oid.as_schema_id(),
1997                    database_id: obj.database_id.unwrap(),
1998                    ..Default::default()
1999                })),
2000            }),
2001            ObjectType::Table => objects.push(PbObject {
2002                object_info: Some(PbObjectInfo::Table(PbTable {
2003                    id: obj.oid.as_table_id(),
2004                    schema_id: obj.schema_id.unwrap(),
2005                    database_id: obj.database_id.unwrap(),
2006                    ..Default::default()
2007                })),
2008            }),
2009            ObjectType::Source => objects.push(PbObject {
2010                object_info: Some(PbObjectInfo::Source(PbSource {
2011                    id: obj.oid.as_source_id(),
2012                    schema_id: obj.schema_id.unwrap(),
2013                    database_id: obj.database_id.unwrap(),
2014                    ..Default::default()
2015                })),
2016            }),
2017            ObjectType::Sink => objects.push(PbObject {
2018                object_info: Some(PbObjectInfo::Sink(PbSink {
2019                    id: obj.oid.as_sink_id(),
2020                    schema_id: obj.schema_id.unwrap(),
2021                    database_id: obj.database_id.unwrap(),
2022                    ..Default::default()
2023                })),
2024            }),
2025            ObjectType::Subscription => objects.push(PbObject {
2026                object_info: Some(PbObjectInfo::Subscription(PbSubscription {
2027                    id: obj.oid.as_subscription_id(),
2028                    schema_id: obj.schema_id.unwrap(),
2029                    database_id: obj.database_id.unwrap(),
2030                    ..Default::default()
2031                })),
2032            }),
2033            ObjectType::View => objects.push(PbObject {
2034                object_info: Some(PbObjectInfo::View(PbView {
2035                    id: obj.oid.as_view_id(),
2036                    schema_id: obj.schema_id.unwrap(),
2037                    database_id: obj.database_id.unwrap(),
2038                    ..Default::default()
2039                })),
2040            }),
2041            ObjectType::Index => {
2042                objects.push(PbObject {
2043                    object_info: Some(PbObjectInfo::Index(PbIndex {
2044                        id: obj.oid.as_index_id(),
2045                        schema_id: obj.schema_id.unwrap(),
2046                        database_id: obj.database_id.unwrap(),
2047                        ..Default::default()
2048                    })),
2049                });
2050                objects.push(PbObject {
2051                    object_info: Some(PbObjectInfo::Table(PbTable {
2052                        id: obj.oid.as_table_id(),
2053                        schema_id: obj.schema_id.unwrap(),
2054                        database_id: obj.database_id.unwrap(),
2055                        ..Default::default()
2056                    })),
2057                });
2058            }
2059            ObjectType::Function => objects.push(PbObject {
2060                object_info: Some(PbObjectInfo::Function(PbFunction {
2061                    id: obj.oid.as_function_id(),
2062                    schema_id: obj.schema_id.unwrap(),
2063                    database_id: obj.database_id.unwrap(),
2064                    ..Default::default()
2065                })),
2066            }),
2067            ObjectType::Connection => objects.push(PbObject {
2068                object_info: Some(PbObjectInfo::Connection(PbConnection {
2069                    id: obj.oid.as_connection_id(),
2070                    schema_id: obj.schema_id.unwrap(),
2071                    database_id: obj.database_id.unwrap(),
2072                    ..Default::default()
2073                })),
2074            }),
2075            ObjectType::Secret => objects.push(PbObject {
2076                object_info: Some(PbObjectInfo::Secret(PbSecret {
2077                    id: obj.oid.as_secret_id(),
2078                    schema_id: obj.schema_id.unwrap(),
2079                    database_id: obj.database_id.unwrap(),
2080                    ..Default::default()
2081                })),
2082            }),
2083        }
2084    }
2085    NotificationInfo::ObjectGroup(PbObjectGroup {
2086        objects,
2087        dependencies: vec![],
2088    })
2089}
2090
2091pub fn extract_external_table_name_from_definition(table_definition: &str) -> Option<String> {
2092    let [mut definition]: [_; 1] = Parser::parse_sql(table_definition)
2093        .context("unable to parse table definition")
2094        .inspect_err(|e| {
2095            tracing::error!(
2096                target: "auto_schema_change",
2097                error = %e.as_report(),
2098                "failed to parse table definition")
2099        })
2100        .unwrap()
2101        .try_into()
2102        .unwrap();
2103    if let SqlStatement::CreateTable { cdc_table_info, .. } = &mut definition {
2104        cdc_table_info
2105            .clone()
2106            .map(|cdc_table_info| cdc_table_info.external_table_name)
2107    } else {
2108        None
2109    }
2110}
2111
2112/// `rename_relation` renames the target relation and its definition,
2113/// it commits the changes to the transaction and returns the updated relations and the old name.
2114pub async fn rename_relation(
2115    txn: &DatabaseTransaction,
2116    object_type: ObjectType,
2117    object_id: ObjectId,
2118    object_name: &str,
2119) -> MetaResult<(Vec<PbObject>, String)> {
2120    use sea_orm::ActiveModelTrait;
2121
2122    use crate::controller::rename::alter_relation_rename;
2123
2124    let mut to_update_relations = vec![];
2125    // rename relation.
2126    macro_rules! rename_relation {
2127        ($entity:ident, $table:ident, $identity:ident, $object_id:expr) => {{
2128            let (mut relation, obj) = $entity::find_by_id($object_id)
2129                .find_also_related(Object)
2130                .one(txn)
2131                .await?
2132                .unwrap();
2133            let obj = obj.unwrap();
2134            let old_name = relation.name.clone();
2135            relation.name = object_name.into();
2136            if obj.obj_type != ObjectType::View {
2137                relation.definition = alter_relation_rename(&relation.definition, object_name);
2138            }
2139            let active_model = $table::ActiveModel {
2140                $identity: Set(relation.$identity),
2141                name: Set(object_name.into()),
2142                definition: Set(relation.definition.clone()),
2143                ..Default::default()
2144            };
2145            active_model.update(txn).await?;
2146            let streaming_job = streaming_job::Entity::find_by_id($object_id.as_raw_id())
2147                .one(txn)
2148                .await?;
2149            to_update_relations.push(PbObject {
2150                object_info: Some(PbObjectInfo::$entity(
2151                    ObjectModel(relation, obj, streaming_job).into(),
2152                )),
2153            });
2154            old_name
2155        }};
2156    }
2157    // TODO: check is there any thing to change for shared source?
2158    let old_name = match object_type {
2159        ObjectType::Table => {
2160            let associated_source_id: Option<SourceId> = Source::find()
2161                .select_only()
2162                .column(source::Column::SourceId)
2163                .filter(source::Column::OptionalAssociatedTableId.eq(object_id))
2164                .into_tuple()
2165                .one(txn)
2166                .await?;
2167            if let Some(source_id) = associated_source_id {
2168                rename_relation!(Source, source, source_id, source_id);
2169            }
2170            rename_relation!(Table, table, table_id, object_id.as_table_id())
2171        }
2172        ObjectType::Source => {
2173            rename_relation!(Source, source, source_id, object_id.as_source_id())
2174        }
2175        ObjectType::Sink => rename_relation!(Sink, sink, sink_id, object_id.as_sink_id()),
2176        ObjectType::Subscription => {
2177            rename_relation!(
2178                Subscription,
2179                subscription,
2180                subscription_id,
2181                object_id.as_subscription_id()
2182            )
2183        }
2184        ObjectType::View => rename_relation!(View, view, view_id, object_id.as_view_id()),
2185        ObjectType::Index => {
2186            let (mut index, obj) = Index::find_by_id(object_id.as_index_id())
2187                .find_also_related(Object)
2188                .one(txn)
2189                .await?
2190                .unwrap();
2191            let streaming_job = streaming_job::Entity::find_by_id(index.index_id.as_job_id())
2192                .one(txn)
2193                .await?;
2194            index.name = object_name.into();
2195            let index_table_id = index.index_table_id;
2196            let old_name = rename_relation!(Table, table, table_id, index_table_id);
2197
2198            // the name of index and its associated table is the same.
2199            let active_model = index::ActiveModel {
2200                index_id: sea_orm::ActiveValue::Set(index.index_id),
2201                name: sea_orm::ActiveValue::Set(object_name.into()),
2202                ..Default::default()
2203            };
2204            active_model.update(txn).await?;
2205            to_update_relations.push(PbObject {
2206                object_info: Some(PbObjectInfo::Index(
2207                    ObjectModel(index, obj.unwrap(), streaming_job).into(),
2208                )),
2209            });
2210            old_name
2211        }
2212        _ => unreachable!("only relation name can be altered."),
2213    };
2214
2215    Ok((to_update_relations, old_name))
2216}
2217
2218pub async fn get_database_resource_group<C>(txn: &C, database_id: DatabaseId) -> MetaResult<String>
2219where
2220    C: ConnectionTrait,
2221{
2222    let database_resource_group: Option<String> = Database::find_by_id(database_id)
2223        .select_only()
2224        .column(database::Column::ResourceGroup)
2225        .into_tuple()
2226        .one(txn)
2227        .await?
2228        .ok_or_else(|| MetaError::catalog_id_not_found("database", database_id))?;
2229
2230    Ok(database_resource_group.unwrap_or_else(|| DEFAULT_RESOURCE_GROUP.to_owned()))
2231}
2232
2233pub async fn get_existing_job_resource_group<C>(
2234    txn: &C,
2235    streaming_job_id: JobId,
2236) -> MetaResult<String>
2237where
2238    C: ConnectionTrait,
2239{
2240    let (job_specific_resource_group, database_resource_group): (Option<String>, Option<String>) =
2241        StreamingJob::find_by_id(streaming_job_id)
2242            .select_only()
2243            .join(JoinType::InnerJoin, streaming_job::Relation::Object.def())
2244            .join(JoinType::InnerJoin, object::Relation::Database2.def())
2245            .column(streaming_job::Column::SpecificResourceGroup)
2246            .column(database::Column::ResourceGroup)
2247            .into_tuple()
2248            .one(txn)
2249            .await?
2250            .ok_or_else(|| MetaError::catalog_id_not_found("streaming job", streaming_job_id))?;
2251
2252    Ok(job_specific_resource_group.unwrap_or_else(|| {
2253        database_resource_group.unwrap_or_else(|| DEFAULT_RESOURCE_GROUP.to_owned())
2254    }))
2255}
2256
2257pub fn filter_workers_by_resource_group(
2258    workers: &HashMap<WorkerId, WorkerNode>,
2259    resource_group: &str,
2260) -> BTreeSet<WorkerId> {
2261    workers
2262        .iter()
2263        .filter(|&(_, worker)| {
2264            worker
2265                .resource_group()
2266                .map(|node_label| node_label.as_str() == resource_group)
2267                .unwrap_or(false)
2268        })
2269        .map(|(id, _)| *id)
2270        .collect()
2271}
2272
2273/// `rename_relation_refer` updates the definition of relations that refer to the target one,
2274/// it commits the changes to the transaction and returns all the updated relations.
2275pub async fn rename_relation_refer(
2276    txn: &DatabaseTransaction,
2277    object_type: ObjectType,
2278    object_id: ObjectId,
2279    object_name: &str,
2280    old_name: &str,
2281) -> MetaResult<Vec<PbObject>> {
2282    use sea_orm::ActiveModelTrait;
2283
2284    use crate::controller::rename::alter_relation_rename_refs;
2285
2286    let mut to_update_relations = vec![];
2287    macro_rules! rename_relation_ref {
2288        ($entity:ident, $table:ident, $identity:ident, $object_id:expr) => {{
2289            let (mut relation, obj) = $entity::find_by_id($object_id)
2290                .find_also_related(Object)
2291                .one(txn)
2292                .await?
2293                .unwrap();
2294            relation.definition =
2295                alter_relation_rename_refs(&relation.definition, old_name, object_name);
2296            let active_model = $table::ActiveModel {
2297                $identity: Set(relation.$identity),
2298                definition: Set(relation.definition.clone()),
2299                ..Default::default()
2300            };
2301            active_model.update(txn).await?;
2302            let streaming_job = streaming_job::Entity::find_by_id($object_id.as_raw_id())
2303                .one(txn)
2304                .await?;
2305            to_update_relations.push(PbObject {
2306                object_info: Some(PbObjectInfo::$entity(
2307                    ObjectModel(relation, obj.unwrap(), streaming_job).into(),
2308                )),
2309            });
2310        }};
2311    }
2312    let objs = get_referring_objects(object_id, object_type, txn).await?;
2313
2314    for obj in objs {
2315        match obj.obj_type {
2316            ObjectType::Table => {
2317                rename_relation_ref!(Table, table, table_id, obj.oid.as_table_id())
2318            }
2319            ObjectType::Sink => {
2320                rename_relation_ref!(Sink, sink, sink_id, obj.oid.as_sink_id())
2321            }
2322            // A CDC table source refers to the shared CDC source it reads from.
2323            ObjectType::Source => {
2324                rename_relation_ref!(Source, source, source_id, obj.oid.as_source_id())
2325            }
2326            ObjectType::Subscription => {
2327                rename_relation_ref!(
2328                    Subscription,
2329                    subscription,
2330                    subscription_id,
2331                    obj.oid.as_subscription_id()
2332                )
2333            }
2334            ObjectType::View => rename_relation_ref!(View, view, view_id, obj.oid.as_view_id()),
2335            ObjectType::Index => {
2336                let index_table_id: Option<TableId> = Index::find_by_id(obj.oid.as_index_id())
2337                    .select_only()
2338                    .column(index::Column::IndexTableId)
2339                    .into_tuple()
2340                    .one(txn)
2341                    .await?;
2342                rename_relation_ref!(Table, table, table_id, index_table_id.unwrap());
2343            }
2344            _ => {
2345                bail!(
2346                    "only the table, source, sink, subscription, view and index will depend on other objects."
2347                )
2348            }
2349        }
2350    }
2351
2352    Ok(to_update_relations)
2353}
2354
2355/// Validate that subscription can be safely deleted, meeting any of the following conditions:
2356/// 1. The upstream table is not referred to by any cross-db mv.
2357/// 2. After deleting the subscription, the upstream table still has at least one subscription.
2358pub async fn validate_subscription_deletion<C>(
2359    txn: &C,
2360    subscription_id: SubscriptionId,
2361) -> MetaResult<()>
2362where
2363    C: ConnectionTrait,
2364{
2365    let upstream_table_id: ObjectId = Subscription::find_by_id(subscription_id)
2366        .select_only()
2367        .column(subscription::Column::DependentTableId)
2368        .into_tuple()
2369        .one(txn)
2370        .await?
2371        .ok_or_else(|| MetaError::catalog_id_not_found("subscription", subscription_id))?;
2372
2373    let cnt = Subscription::find()
2374        .filter(subscription::Column::DependentTableId.eq(upstream_table_id))
2375        .count(txn)
2376        .await?;
2377    if cnt > 1 {
2378        // Ensure that at least one subscription is remained for the upstream table
2379        // once the subscription is dropped.
2380        return Ok(());
2381    }
2382
2383    // Ensure that the upstream table is not referred by any cross-db mv.
2384    let obj_alias = Alias::new("o1");
2385    let used_by_alias = Alias::new("o2");
2386    let count = ObjectDependency::find()
2387        .join_as(
2388            JoinType::InnerJoin,
2389            object_dependency::Relation::Object2.def(),
2390            obj_alias.clone(),
2391        )
2392        .join_as(
2393            JoinType::InnerJoin,
2394            object_dependency::Relation::Object1.def(),
2395            used_by_alias.clone(),
2396        )
2397        .filter(
2398            object_dependency::Column::Oid
2399                .eq(upstream_table_id)
2400                .and(object_dependency::Column::UsedBy.ne(subscription_id))
2401                .and(
2402                    Expr::col((obj_alias, object::Column::DatabaseId))
2403                        .ne(Expr::col((used_by_alias, object::Column::DatabaseId))),
2404                ),
2405        )
2406        .count(txn)
2407        .await?;
2408
2409    if count != 0 {
2410        return Err(MetaError::permission_denied(format!(
2411            "Referenced by {} cross-db objects.",
2412            count
2413        )));
2414    }
2415
2416    Ok(())
2417}
2418
2419pub async fn fetch_target_fragments<C>(
2420    txn: &C,
2421    src_fragment_id: impl IntoIterator<Item = FragmentId>,
2422) -> MetaResult<HashMap<FragmentId, Vec<FragmentId>>>
2423where
2424    C: ConnectionTrait,
2425{
2426    let source_target_fragments: Vec<(FragmentId, FragmentId)> = FragmentRelation::find()
2427        .select_only()
2428        .columns([
2429            fragment_relation::Column::SourceFragmentId,
2430            fragment_relation::Column::TargetFragmentId,
2431        ])
2432        .filter(fragment_relation::Column::SourceFragmentId.is_in(src_fragment_id))
2433        .into_tuple()
2434        .all(txn)
2435        .await?;
2436
2437    let source_target_fragments = source_target_fragments.into_iter().into_group_map();
2438
2439    Ok(source_target_fragments)
2440}
2441
2442pub async fn get_sink_fragment_by_ids<C>(
2443    txn: &C,
2444    sink_ids: Vec<SinkId>,
2445) -> MetaResult<HashMap<SinkId, FragmentId>>
2446where
2447    C: ConnectionTrait,
2448{
2449    let sink_num = sink_ids.len();
2450    let sink_fragment_ids: Vec<(SinkId, FragmentId)> = Fragment::find()
2451        .select_only()
2452        .columns([fragment::Column::JobId, fragment::Column::FragmentId])
2453        .filter(
2454            fragment::Column::JobId
2455                .is_in(sink_ids)
2456                .and(FragmentTypeMask::intersects(FragmentTypeFlag::Sink)),
2457        )
2458        .into_tuple()
2459        .all(txn)
2460        .await?;
2461
2462    if sink_fragment_ids.len() != sink_num {
2463        return Err(anyhow::anyhow!(
2464            "expected exactly one sink fragment for each sink, but got {} fragments for {} sinks",
2465            sink_fragment_ids.len(),
2466            sink_num
2467        )
2468        .into());
2469    }
2470
2471    Ok(sink_fragment_ids.into_iter().collect())
2472}
2473
2474pub async fn has_table_been_migrated<C>(txn: &C, table_id: TableId) -> MetaResult<bool>
2475where
2476    C: ConnectionTrait,
2477{
2478    let mview_fragment: Vec<i32> = Fragment::find()
2479        .select_only()
2480        .column(fragment::Column::FragmentTypeMask)
2481        .filter(
2482            fragment::Column::JobId
2483                .eq(table_id)
2484                .and(FragmentTypeMask::intersects(FragmentTypeFlag::Mview)),
2485        )
2486        .into_tuple()
2487        .all(txn)
2488        .await?;
2489
2490    let mview_fragment_len = mview_fragment.len();
2491    if mview_fragment_len != 1 {
2492        bail!(
2493            "expected exactly one mview fragment for table {}, found {}",
2494            table_id,
2495            mview_fragment_len
2496        );
2497    }
2498
2499    let mview_fragment = mview_fragment.into_iter().next().unwrap();
2500    let migrated =
2501        FragmentTypeMask::from(mview_fragment).contains(FragmentTypeFlag::UpstreamSinkUnion);
2502
2503    Ok(migrated)
2504}
2505
2506pub async fn check_if_belongs_to_iceberg_table<C>(txn: &C, job_id: JobId) -> MetaResult<bool>
2507where
2508    C: ConnectionTrait,
2509{
2510    if let Some(engine) = Table::find_by_id(job_id.as_mv_table_id())
2511        .select_only()
2512        .column(table::Column::Engine)
2513        .into_tuple::<table::Engine>()
2514        .one(txn)
2515        .await?
2516        && engine == table::Engine::Iceberg
2517    {
2518        return Ok(true);
2519    }
2520    if let Some(parent_oid) = Object::find_by_id(job_id)
2521        .select_only()
2522        .column(object::Column::BelongToOid)
2523        .into_tuple::<Option<ObjectId>>()
2524        .one(txn)
2525        .await?
2526        .flatten()
2527        && Table::find_by_id(parent_oid.as_table_id())
2528            .filter(table::Column::Engine.eq(table::Engine::Iceberg))
2529            .one(txn)
2530            .await?
2531            .is_some()
2532    {
2533        return Ok(true);
2534    }
2535    Ok(false)
2536}
2537
2538pub async fn find_dirty_iceberg_table_jobs<C>(
2539    txn: &C,
2540    database_id: Option<DatabaseId>,
2541) -> MetaResult<Vec<PartialObject>>
2542where
2543    C: ConnectionTrait,
2544{
2545    let mut filter_condition = streaming_job::Column::JobStatus
2546        .ne(JobStatus::Created)
2547        .and(object::Column::ObjType.is_in([ObjectType::Table, ObjectType::Sink]))
2548        .and(streaming_job::Column::CreateType.eq(CreateType::Background));
2549    if let Some(database_id) = database_id {
2550        filter_condition = filter_condition.and(object::Column::DatabaseId.eq(database_id));
2551    }
2552    let creating_table_sink_jobs: Vec<PartialObject> = StreamingJob::find()
2553        .select_only()
2554        .columns([
2555            object::Column::Oid,
2556            object::Column::ObjType,
2557            object::Column::SchemaId,
2558            object::Column::DatabaseId,
2559        ])
2560        .join(JoinType::InnerJoin, streaming_job::Relation::Object.def())
2561        .filter(filter_condition)
2562        .into_partial_model()
2563        .all(txn)
2564        .await?;
2565
2566    let mut dirty_iceberg_table_jobs = vec![];
2567    for job in creating_table_sink_jobs {
2568        if check_if_belongs_to_iceberg_table(txn, job.oid.as_job_id()).await? {
2569            tracing::info!("Found dirty iceberg job with id: {}", job.oid);
2570            dirty_iceberg_table_jobs.push(job);
2571        }
2572    }
2573
2574    Ok(dirty_iceberg_table_jobs)
2575}
2576
2577pub fn build_select_node_list(
2578    from: &[ColumnCatalog],
2579    to: &[ColumnCatalog],
2580) -> MetaResult<Vec<PbExprNode>> {
2581    let mut exprs = Vec::with_capacity(to.len());
2582    let idx_by_col_id = from
2583        .iter()
2584        .enumerate()
2585        .map(|(idx, col)| (col.column_desc.as_ref().unwrap().column_id, idx))
2586        .collect::<HashMap<_, _>>();
2587
2588    for to_col in to {
2589        let to_col = to_col.column_desc.as_ref().unwrap();
2590        let to_col_type_ref = to_col.column_type.as_ref().unwrap();
2591        let to_col_type = DataType::from(to_col_type_ref);
2592        if let Some(from_idx) = idx_by_col_id.get(&to_col.column_id) {
2593            let from_col_type = DataType::from(
2594                from[*from_idx]
2595                    .column_desc
2596                    .as_ref()
2597                    .unwrap()
2598                    .column_type
2599                    .as_ref()
2600                    .unwrap(),
2601            );
2602            if !to_col_type.equals_datatype(&from_col_type) {
2603                return Err(anyhow!(
2604                    "Column type mismatch: {:?} != {:?}",
2605                    from_col_type,
2606                    to_col_type
2607                )
2608                .into());
2609            }
2610            exprs.push(PbExprNode {
2611                function_type: expr_node::Type::Unspecified.into(),
2612                return_type: Some(to_col_type_ref.clone()),
2613                rex_node: Some(expr_node::RexNode::InputRef(*from_idx as _)),
2614            });
2615        } else {
2616            let to_default_node =
2617                if let Some(GeneratedOrDefaultColumn::DefaultColumn(DefaultColumnDesc {
2618                    expr,
2619                    ..
2620                })) = &to_col.generated_or_default_column
2621                {
2622                    expr.clone().unwrap()
2623                } else {
2624                    let null = Datum::None.to_protobuf();
2625                    PbExprNode {
2626                        function_type: expr_node::Type::Unspecified.into(),
2627                        return_type: Some(to_col_type_ref.clone()),
2628                        rex_node: Some(expr_node::RexNode::Constant(null)),
2629                    }
2630                };
2631            exprs.push(to_default_node);
2632        }
2633    }
2634
2635    Ok(exprs)
2636}
2637
2638#[derive(Clone, Debug, Default)]
2639pub struct StreamingJobExtraInfo {
2640    pub timezone: Option<String>,
2641    pub config_override: Arc<str>,
2642    pub job_definition: String,
2643    pub backfill_orders: Option<BackfillOrders>,
2644    pub refresh_interval_sec: Option<u64>,
2645}
2646
2647impl StreamingJobExtraInfo {
2648    pub fn stream_context(&self) -> StreamContext {
2649        StreamContext {
2650            timezone: self.timezone.clone(),
2651            config_override: self.config_override.clone(),
2652        }
2653    }
2654}
2655
2656/// Tuple of (`job_id`, `timezone`, `config_override`, `backfill_orders`, `refresh_interval_sec`)
2657type StreamingJobExtraInfoRow = (
2658    JobId,
2659    Option<String>,
2660    Option<String>,
2661    Option<BackfillOrders>,
2662    Option<i64>,
2663);
2664
2665pub async fn get_streaming_job_extra_info<C>(
2666    txn: &C,
2667    job_ids: Vec<JobId>,
2668) -> MetaResult<HashMap<JobId, StreamingJobExtraInfo>>
2669where
2670    C: ConnectionTrait,
2671{
2672    let pairs: Vec<StreamingJobExtraInfoRow> = StreamingJob::find()
2673        .select_only()
2674        .columns([
2675            streaming_job::Column::JobId,
2676            streaming_job::Column::Timezone,
2677            streaming_job::Column::ConfigOverride,
2678            streaming_job::Column::BackfillOrders,
2679            streaming_job::Column::RefreshIntervalSec,
2680        ])
2681        .filter(streaming_job::Column::JobId.is_in(job_ids.clone()))
2682        .into_tuple()
2683        .all(txn)
2684        .await?;
2685
2686    let job_ids = job_ids.into_iter().collect();
2687
2688    let mut definitions = resolve_streaming_job_definition(txn, &job_ids).await?;
2689
2690    let result = pairs
2691        .into_iter()
2692        .map(
2693            |(job_id, timezone, config_override, backfill_orders, refresh_interval_sec)| {
2694                let job_definition = definitions.remove(&job_id).unwrap_or_default();
2695                (
2696                    job_id,
2697                    StreamingJobExtraInfo {
2698                        timezone,
2699                        config_override: config_override.unwrap_or_default().into(),
2700                        job_definition,
2701                        backfill_orders,
2702                        refresh_interval_sec: refresh_interval_sec.map(|s| s as u64),
2703                    },
2704                )
2705            },
2706        )
2707        .collect();
2708
2709    Ok(result)
2710}
2711
2712#[cfg(test)]
2713mod tests {
2714    use super::*;
2715
2716    #[test]
2717    fn test_extract_cdc_table_name() {
2718        let ddl1 = "CREATE TABLE t1 () FROM pg_source TABLE 'public.t1'";
2719        let ddl2 = "CREATE TABLE t2 (v1 int) FROM pg_source TABLE 'mydb.t2'";
2720        assert_eq!(
2721            extract_external_table_name_from_definition(ddl1),
2722            Some("public.t1".into())
2723        );
2724        assert_eq!(
2725            extract_external_table_name_from_definition(ddl2),
2726            Some("mydb.t2".into())
2727        );
2728    }
2729}