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