1use 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
81pub 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
185async 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
280pub 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
288pub 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
296pub 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
420pub 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
478pub 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
514pub 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 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
549pub 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
583pub 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
595pub 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
611pub 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
640pub 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
716pub 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
732pub 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
803pub 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
829pub 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 _ => 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
1019pub 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
1051pub 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
1067pub 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
1088pub 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
1103pub 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
1224pub 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
1308pub(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
1354pub 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
1372pub 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
1416pub(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 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
1462pub 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 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 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
1570pub 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
1718pub 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 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 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
1853pub 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
1908pub(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
2042pub 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 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 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 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
2203pub 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
2281pub 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 return Ok(());
2307 }
2308
2309 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
2582type 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}