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