Skip to main content

common_meta/rpc/
ddl.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[cfg(feature = "enterprise")]
16pub mod trigger;
17
18use std::collections::{HashMap, HashSet};
19use std::result;
20use std::time::Duration;
21
22use api::helper::{from_pb_time_ranges, to_pb_time_ranges};
23use api::v1::alter_database_expr::Kind as PbAlterDatabaseKind;
24use api::v1::meta::ddl_task_request::Task;
25use api::v1::meta::{
26    AlterDatabaseTask as PbAlterDatabaseTask, AlterTableTask as PbAlterTableTask,
27    AlterTableTasks as PbAlterTableTasks, CommentOnTask as PbCommentOnTask,
28    CreateDatabaseTask as PbCreateDatabaseTask, CreateFlowTask as PbCreateFlowTask,
29    CreateTableTask as PbCreateTableTask, CreateTableTasks as PbCreateTableTasks,
30    CreateViewTask as PbCreateViewTask, DdlTaskRequest as PbDdlTaskRequest,
31    DdlTaskResponse as PbDdlTaskResponse, DropDatabaseTask as PbDropDatabaseTask,
32    DropFlowTask as PbDropFlowTask, DropTableTask as PbDropTableTask,
33    DropTableTasks as PbDropTableTasks, DropViewTask as PbDropViewTask, Partition, ProcedureId,
34    PurgeDroppedTableTask as PbPurgeDroppedTableTask, TruncateTableTask as PbTruncateTableTask,
35    UndropTableTask as PbUndropTableTask,
36};
37use api::v1::{
38    AlterDatabaseExpr, AlterTableExpr, CommentObjectType as PbCommentObjectType, CommentOnExpr,
39    CreateDatabaseExpr, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropDatabaseExpr,
40    DropFlowExpr, DropTableExpr, DropViewExpr, EvalInterval, ExpireAfter, Option as PbOption,
41    QueryContext as PbQueryContext, TruncateTableExpr,
42};
43use base64::Engine as _;
44use base64::engine::general_purpose;
45use common_base::protocol::Channel;
46use common_catalog::{format_full_flow_name, format_full_table_name};
47use common_error::ext::BoxedError;
48pub use common_event_recorder::TriggerReason;
49use common_time::{DatabaseTimeToLive, Timestamp};
50use prost::Message;
51use serde::{Deserialize, Serialize};
52use serde_with::{DefaultOnNull, serde_as};
53use snafu::{OptionExt, ResultExt};
54use table::metadata::{TableId, TableInfo};
55use table::requests::validate_database_option;
56use table::table_name::TableName;
57use table::table_reference::TableReference;
58
59use crate::error::{
60    self, ConvertTimeRangesSnafu, ExternalSnafu, InvalidSetDatabaseOptionSnafu,
61    InvalidUnsetDatabaseOptionSnafu, Result,
62};
63use crate::flow_name::FlowName;
64use crate::instruction::CacheIdent;
65use crate::key::FlowId;
66use crate::key::flow::flow_name::FlowNameManager;
67use crate::key::table_name::{TableNameKey, TableNameManager};
68
69/// Reserved query-context extension key for the frontend peer address that submitted a DDL request.
70pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr";
71/// Reserved query-context extension key for the authenticated database creator.
72pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator";
73/// Internal gRPC metadata key for the authenticated database creator.
74pub const CREATE_DATABASE_CREATOR_METADATA_KEY: &str =
75    "x-greptime-internal-create-database-creator-bin";
76
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
78pub struct CreatorGrantIntent {
79    pub username: String,
80    pub created_at_ns: i64,
81}
82
83/// DDL tasks
84#[derive(Debug, Clone)]
85pub enum DdlTask {
86    CreateTable(CreateTableTask),
87    DropTable(DropTableTask),
88    UndropTable(UndropTableTask),
89    PurgeDroppedTable(PurgeDroppedTableTask),
90    AlterTable(AlterTableTask),
91    TruncateTable(TruncateTableTask),
92    CreateLogicalTables(Vec<CreateTableTask>),
93    DropLogicalTables(Vec<DropTableTask>),
94    AlterLogicalTables(Vec<AlterTableTask>),
95    CreateDatabase(CreateDatabaseTask),
96    DropDatabase(DropDatabaseTask),
97    AlterDatabase(AlterDatabaseTask),
98    CreateFlow(CreateFlowTask),
99    DropFlow(DropFlowTask),
100    #[cfg(feature = "enterprise")]
101    DropTrigger(trigger::DropTriggerTask),
102    CreateView(CreateViewTask),
103    DropView(DropViewTask),
104    #[cfg(feature = "enterprise")]
105    CreateTrigger(trigger::CreateTriggerTask),
106    CommentOn(CommentOnTask),
107}
108
109impl DdlTask {
110    /// Creates a [`DdlTask`] to create a flow.
111    pub fn new_create_flow(expr: CreateFlowTask) -> Self {
112        DdlTask::CreateFlow(expr)
113    }
114
115    /// Creates a [`DdlTask`] to drop a flow.
116    pub fn new_drop_flow(expr: DropFlowTask) -> Self {
117        DdlTask::DropFlow(expr)
118    }
119
120    /// Creates a [`DdlTask`] to drop a view.
121    pub fn new_drop_view(expr: DropViewTask) -> Self {
122        DdlTask::DropView(expr)
123    }
124
125    /// Creates a [`DdlTask`] to create a table.
126    pub fn new_create_table(
127        expr: CreateTableExpr,
128        partitions: Vec<Partition>,
129        table_info: TableInfo,
130    ) -> Self {
131        DdlTask::CreateTable(CreateTableTask::new(expr, partitions, table_info))
132    }
133
134    /// Creates a [`DdlTask`] to create several logical tables.
135    pub fn new_create_logical_tables(table_data: Vec<(CreateTableExpr, TableInfo)>) -> Self {
136        DdlTask::CreateLogicalTables(
137            table_data
138                .into_iter()
139                .map(|(expr, table_info)| CreateTableTask::new(expr, Vec::new(), table_info))
140                .collect(),
141        )
142    }
143
144    /// Creates a [`DdlTask`] to alter several logical tables.
145    pub fn new_alter_logical_tables(table_data: Vec<AlterTableExpr>) -> Self {
146        DdlTask::AlterLogicalTables(
147            table_data
148                .into_iter()
149                .map(|alter_table| AlterTableTask { alter_table })
150                .collect(),
151        )
152    }
153
154    /// Creates a [`DdlTask`] to drop a table.
155    pub fn new_drop_table(
156        catalog: String,
157        schema: String,
158        table: String,
159        table_id: TableId,
160        drop_if_exists: bool,
161    ) -> Self {
162        DdlTask::DropTable(DropTableTask {
163            catalog,
164            schema,
165            table,
166            table_id,
167            drop_if_exists,
168        })
169    }
170
171    /// Creates a [`DdlTask`] to undrop a table.
172    pub fn new_undrop_table(table_id: TableId) -> Self {
173        DdlTask::UndropTable(UndropTableTask { table_id })
174    }
175
176    /// Creates a [`DdlTask`] to purge a dropped table.
177    pub fn new_purge_dropped_table(table_id: TableId) -> Self {
178        DdlTask::PurgeDroppedTable(PurgeDroppedTableTask { table_id })
179    }
180
181    /// Creates a [`DdlTask`] to create a database.
182    pub fn new_create_database(
183        catalog: String,
184        schema: String,
185        create_if_not_exists: bool,
186        options: HashMap<String, String>,
187        creator: Option<CreatorGrantIntent>,
188    ) -> Self {
189        DdlTask::CreateDatabase(CreateDatabaseTask {
190            catalog,
191            schema,
192            create_if_not_exists,
193            options,
194            creator,
195        })
196    }
197
198    /// Creates a [`DdlTask`] to drop a database.
199    pub fn new_drop_database(catalog: String, schema: String, drop_if_exists: bool) -> Self {
200        DdlTask::DropDatabase(DropDatabaseTask {
201            catalog,
202            schema,
203            drop_if_exists,
204        })
205    }
206
207    /// Creates a [`DdlTask`] to alter a database.
208    pub fn new_alter_database(alter_expr: AlterDatabaseExpr) -> Self {
209        DdlTask::AlterDatabase(AlterDatabaseTask { alter_expr })
210    }
211
212    /// Creates a [`DdlTask`] to alter a table.
213    pub fn new_alter_table(alter_table: AlterTableExpr) -> Self {
214        DdlTask::AlterTable(AlterTableTask { alter_table })
215    }
216
217    /// Creates a [`DdlTask`] to truncate a table.
218    pub fn new_truncate_table(
219        catalog: String,
220        schema: String,
221        table: String,
222        table_id: TableId,
223        time_ranges: Vec<(Timestamp, Timestamp)>,
224    ) -> Self {
225        DdlTask::TruncateTable(TruncateTableTask {
226            catalog,
227            schema,
228            table,
229            table_id,
230            time_ranges,
231        })
232    }
233
234    /// Creates a [`DdlTask`] to create a view.
235    pub fn new_create_view(create_view: CreateViewExpr, view_info: TableInfo) -> Self {
236        DdlTask::CreateView(CreateViewTask {
237            create_view,
238            view_info,
239        })
240    }
241
242    /// Creates a [`DdlTask`] to comment on a table, column, or flow.
243    pub fn new_comment_on(task: CommentOnTask) -> Self {
244        DdlTask::CommentOn(task)
245    }
246}
247
248impl TryFrom<Task> for DdlTask {
249    type Error = error::Error;
250    fn try_from(task: Task) -> Result<Self> {
251        match task {
252            Task::CreateTableTask(create_table) => {
253                Ok(DdlTask::CreateTable(create_table.try_into()?))
254            }
255            Task::DropTableTask(drop_table) => Ok(DdlTask::DropTable(drop_table.try_into()?)),
256            Task::UndropTableTask(undrop_table) => {
257                Ok(DdlTask::UndropTable(undrop_table.try_into()?))
258            }
259            Task::PurgeDroppedTableTask(purge_dropped_table) => {
260                Ok(DdlTask::PurgeDroppedTable(purge_dropped_table.try_into()?))
261            }
262            Task::AlterTableTask(alter_table) => Ok(DdlTask::AlterTable(alter_table.try_into()?)),
263            Task::TruncateTableTask(truncate_table) => {
264                Ok(DdlTask::TruncateTable(truncate_table.try_into()?))
265            }
266            Task::CreateTableTasks(create_tables) => {
267                let tasks = create_tables
268                    .tasks
269                    .into_iter()
270                    .map(|task| task.try_into())
271                    .collect::<Result<Vec<_>>>()?;
272
273                Ok(DdlTask::CreateLogicalTables(tasks))
274            }
275            Task::DropTableTasks(drop_tables) => {
276                let tasks = drop_tables
277                    .tasks
278                    .into_iter()
279                    .map(|task| task.try_into())
280                    .collect::<Result<Vec<_>>>()?;
281
282                Ok(DdlTask::DropLogicalTables(tasks))
283            }
284            Task::AlterTableTasks(alter_tables) => {
285                let tasks = alter_tables
286                    .tasks
287                    .into_iter()
288                    .map(|task| task.try_into())
289                    .collect::<Result<Vec<_>>>()?;
290
291                Ok(DdlTask::AlterLogicalTables(tasks))
292            }
293            Task::CreateDatabaseTask(create_database) => {
294                Ok(DdlTask::CreateDatabase(create_database.try_into()?))
295            }
296            Task::DropDatabaseTask(drop_database) => {
297                Ok(DdlTask::DropDatabase(drop_database.try_into()?))
298            }
299            Task::AlterDatabaseTask(alter_database) => {
300                Ok(DdlTask::AlterDatabase(alter_database.try_into()?))
301            }
302            Task::CreateFlowTask(create_flow) => Ok(DdlTask::CreateFlow(create_flow.try_into()?)),
303            Task::DropFlowTask(drop_flow) => Ok(DdlTask::DropFlow(drop_flow.try_into()?)),
304            Task::CreateViewTask(create_view) => Ok(DdlTask::CreateView(create_view.try_into()?)),
305            Task::DropViewTask(drop_view) => Ok(DdlTask::DropView(drop_view.try_into()?)),
306            Task::CreateTriggerTask(create_trigger) => {
307                #[cfg(feature = "enterprise")]
308                return Ok(DdlTask::CreateTrigger(create_trigger.try_into()?));
309                #[cfg(not(feature = "enterprise"))]
310                {
311                    let _ = create_trigger;
312                    crate::error::UnsupportedSnafu {
313                        operation: "create trigger",
314                    }
315                    .fail()
316                }
317            }
318            Task::DropTriggerTask(drop_trigger) => {
319                #[cfg(feature = "enterprise")]
320                return Ok(DdlTask::DropTrigger(drop_trigger.try_into()?));
321                #[cfg(not(feature = "enterprise"))]
322                {
323                    let _ = drop_trigger;
324                    crate::error::UnsupportedSnafu {
325                        operation: "drop trigger",
326                    }
327                    .fail()
328                }
329            }
330            Task::CommentOnTask(comment_on) => Ok(DdlTask::CommentOn(comment_on.try_into()?)),
331        }
332    }
333}
334
335#[derive(Clone)]
336pub struct SubmitDdlTaskRequest {
337    pub wait: bool,
338    pub timeout: Duration,
339    pub task: DdlTask,
340}
341
342impl SubmitDdlTaskRequest {
343    /// The default constructor for [`SubmitDdlTaskRequest`].
344    pub fn new(task: DdlTask) -> Self {
345        Self {
346            wait: Self::default_wait(),
347            timeout: Self::default_timeout(),
348            task,
349        }
350    }
351
352    /// The default timeout for a DDL task.
353    pub fn default_timeout() -> Duration {
354        Duration::from_secs(60)
355    }
356
357    /// The default wait for a DDL task.
358    pub fn default_wait() -> bool {
359        true
360    }
361}
362
363fn ddl_timeout_secs(timeout: Duration) -> u32 {
364    timeout
365        .as_nanos()
366        .div_ceil(Duration::from_secs(1).as_nanos())
367        .try_into()
368        .unwrap_or(u32::MAX)
369}
370
371impl TryFrom<SubmitDdlTaskRequest> for PbDdlTaskRequest {
372    type Error = error::Error;
373
374    fn try_from(request: SubmitDdlTaskRequest) -> Result<Self> {
375        let SubmitDdlTaskRequest {
376            wait,
377            timeout,
378            task,
379        } = request;
380
381        let task = match task {
382            DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?),
383            DdlTask::DropTable(task) => Task::DropTableTask(task.into()),
384            DdlTask::UndropTable(task) => Task::UndropTableTask(task.into()),
385            DdlTask::PurgeDroppedTable(task) => Task::PurgeDroppedTableTask(task.into()),
386            DdlTask::AlterTable(task) => Task::AlterTableTask(task.try_into()?),
387            DdlTask::TruncateTable(task) => Task::TruncateTableTask(task.try_into()?),
388            DdlTask::CreateLogicalTables(tasks) => {
389                let tasks = tasks
390                    .into_iter()
391                    .map(|task| task.try_into())
392                    .collect::<Result<Vec<_>>>()?;
393
394                Task::CreateTableTasks(PbCreateTableTasks { tasks })
395            }
396            DdlTask::DropLogicalTables(tasks) => {
397                let tasks = tasks
398                    .into_iter()
399                    .map(|task| task.into())
400                    .collect::<Vec<_>>();
401
402                Task::DropTableTasks(PbDropTableTasks { tasks })
403            }
404            DdlTask::AlterLogicalTables(tasks) => {
405                let tasks = tasks
406                    .into_iter()
407                    .map(|task| task.try_into())
408                    .collect::<Result<Vec<_>>>()?;
409
410                Task::AlterTableTasks(PbAlterTableTasks { tasks })
411            }
412            DdlTask::CreateDatabase(task) => Task::CreateDatabaseTask(task.try_into()?),
413            DdlTask::DropDatabase(task) => Task::DropDatabaseTask(task.try_into()?),
414            DdlTask::AlterDatabase(task) => Task::AlterDatabaseTask(task.try_into()?),
415            DdlTask::CreateFlow(task) => Task::CreateFlowTask(task.into()),
416            DdlTask::DropFlow(task) => Task::DropFlowTask(task.into()),
417            DdlTask::CreateView(task) => Task::CreateViewTask(task.try_into()?),
418            DdlTask::DropView(task) => Task::DropViewTask(task.into()),
419            #[cfg(feature = "enterprise")]
420            DdlTask::CreateTrigger(task) => Task::CreateTriggerTask(task.try_into()?),
421            #[cfg(feature = "enterprise")]
422            DdlTask::DropTrigger(task) => Task::DropTriggerTask(task.into()),
423            DdlTask::CommentOn(task) => Task::CommentOnTask(task.into()),
424        };
425
426        Ok(Self {
427            header: None,
428            query_context: None,
429            timeout_secs: ddl_timeout_secs(timeout),
430            wait,
431            task: Some(task),
432            event_context: None,
433            actor: None,
434        })
435    }
436}
437
438#[derive(Debug, Default)]
439pub struct SubmitDdlTaskResponse {
440    pub key: Vec<u8>,
441    // `table_id`s for `CREATE TABLE` or `CREATE LOGICAL TABLES` task.
442    pub table_ids: Vec<TableId>,
443}
444
445impl TryFrom<PbDdlTaskResponse> for SubmitDdlTaskResponse {
446    type Error = error::Error;
447
448    fn try_from(resp: PbDdlTaskResponse) -> Result<Self> {
449        let table_ids = resp.table_ids.into_iter().map(|t| t.id).collect();
450        Ok(Self {
451            key: resp.pid.map(|pid| pid.key).unwrap_or_default(),
452            table_ids,
453        })
454    }
455}
456
457impl From<SubmitDdlTaskResponse> for PbDdlTaskResponse {
458    fn from(val: SubmitDdlTaskResponse) -> Self {
459        Self {
460            pid: Some(ProcedureId { key: val.key }),
461            table_ids: val
462                .table_ids
463                .into_iter()
464                .map(|id| api::v1::TableId { id })
465                .collect(),
466            ..Default::default()
467        }
468    }
469}
470
471/// A `CREATE VIEW` task.
472#[derive(Debug, PartialEq, Clone)]
473pub struct CreateViewTask {
474    pub create_view: CreateViewExpr,
475    pub view_info: TableInfo,
476}
477
478impl CreateViewTask {
479    /// Returns the [`TableReference`] of view.
480    pub fn table_ref(&self) -> TableReference<'_> {
481        TableReference {
482            catalog: &self.create_view.catalog_name,
483            schema: &self.create_view.schema_name,
484            table: &self.create_view.view_name,
485        }
486    }
487
488    /// Returns the encoded logical plan
489    pub fn raw_logical_plan(&self) -> &Vec<u8> {
490        &self.create_view.logical_plan
491    }
492
493    /// Returns the view definition in SQL
494    pub fn view_definition(&self) -> &str {
495        &self.create_view.definition
496    }
497
498    /// Returns the resolved table names in view's logical plan
499    pub fn table_names(&self) -> HashSet<TableName> {
500        self.create_view
501            .table_names
502            .iter()
503            .map(|t| t.clone().into())
504            .collect()
505    }
506
507    /// Returns the view's columns
508    pub fn columns(&self) -> &Vec<String> {
509        &self.create_view.columns
510    }
511
512    /// Returns the original logical plan's columns
513    pub fn plan_columns(&self) -> &Vec<String> {
514        &self.create_view.plan_columns
515    }
516}
517
518impl TryFrom<PbCreateViewTask> for CreateViewTask {
519    type Error = error::Error;
520
521    fn try_from(pb: PbCreateViewTask) -> Result<Self> {
522        let view_info = serde_json::from_slice(&pb.view_info).context(error::SerdeJsonSnafu)?;
523
524        Ok(CreateViewTask {
525            create_view: pb.create_view.context(error::InvalidProtoMsgSnafu {
526                err_msg: "expected create view",
527            })?,
528            view_info,
529        })
530    }
531}
532
533impl TryFrom<CreateViewTask> for PbCreateViewTask {
534    type Error = error::Error;
535
536    fn try_from(task: CreateViewTask) -> Result<PbCreateViewTask> {
537        Ok(PbCreateViewTask {
538            create_view: Some(task.create_view),
539            view_info: serde_json::to_vec(&task.view_info).context(error::SerdeJsonSnafu)?,
540        })
541    }
542}
543
544impl Serialize for CreateViewTask {
545    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
546    where
547        S: serde::Serializer,
548    {
549        let view_info = serde_json::to_vec(&self.view_info)
550            .map_err(|err| serde::ser::Error::custom(err.to_string()))?;
551
552        let pb = PbCreateViewTask {
553            create_view: Some(self.create_view.clone()),
554            view_info,
555        };
556        let buf = pb.encode_to_vec();
557        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
558        serializer.serialize_str(&encoded)
559    }
560}
561
562impl<'de> Deserialize<'de> for CreateViewTask {
563    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
564    where
565        D: serde::Deserializer<'de>,
566    {
567        let encoded = String::deserialize(deserializer)?;
568        let buf = general_purpose::STANDARD_NO_PAD
569            .decode(encoded)
570            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
571        let expr: PbCreateViewTask = PbCreateViewTask::decode(&*buf)
572            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
573
574        let expr = CreateViewTask::try_from(expr)
575            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
576
577        Ok(expr)
578    }
579}
580
581/// A `DROP VIEW` task.
582#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
583pub struct DropViewTask {
584    pub catalog: String,
585    pub schema: String,
586    pub view: String,
587    pub view_id: TableId,
588    pub drop_if_exists: bool,
589}
590
591impl DropViewTask {
592    /// Returns the [`TableReference`] of view.
593    pub fn table_ref(&self) -> TableReference<'_> {
594        TableReference {
595            catalog: &self.catalog,
596            schema: &self.schema,
597            table: &self.view,
598        }
599    }
600}
601
602impl TryFrom<PbDropViewTask> for DropViewTask {
603    type Error = error::Error;
604
605    fn try_from(pb: PbDropViewTask) -> Result<Self> {
606        let expr = pb.drop_view.context(error::InvalidProtoMsgSnafu {
607            err_msg: "expected drop view",
608        })?;
609
610        Ok(DropViewTask {
611            catalog: expr.catalog_name,
612            schema: expr.schema_name,
613            view: expr.view_name,
614            view_id: expr
615                .view_id
616                .context(error::InvalidProtoMsgSnafu {
617                    err_msg: "expected view_id",
618                })?
619                .id,
620            drop_if_exists: expr.drop_if_exists,
621        })
622    }
623}
624
625impl From<DropViewTask> for PbDropViewTask {
626    fn from(task: DropViewTask) -> Self {
627        PbDropViewTask {
628            drop_view: Some(DropViewExpr {
629                catalog_name: task.catalog,
630                schema_name: task.schema,
631                view_name: task.view,
632                view_id: Some(api::v1::TableId { id: task.view_id }),
633                drop_if_exists: task.drop_if_exists,
634            }),
635        }
636    }
637}
638
639#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
640pub struct DropTableTask {
641    pub catalog: String,
642    pub schema: String,
643    pub table: String,
644    pub table_id: TableId,
645    #[serde(default)]
646    pub drop_if_exists: bool,
647}
648
649#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
650pub struct UndropTableTask {
651    pub table_id: TableId,
652}
653
654#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
655pub struct PurgeDroppedTableTask {
656    pub table_id: TableId,
657}
658
659impl TryFrom<PbUndropTableTask> for UndropTableTask {
660    type Error = error::Error;
661
662    fn try_from(pb: PbUndropTableTask) -> Result<Self> {
663        Ok(Self {
664            table_id: pb
665                .table_id
666                .context(error::InvalidProtoMsgSnafu {
667                    err_msg: "expected table_id",
668                })?
669                .id,
670        })
671    }
672}
673
674impl From<UndropTableTask> for PbUndropTableTask {
675    fn from(task: UndropTableTask) -> Self {
676        Self {
677            table_id: Some(api::v1::TableId { id: task.table_id }),
678        }
679    }
680}
681
682impl TryFrom<PbPurgeDroppedTableTask> for PurgeDroppedTableTask {
683    type Error = error::Error;
684
685    fn try_from(pb: PbPurgeDroppedTableTask) -> Result<Self> {
686        Ok(Self {
687            table_id: pb
688                .table_id
689                .context(error::InvalidProtoMsgSnafu {
690                    err_msg: "expected table_id",
691                })?
692                .id,
693        })
694    }
695}
696
697impl From<PurgeDroppedTableTask> for PbPurgeDroppedTableTask {
698    fn from(task: PurgeDroppedTableTask) -> Self {
699        Self {
700            table_id: Some(api::v1::TableId { id: task.table_id }),
701        }
702    }
703}
704
705impl DropTableTask {
706    pub fn table_ref(&self) -> TableReference<'_> {
707        TableReference {
708            catalog: &self.catalog,
709            schema: &self.schema,
710            table: &self.table,
711        }
712    }
713
714    pub fn table_name(&self) -> TableName {
715        TableName {
716            catalog_name: self.catalog.clone(),
717            schema_name: self.schema.clone(),
718            table_name: self.table.clone(),
719        }
720    }
721}
722
723impl TryFrom<PbDropTableTask> for DropTableTask {
724    type Error = error::Error;
725
726    fn try_from(pb: PbDropTableTask) -> Result<Self> {
727        let drop_table = pb.drop_table.context(error::InvalidProtoMsgSnafu {
728            err_msg: "expected drop table",
729        })?;
730
731        Ok(Self {
732            catalog: drop_table.catalog_name,
733            schema: drop_table.schema_name,
734            table: drop_table.table_name,
735            table_id: drop_table
736                .table_id
737                .context(error::InvalidProtoMsgSnafu {
738                    err_msg: "expected table_id",
739                })?
740                .id,
741            drop_if_exists: drop_table.drop_if_exists,
742        })
743    }
744}
745
746impl From<DropTableTask> for PbDropTableTask {
747    fn from(task: DropTableTask) -> Self {
748        PbDropTableTask {
749            drop_table: Some(DropTableExpr {
750                catalog_name: task.catalog,
751                schema_name: task.schema,
752                table_name: task.table,
753                table_id: Some(api::v1::TableId { id: task.table_id }),
754                drop_if_exists: task.drop_if_exists,
755            }),
756        }
757    }
758}
759
760#[derive(Debug, PartialEq, Clone)]
761pub struct CreateTableTask {
762    pub create_table: CreateTableExpr,
763    pub partitions: Vec<Partition>,
764    pub table_info: TableInfo,
765}
766
767impl TryFrom<PbCreateTableTask> for CreateTableTask {
768    type Error = error::Error;
769
770    fn try_from(pb: PbCreateTableTask) -> Result<Self> {
771        let table_info = serde_json::from_slice(&pb.table_info).context(error::SerdeJsonSnafu)?;
772
773        Ok(CreateTableTask::new(
774            pb.create_table.context(error::InvalidProtoMsgSnafu {
775                err_msg: "expected create table",
776            })?,
777            pb.partitions,
778            table_info,
779        ))
780    }
781}
782
783impl TryFrom<CreateTableTask> for PbCreateTableTask {
784    type Error = error::Error;
785
786    fn try_from(task: CreateTableTask) -> Result<Self> {
787        Ok(PbCreateTableTask {
788            table_info: serde_json::to_vec(&task.table_info).context(error::SerdeJsonSnafu)?,
789            create_table: Some(task.create_table),
790            partitions: task.partitions,
791        })
792    }
793}
794
795impl CreateTableTask {
796    pub fn new(
797        expr: CreateTableExpr,
798        partitions: Vec<Partition>,
799        table_info: TableInfo,
800    ) -> CreateTableTask {
801        CreateTableTask {
802            create_table: expr,
803            partitions,
804            table_info,
805        }
806    }
807
808    pub fn table_name(&self) -> TableName {
809        let table = &self.create_table;
810
811        TableName {
812            catalog_name: table.catalog_name.clone(),
813            schema_name: table.schema_name.clone(),
814            table_name: table.table_name.clone(),
815        }
816    }
817
818    pub fn table_ref(&self) -> TableReference<'_> {
819        let table = &self.create_table;
820
821        TableReference {
822            catalog: &table.catalog_name,
823            schema: &table.schema_name,
824            table: &table.table_name,
825        }
826    }
827
828    /// Sets the `table_info`'s table_id.
829    pub fn set_table_id(&mut self, table_id: TableId) {
830        self.table_info.ident.table_id = table_id;
831    }
832
833    /// Sort the columns in [CreateTableExpr] and [TableInfo].
834    ///
835    /// This function won't do any check or verification. Caller should
836    /// ensure this task is valid.
837    pub fn sort_columns(&mut self) {
838        // sort create table expr
839        // sort column_defs by name
840        self.create_table
841            .column_defs
842            .sort_unstable_by(|a, b| a.name.cmp(&b.name));
843
844        self.table_info.sort_columns();
845    }
846}
847
848impl Serialize for CreateTableTask {
849    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
850    where
851        S: serde::Serializer,
852    {
853        let table_info = serde_json::to_vec(&self.table_info)
854            .map_err(|err| serde::ser::Error::custom(err.to_string()))?;
855
856        let pb = PbCreateTableTask {
857            create_table: Some(self.create_table.clone()),
858            partitions: self.partitions.clone(),
859            table_info,
860        };
861        let buf = pb.encode_to_vec();
862        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
863        serializer.serialize_str(&encoded)
864    }
865}
866
867impl<'de> Deserialize<'de> for CreateTableTask {
868    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
869    where
870        D: serde::Deserializer<'de>,
871    {
872        let encoded = String::deserialize(deserializer)?;
873        let buf = general_purpose::STANDARD_NO_PAD
874            .decode(encoded)
875            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
876        let expr: PbCreateTableTask = PbCreateTableTask::decode(&*buf)
877            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
878
879        let expr = CreateTableTask::try_from(expr)
880            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
881
882        Ok(expr)
883    }
884}
885
886#[derive(Debug, PartialEq, Clone)]
887pub struct AlterTableTask {
888    // TODO(CookiePieWw): Replace proto struct with user-defined struct
889    pub alter_table: AlterTableExpr,
890}
891
892impl AlterTableTask {
893    pub fn validate(&self) -> Result<()> {
894        self.alter_table
895            .kind
896            .as_ref()
897            .context(error::UnexpectedSnafu {
898                err_msg: "'kind' is absent",
899            })?;
900        Ok(())
901    }
902
903    pub fn table_ref(&self) -> TableReference<'_> {
904        TableReference {
905            catalog: &self.alter_table.catalog_name,
906            schema: &self.alter_table.schema_name,
907            table: &self.alter_table.table_name,
908        }
909    }
910
911    pub fn table_name(&self) -> TableName {
912        let table = &self.alter_table;
913
914        TableName {
915            catalog_name: table.catalog_name.clone(),
916            schema_name: table.schema_name.clone(),
917            table_name: table.table_name.clone(),
918        }
919    }
920}
921
922impl TryFrom<PbAlterTableTask> for AlterTableTask {
923    type Error = error::Error;
924
925    fn try_from(pb: PbAlterTableTask) -> Result<Self> {
926        let alter_table = pb.alter_table.context(error::InvalidProtoMsgSnafu {
927            err_msg: "expected alter_table",
928        })?;
929
930        Ok(AlterTableTask { alter_table })
931    }
932}
933
934impl TryFrom<AlterTableTask> for PbAlterTableTask {
935    type Error = error::Error;
936
937    fn try_from(task: AlterTableTask) -> Result<Self> {
938        Ok(PbAlterTableTask {
939            alter_table: Some(task.alter_table),
940        })
941    }
942}
943
944impl Serialize for AlterTableTask {
945    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
946    where
947        S: serde::Serializer,
948    {
949        let pb = PbAlterTableTask {
950            alter_table: Some(self.alter_table.clone()),
951        };
952        let buf = pb.encode_to_vec();
953        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
954        serializer.serialize_str(&encoded)
955    }
956}
957
958impl<'de> Deserialize<'de> for AlterTableTask {
959    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
960    where
961        D: serde::Deserializer<'de>,
962    {
963        let encoded = String::deserialize(deserializer)?;
964        let buf = general_purpose::STANDARD_NO_PAD
965            .decode(encoded)
966            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
967        let expr: PbAlterTableTask = PbAlterTableTask::decode(&*buf)
968            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
969
970        let expr = AlterTableTask::try_from(expr)
971            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
972
973        Ok(expr)
974    }
975}
976
977#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
978pub struct TruncateTableTask {
979    pub catalog: String,
980    pub schema: String,
981    pub table: String,
982    pub table_id: TableId,
983    pub time_ranges: Vec<(Timestamp, Timestamp)>,
984}
985
986impl TruncateTableTask {
987    pub fn table_ref(&self) -> TableReference<'_> {
988        TableReference {
989            catalog: &self.catalog,
990            schema: &self.schema,
991            table: &self.table,
992        }
993    }
994
995    pub fn table_name(&self) -> TableName {
996        TableName {
997            catalog_name: self.catalog.clone(),
998            schema_name: self.schema.clone(),
999            table_name: self.table.clone(),
1000        }
1001    }
1002}
1003
1004impl TryFrom<PbTruncateTableTask> for TruncateTableTask {
1005    type Error = error::Error;
1006
1007    fn try_from(pb: PbTruncateTableTask) -> Result<Self> {
1008        let truncate_table = pb.truncate_table.context(error::InvalidProtoMsgSnafu {
1009            err_msg: "expected truncate table",
1010        })?;
1011
1012        Ok(Self {
1013            catalog: truncate_table.catalog_name,
1014            schema: truncate_table.schema_name,
1015            table: truncate_table.table_name,
1016            table_id: truncate_table
1017                .table_id
1018                .context(error::InvalidProtoMsgSnafu {
1019                    err_msg: "expected table_id",
1020                })?
1021                .id,
1022            time_ranges: truncate_table
1023                .time_ranges
1024                .map(from_pb_time_ranges)
1025                .transpose()
1026                .map_err(BoxedError::new)
1027                .context(ExternalSnafu)?
1028                .unwrap_or_default(),
1029        })
1030    }
1031}
1032
1033impl TryFrom<TruncateTableTask> for PbTruncateTableTask {
1034    type Error = error::Error;
1035
1036    fn try_from(task: TruncateTableTask) -> Result<Self> {
1037        Ok(PbTruncateTableTask {
1038            truncate_table: Some(TruncateTableExpr {
1039                catalog_name: task.catalog,
1040                schema_name: task.schema,
1041                table_name: task.table,
1042                table_id: Some(api::v1::TableId { id: task.table_id }),
1043                time_ranges: Some(
1044                    to_pb_time_ranges(&task.time_ranges).context(ConvertTimeRangesSnafu)?,
1045                ),
1046            }),
1047        })
1048    }
1049}
1050
1051#[serde_as]
1052#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
1053pub struct CreateDatabaseTask {
1054    pub catalog: String,
1055    pub schema: String,
1056    pub create_if_not_exists: bool,
1057    #[serde_as(deserialize_as = "DefaultOnNull")]
1058    pub options: HashMap<String, String>,
1059    #[serde(default)]
1060    pub creator: Option<CreatorGrantIntent>,
1061}
1062
1063impl TryFrom<PbCreateDatabaseTask> for CreateDatabaseTask {
1064    type Error = error::Error;
1065
1066    fn try_from(pb: PbCreateDatabaseTask) -> Result<Self> {
1067        let CreateDatabaseExpr {
1068            catalog_name,
1069            schema_name,
1070            create_if_not_exists,
1071            options,
1072        } = pb.create_database.context(error::InvalidProtoMsgSnafu {
1073            err_msg: "expected create database",
1074        })?;
1075
1076        Ok(CreateDatabaseTask {
1077            catalog: catalog_name,
1078            schema: schema_name,
1079            create_if_not_exists,
1080            options,
1081            creator: None,
1082        })
1083    }
1084}
1085
1086impl TryFrom<CreateDatabaseTask> for PbCreateDatabaseTask {
1087    type Error = error::Error;
1088
1089    fn try_from(
1090        CreateDatabaseTask {
1091            catalog,
1092            schema,
1093            create_if_not_exists,
1094            options,
1095            creator: _,
1096        }: CreateDatabaseTask,
1097    ) -> Result<Self> {
1098        Ok(PbCreateDatabaseTask {
1099            create_database: Some(CreateDatabaseExpr {
1100                catalog_name: catalog,
1101                schema_name: schema,
1102                create_if_not_exists,
1103                options,
1104            }),
1105        })
1106    }
1107}
1108
1109#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
1110pub struct DropDatabaseTask {
1111    pub catalog: String,
1112    pub schema: String,
1113    pub drop_if_exists: bool,
1114}
1115
1116impl TryFrom<PbDropDatabaseTask> for DropDatabaseTask {
1117    type Error = error::Error;
1118
1119    fn try_from(pb: PbDropDatabaseTask) -> Result<Self> {
1120        let DropDatabaseExpr {
1121            catalog_name,
1122            schema_name,
1123            drop_if_exists,
1124        } = pb.drop_database.context(error::InvalidProtoMsgSnafu {
1125            err_msg: "expected drop database",
1126        })?;
1127
1128        Ok(DropDatabaseTask {
1129            catalog: catalog_name,
1130            schema: schema_name,
1131            drop_if_exists,
1132        })
1133    }
1134}
1135
1136impl TryFrom<DropDatabaseTask> for PbDropDatabaseTask {
1137    type Error = error::Error;
1138
1139    fn try_from(
1140        DropDatabaseTask {
1141            catalog,
1142            schema,
1143            drop_if_exists,
1144        }: DropDatabaseTask,
1145    ) -> Result<Self> {
1146        Ok(PbDropDatabaseTask {
1147            drop_database: Some(DropDatabaseExpr {
1148                catalog_name: catalog,
1149                schema_name: schema,
1150                drop_if_exists,
1151            }),
1152        })
1153    }
1154}
1155
1156#[derive(Debug, PartialEq, Clone)]
1157pub struct AlterDatabaseTask {
1158    pub alter_expr: AlterDatabaseExpr,
1159}
1160
1161impl TryFrom<AlterDatabaseTask> for PbAlterDatabaseTask {
1162    type Error = error::Error;
1163
1164    fn try_from(task: AlterDatabaseTask) -> Result<Self> {
1165        Ok(PbAlterDatabaseTask {
1166            task: Some(task.alter_expr),
1167        })
1168    }
1169}
1170
1171impl TryFrom<PbAlterDatabaseTask> for AlterDatabaseTask {
1172    type Error = error::Error;
1173
1174    fn try_from(pb: PbAlterDatabaseTask) -> Result<Self> {
1175        let alter_expr = pb.task.context(error::InvalidProtoMsgSnafu {
1176            err_msg: "expected alter database",
1177        })?;
1178
1179        Ok(AlterDatabaseTask { alter_expr })
1180    }
1181}
1182
1183impl TryFrom<PbAlterDatabaseKind> for AlterDatabaseKind {
1184    type Error = error::Error;
1185
1186    fn try_from(pb: PbAlterDatabaseKind) -> Result<Self> {
1187        match pb {
1188            PbAlterDatabaseKind::SetDatabaseOptions(options) => {
1189                Ok(AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(
1190                    options
1191                        .set_database_options
1192                        .into_iter()
1193                        .map(SetDatabaseOption::try_from)
1194                        .collect::<Result<Vec<_>>>()?,
1195                )))
1196            }
1197            PbAlterDatabaseKind::UnsetDatabaseOptions(options) => Ok(
1198                AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(
1199                    options
1200                        .keys
1201                        .iter()
1202                        .map(|key| UnsetDatabaseOption::try_from(key.as_str()))
1203                        .collect::<Result<Vec<_>>>()?,
1204                )),
1205            ),
1206        }
1207    }
1208}
1209
1210const TTL_KEY: &str = "ttl";
1211
1212impl TryFrom<PbOption> for SetDatabaseOption {
1213    type Error = error::Error;
1214
1215    fn try_from(PbOption { key, value }: PbOption) -> Result<Self> {
1216        let key_lower = key.to_ascii_lowercase();
1217        match key_lower.as_str() {
1218            TTL_KEY => {
1219                let ttl = DatabaseTimeToLive::from_humantime_or_str(&value)
1220                    .map_err(|_| InvalidSetDatabaseOptionSnafu { key, value }.build())?;
1221
1222                Ok(SetDatabaseOption::Ttl(ttl))
1223            }
1224            _ => {
1225                if validate_database_option(&key_lower) {
1226                    Ok(SetDatabaseOption::Other(key_lower, value))
1227                } else {
1228                    InvalidSetDatabaseOptionSnafu { key, value }.fail()
1229                }
1230            }
1231        }
1232    }
1233}
1234
1235#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
1236pub enum SetDatabaseOption {
1237    Ttl(DatabaseTimeToLive),
1238    Other(String, String),
1239}
1240
1241#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
1242pub enum UnsetDatabaseOption {
1243    Ttl,
1244    Other(String),
1245}
1246
1247impl TryFrom<&str> for UnsetDatabaseOption {
1248    type Error = error::Error;
1249
1250    fn try_from(key: &str) -> Result<Self> {
1251        let key_lower = key.to_ascii_lowercase();
1252        match key_lower.as_str() {
1253            TTL_KEY => Ok(UnsetDatabaseOption::Ttl),
1254            _ => {
1255                if validate_database_option(&key_lower) {
1256                    Ok(UnsetDatabaseOption::Other(key_lower))
1257                } else {
1258                    InvalidUnsetDatabaseOptionSnafu { key }.fail()
1259                }
1260            }
1261        }
1262    }
1263}
1264
1265#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
1266pub struct SetDatabaseOptions(pub Vec<SetDatabaseOption>);
1267
1268#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
1269pub struct UnsetDatabaseOptions(pub Vec<UnsetDatabaseOption>);
1270
1271#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
1272pub enum AlterDatabaseKind {
1273    SetDatabaseOptions(SetDatabaseOptions),
1274    UnsetDatabaseOptions(UnsetDatabaseOptions),
1275}
1276
1277impl AlterDatabaseTask {
1278    pub fn catalog(&self) -> &str {
1279        &self.alter_expr.catalog_name
1280    }
1281
1282    pub fn schema(&self) -> &str {
1283        &self.alter_expr.catalog_name
1284    }
1285}
1286
1287/// Create flow
1288#[derive(Debug, Clone, Serialize, Deserialize)]
1289pub struct CreateFlowTask {
1290    pub catalog_name: String,
1291    pub flow_name: String,
1292    pub source_table_names: Vec<TableName>,
1293    pub sink_table_name: TableName,
1294    pub or_replace: bool,
1295    pub create_if_not_exists: bool,
1296    /// Duration in seconds. Data older than this duration will not be used.
1297    pub expire_after: Option<i64>,
1298    pub eval_interval_secs: Option<i64>,
1299    /// Phase offset of the evaluation schedule within `eval_interval_secs`,
1300    /// in seconds. Must be in `[0, eval_interval_secs)`. `None` means a zero
1301    /// offset (epoch-anchored schedule).
1302    /// Transported through the transient option map (no proto field), see
1303    /// `INTERNAL_EVAL_OFFSET_KEY`.
1304    #[serde(default)]
1305    pub eval_offset_secs: Option<i64>,
1306    pub comment: String,
1307    pub sql: String,
1308    pub flow_options: HashMap<String, String>,
1309    /// Typed schedule configuration resolved during `on_prepare`.
1310    /// Not populated from proto; set by the procedure layer after
1311    /// defaults are resolved.
1312    #[serde(default)]
1313    pub eval_schedule: Option<crate::key::flow::flow_info::FlowScheduleConfig>,
1314}
1315
1316impl TryFrom<PbCreateFlowTask> for CreateFlowTask {
1317    type Error = error::Error;
1318
1319    fn try_from(pb: PbCreateFlowTask) -> Result<Self> {
1320        let CreateFlowExpr {
1321            catalog_name,
1322            flow_name,
1323            source_table_names,
1324            sink_table_name,
1325            or_replace,
1326            create_if_not_exists,
1327            expire_after,
1328            eval_interval,
1329            comment,
1330            sql,
1331            mut flow_options,
1332        } = pb.create_flow.context(error::InvalidProtoMsgSnafu {
1333            err_msg: "expected create_flow",
1334        })?;
1335
1336        // Parse and strip the trusted transient offset key inserted by the
1337        // operator after user option validation. It must never persist in
1338        // user-visible options.
1339        let eval_offset_secs =
1340            match flow_options.remove(crate::ddl::create_flow::INTERNAL_EVAL_OFFSET_KEY) {
1341                Some(value) => Some(value.parse::<i64>().map_err(|_| {
1342                    error::UnexpectedSnafu {
1343                        err_msg: format!(
1344                            "Invalid internal eval offset payload '{value}': expected whole seconds"
1345                        ),
1346                    }
1347                    .build()
1348                })?),
1349                None => None,
1350            };
1351
1352        Ok(CreateFlowTask {
1353            catalog_name,
1354            flow_name,
1355            source_table_names: source_table_names.into_iter().map(Into::into).collect(),
1356            sink_table_name: sink_table_name
1357                .context(error::InvalidProtoMsgSnafu {
1358                    err_msg: "expected sink_table_name",
1359                })?
1360                .into(),
1361            or_replace,
1362            create_if_not_exists,
1363            expire_after: expire_after.map(|e| e.value),
1364            eval_interval_secs: eval_interval.map(|e| e.seconds),
1365            eval_offset_secs,
1366            comment,
1367            sql,
1368            flow_options,
1369            eval_schedule: None,
1370        })
1371    }
1372}
1373
1374impl From<CreateFlowTask> for PbCreateFlowTask {
1375    fn from(
1376        CreateFlowTask {
1377            catalog_name,
1378            flow_name,
1379            source_table_names,
1380            sink_table_name,
1381            or_replace,
1382            create_if_not_exists,
1383            expire_after,
1384            eval_interval_secs: eval_interval,
1385            eval_offset_secs,
1386            comment,
1387            sql,
1388            mut flow_options,
1389            ..
1390        }: CreateFlowTask,
1391    ) -> Self {
1392        // Re-insert the transient offset key so the proto round-trip (e.g. DDL
1393        // task submission between frontend and metasrv) preserves the offset.
1394        if let Some(offset_secs) = eval_offset_secs {
1395            flow_options.insert(
1396                crate::ddl::create_flow::INTERNAL_EVAL_OFFSET_KEY.to_string(),
1397                offset_secs.to_string(),
1398            );
1399        }
1400        PbCreateFlowTask {
1401            create_flow: Some(CreateFlowExpr {
1402                catalog_name,
1403                flow_name,
1404                source_table_names: source_table_names.into_iter().map(Into::into).collect(),
1405                sink_table_name: Some(sink_table_name.into()),
1406                or_replace,
1407                create_if_not_exists,
1408                expire_after: expire_after.map(|value| ExpireAfter { value }),
1409                eval_interval: eval_interval.map(|seconds| EvalInterval { seconds }),
1410                comment,
1411                sql,
1412                flow_options,
1413            }),
1414        }
1415    }
1416}
1417
1418/// Drop flow
1419#[derive(Debug, Clone, Serialize, Deserialize)]
1420pub struct DropFlowTask {
1421    pub catalog_name: String,
1422    pub flow_name: String,
1423    pub flow_id: FlowId,
1424    pub drop_if_exists: bool,
1425}
1426
1427impl TryFrom<PbDropFlowTask> for DropFlowTask {
1428    type Error = error::Error;
1429
1430    fn try_from(pb: PbDropFlowTask) -> Result<Self> {
1431        let DropFlowExpr {
1432            catalog_name,
1433            flow_name,
1434            flow_id,
1435            drop_if_exists,
1436        } = pb.drop_flow.context(error::InvalidProtoMsgSnafu {
1437            err_msg: "expected drop_flow",
1438        })?;
1439        let flow_id = flow_id
1440            .context(error::InvalidProtoMsgSnafu {
1441                err_msg: "expected flow_id",
1442            })?
1443            .id;
1444        Ok(DropFlowTask {
1445            catalog_name,
1446            flow_name,
1447            flow_id,
1448            drop_if_exists,
1449        })
1450    }
1451}
1452
1453impl From<DropFlowTask> for PbDropFlowTask {
1454    fn from(
1455        DropFlowTask {
1456            catalog_name,
1457            flow_name,
1458            flow_id,
1459            drop_if_exists,
1460        }: DropFlowTask,
1461    ) -> Self {
1462        PbDropFlowTask {
1463            drop_flow: Some(DropFlowExpr {
1464                catalog_name,
1465                flow_name,
1466                flow_id: Some(api::v1::FlowId { id: flow_id }),
1467                drop_if_exists,
1468            }),
1469        }
1470    }
1471}
1472
1473/// Represents the ID of the object being commented on (Table or Flow).
1474#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1475pub enum CommentObjectId {
1476    Table(TableId),
1477    Flow(FlowId),
1478}
1479
1480/// Comment on table, column, or flow
1481#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1482pub struct CommentOnTask {
1483    pub catalog_name: String,
1484    pub schema_name: String,
1485    pub object_type: CommentObjectType,
1486    pub object_name: String,
1487    /// Column name (only for Column comments)
1488    pub column_name: Option<String>,
1489    /// Object ID (Table or Flow) for validation and cache invalidation
1490    pub object_id: Option<CommentObjectId>,
1491    pub comment: Option<String>,
1492}
1493
1494#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1495pub enum CommentObjectType {
1496    Table,
1497    Column,
1498    Flow,
1499}
1500
1501impl CommentOnTask {
1502    pub fn table_id(&self) -> Option<TableId> {
1503        match self.object_id.as_ref() {
1504            Some(CommentObjectId::Table(table_id)) => Some(*table_id),
1505            _ => None,
1506        }
1507    }
1508
1509    pub fn flow_id(&self) -> Option<FlowId> {
1510        match self.object_id.as_ref() {
1511            Some(CommentObjectId::Flow(flow_id)) => Some(*flow_id),
1512            _ => None,
1513        }
1514    }
1515
1516    fn set_table_id(&mut self, table_id: TableId) {
1517        self.object_id = Some(CommentObjectId::Table(table_id));
1518    }
1519
1520    fn set_flow_id(&mut self, flow_id: FlowId) {
1521        self.object_id = Some(CommentObjectId::Flow(flow_id));
1522    }
1523
1524    /// Returns the cache identifiers for the object being commented on.
1525    pub fn cache_idents(&self) -> Vec<CacheIdent> {
1526        match self.object_type {
1527            CommentObjectType::Table | CommentObjectType::Column => {
1528                let mut cache_idents = Vec::with_capacity(2);
1529                if let Some(CommentObjectId::Table(table_id)) = self.object_id.as_ref() {
1530                    cache_idents.push(CacheIdent::TableId(*table_id));
1531                }
1532                cache_idents.push(CacheIdent::TableName(TableName {
1533                    catalog_name: self.catalog_name.clone(),
1534                    schema_name: self.schema_name.clone(),
1535                    table_name: self.object_name.clone(),
1536                }));
1537                cache_idents
1538            }
1539            CommentObjectType::Flow => {
1540                let mut cache_idents = Vec::with_capacity(2);
1541                if let Some(CommentObjectId::Flow(flow_id)) = self.object_id.as_ref() {
1542                    cache_idents.push(CacheIdent::FlowId(*flow_id));
1543                }
1544                cache_idents.push(CacheIdent::FlowName(FlowName {
1545                    catalog_name: self.catalog_name.clone(),
1546                    flow_name: self.object_name.clone(),
1547                }));
1548                cache_idents
1549            }
1550        }
1551    }
1552
1553    /// Enriches the `object_id` field of the `CommentOnTask`
1554    /// by looking up the corresponding table or flow ID using the provided managers.
1555    pub async fn enrich_object_id(
1556        &mut self,
1557        table_name_manager: &TableNameManager,
1558        flow_name_manager: &FlowNameManager,
1559    ) -> Result<()> {
1560        match self.object_type {
1561            CommentObjectType::Table | CommentObjectType::Column => {
1562                let table_id = table_name_manager
1563                    .get(TableNameKey::new(
1564                        &self.catalog_name,
1565                        &self.schema_name,
1566                        &self.object_name,
1567                    ))
1568                    .await?
1569                    .with_context(|| error::TableNotFoundSnafu {
1570                        table_name: format_full_table_name(
1571                            &self.catalog_name,
1572                            &self.schema_name,
1573                            &self.object_name,
1574                        ),
1575                    })?
1576                    .table_id();
1577
1578                self.set_table_id(table_id);
1579            }
1580            CommentObjectType::Flow => {
1581                let flow_id = flow_name_manager
1582                    .get(&self.catalog_name, &self.object_name)
1583                    .await?
1584                    .with_context(|| error::FlowNotFoundSnafu {
1585                        flow_name: format_full_flow_name(&self.catalog_name, &self.object_name),
1586                    })?
1587                    .flow_id();
1588
1589                self.set_flow_id(flow_id);
1590            }
1591        }
1592
1593        Ok(())
1594    }
1595}
1596
1597// Proto conversions for CommentObjectType
1598impl From<CommentObjectType> for PbCommentObjectType {
1599    fn from(object_type: CommentObjectType) -> Self {
1600        match object_type {
1601            CommentObjectType::Table => PbCommentObjectType::Table,
1602            CommentObjectType::Column => PbCommentObjectType::Column,
1603            CommentObjectType::Flow => PbCommentObjectType::Flow,
1604        }
1605    }
1606}
1607
1608impl TryFrom<i32> for CommentObjectType {
1609    type Error = error::Error;
1610
1611    fn try_from(value: i32) -> Result<Self> {
1612        match value {
1613            0 => Ok(CommentObjectType::Table),
1614            1 => Ok(CommentObjectType::Column),
1615            2 => Ok(CommentObjectType::Flow),
1616            _ => error::InvalidProtoMsgSnafu {
1617                err_msg: format!(
1618                    "Invalid CommentObjectType value: {}. Valid values are: 0 (Table), 1 (Column), 2 (Flow)",
1619                    value
1620                ),
1621            }
1622            .fail(),
1623        }
1624    }
1625}
1626
1627// Proto conversions for CommentOnTask
1628impl TryFrom<PbCommentOnTask> for CommentOnTask {
1629    type Error = error::Error;
1630
1631    fn try_from(pb: PbCommentOnTask) -> Result<Self> {
1632        let comment_on = pb.comment_on.context(error::InvalidProtoMsgSnafu {
1633            err_msg: "expected comment_on",
1634        })?;
1635
1636        Ok(CommentOnTask {
1637            catalog_name: comment_on.catalog_name,
1638            schema_name: comment_on.schema_name,
1639            object_type: comment_on.object_type.try_into()?,
1640            object_name: comment_on.object_name,
1641            column_name: if comment_on.column_name.is_empty() {
1642                None
1643            } else {
1644                Some(comment_on.column_name)
1645            },
1646            comment: if comment_on.comment.is_empty() {
1647                None
1648            } else {
1649                Some(comment_on.comment)
1650            },
1651            object_id: None,
1652        })
1653    }
1654}
1655
1656impl From<CommentOnTask> for PbCommentOnTask {
1657    fn from(task: CommentOnTask) -> Self {
1658        let pb_object_type: PbCommentObjectType = task.object_type.into();
1659        PbCommentOnTask {
1660            comment_on: Some(CommentOnExpr {
1661                catalog_name: task.catalog_name,
1662                schema_name: task.schema_name,
1663                object_type: pb_object_type as i32,
1664                object_name: task.object_name,
1665                column_name: task.column_name.unwrap_or_default(),
1666                comment: task.comment.unwrap_or_default(),
1667            }),
1668        }
1669    }
1670}
1671
1672#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
1673pub struct QueryContext {
1674    pub current_catalog: String,
1675    pub current_schema: String,
1676    pub timezone: String,
1677    pub extensions: HashMap<String, String>,
1678    pub channel: u8,
1679    /// Maps region id -> snapshot upper bound sequence for that region.
1680    #[serde(default)]
1681    pub snapshot_seqs: HashMap<u64, u64>,
1682    /// Maps region id -> minimal SST sequence allowed for that region.
1683    #[serde(default)]
1684    pub sst_min_sequences: HashMap<u64, u64>,
1685}
1686
1687impl QueryContext {
1688    /// Get the current catalog
1689    pub fn current_catalog(&self) -> &str {
1690        &self.current_catalog
1691    }
1692
1693    /// Get the current schema
1694    pub fn current_schema(&self) -> &str {
1695        &self.current_schema
1696    }
1697
1698    /// Get the timezone
1699    pub fn timezone(&self) -> &str {
1700        &self.timezone
1701    }
1702
1703    /// Get the extensions
1704    pub fn extensions(&self) -> &HashMap<String, String> {
1705        &self.extensions
1706    }
1707
1708    /// Get the channel
1709    pub fn channel(&self) -> u8 {
1710        self.channel
1711    }
1712
1713    /// Returns the protocol derived from the typed query channel.
1714    pub fn protocol(&self) -> Option<String> {
1715        let channel = Channel::from(u32::from(self.channel));
1716        (channel != Channel::Unknown).then(|| channel.as_ref().to_string())
1717    }
1718
1719    pub fn snapshot_seqs(&self) -> &HashMap<u64, u64> {
1720        &self.snapshot_seqs
1721    }
1722
1723    pub fn sst_min_sequences(&self) -> &HashMap<u64, u64> {
1724        &self.sst_min_sequences
1725    }
1726}
1727
1728/// Lightweight query context for flow operations containing only essential fields.
1729/// This is a subset of QueryContext that includes only the fields actually needed
1730/// for flow creation and execution.
1731#[derive(Debug, Clone, Serialize, PartialEq)]
1732pub struct FlowQueryContext {
1733    /// Current catalog name used for flow metadata and execution.
1734    pub catalog: String,
1735    /// Current schema name used for table resolution during flow execution.
1736    pub schema: String,
1737    /// Timezone used for timestamp evaluation in the flow.
1738    pub timezone: String,
1739    /// Query extensions carried into flow execution.
1740    #[serde(default)]
1741    pub extensions: HashMap<String, String>,
1742    /// Request channel propagated from the original query context.
1743    #[serde(default)]
1744    pub channel: u8,
1745    /// Per-region snapshot upper bounds bound during query planning/execution.
1746    #[serde(default)]
1747    pub snapshot_seqs: HashMap<u64, u64>,
1748    /// Per-region lower SST scan bounds carried with the flow context.
1749    #[serde(default)]
1750    pub sst_min_sequences: HashMap<u64, u64>,
1751}
1752
1753impl<'de> Deserialize<'de> for FlowQueryContext {
1754    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
1755    where
1756        D: serde::Deserializer<'de>,
1757    {
1758        // Support both QueryContext format and FlowQueryContext format
1759        #[derive(Deserialize)]
1760        #[serde(untagged)]
1761        enum ContextCompat {
1762            Flow(FlowQueryContextHelper),
1763            Full(QueryContext),
1764        }
1765
1766        #[derive(Deserialize)]
1767        struct FlowQueryContextHelper {
1768            catalog: String,
1769            schema: String,
1770            timezone: String,
1771            #[serde(default)]
1772            extensions: HashMap<String, String>,
1773            #[serde(default)]
1774            channel: u8,
1775            #[serde(default)]
1776            snapshot_seqs: HashMap<u64, u64>,
1777            #[serde(default)]
1778            sst_min_sequences: HashMap<u64, u64>,
1779        }
1780
1781        match ContextCompat::deserialize(deserializer)? {
1782            ContextCompat::Flow(helper) => Ok(FlowQueryContext {
1783                catalog: helper.catalog,
1784                schema: helper.schema,
1785                timezone: helper.timezone,
1786                extensions: helper.extensions,
1787                channel: helper.channel,
1788                snapshot_seqs: helper.snapshot_seqs,
1789                sst_min_sequences: helper.sst_min_sequences,
1790            }),
1791            ContextCompat::Full(full_ctx) => Ok(full_ctx.into()),
1792        }
1793    }
1794}
1795
1796impl From<PbQueryContext> for QueryContext {
1797    fn from(pb_ctx: PbQueryContext) -> Self {
1798        let (snapshot_seqs, sst_min_sequences) = pb_ctx
1799            .snapshot_seqs
1800            .map(|seqs| (seqs.snapshot_seqs, seqs.sst_min_sequences))
1801            .unwrap_or_default();
1802
1803        Self {
1804            current_catalog: pb_ctx.current_catalog,
1805            current_schema: pb_ctx.current_schema,
1806            timezone: pb_ctx.timezone,
1807            extensions: pb_ctx.extensions,
1808            channel: pb_ctx.channel as u8,
1809            snapshot_seqs,
1810            sst_min_sequences,
1811        }
1812    }
1813}
1814
1815impl From<QueryContext> for PbQueryContext {
1816    fn from(
1817        QueryContext {
1818            current_catalog,
1819            current_schema,
1820            timezone,
1821            extensions,
1822            channel,
1823            snapshot_seqs,
1824            sst_min_sequences,
1825        }: QueryContext,
1826    ) -> Self {
1827        PbQueryContext {
1828            current_catalog,
1829            current_schema,
1830            timezone,
1831            extensions,
1832            channel: channel as u32,
1833            snapshot_seqs: (!snapshot_seqs.is_empty() || !sst_min_sequences.is_empty()).then_some(
1834                api::v1::SnapshotSequences {
1835                    snapshot_seqs,
1836                    sst_min_sequences,
1837                },
1838            ),
1839            explain: None,
1840        }
1841    }
1842}
1843
1844impl From<QueryContext> for FlowQueryContext {
1845    fn from(ctx: QueryContext) -> Self {
1846        Self {
1847            catalog: ctx.current_catalog,
1848            schema: ctx.current_schema,
1849            timezone: ctx.timezone,
1850            extensions: ctx.extensions,
1851            channel: ctx.channel,
1852            snapshot_seqs: ctx.snapshot_seqs,
1853            sst_min_sequences: ctx.sst_min_sequences,
1854        }
1855    }
1856}
1857
1858impl From<FlowQueryContext> for QueryContext {
1859    fn from(flow_ctx: FlowQueryContext) -> Self {
1860        Self {
1861            current_catalog: flow_ctx.catalog,
1862            current_schema: flow_ctx.schema,
1863            timezone: flow_ctx.timezone,
1864            extensions: flow_ctx.extensions,
1865            channel: flow_ctx.channel,
1866            snapshot_seqs: flow_ctx.snapshot_seqs,
1867            sst_min_sequences: flow_ctx.sst_min_sequences,
1868        }
1869    }
1870}
1871
1872impl From<FlowQueryContext> for PbQueryContext {
1873    fn from(flow_ctx: FlowQueryContext) -> Self {
1874        let query_ctx: QueryContext = flow_ctx.into();
1875        query_ctx.into()
1876    }
1877}
1878
1879#[cfg(test)]
1880mod tests {
1881    use std::sync::Arc;
1882
1883    use api::v1::{AlterTableExpr, ColumnDef, CreateTableExpr, SemanticType};
1884    use datatypes::schema::{ColumnSchema, Schema, SchemaBuilder};
1885    use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
1886    use store_api::storage::ConcreteDataType;
1887    use table::metadata::{TableInfo, TableMeta, TableType};
1888    use table::test_util::table_info::test_table_info;
1889
1890    use super::{AlterTableTask, CreateTableTask, *};
1891
1892    #[test]
1893    fn test_ddl_timeout_secs() {
1894        assert_eq!(ddl_timeout_secs(Duration::ZERO), 0);
1895        assert_eq!(ddl_timeout_secs(Duration::from_nanos(1)), 1);
1896        assert_eq!(ddl_timeout_secs(Duration::from_secs(1)), 1);
1897        assert_eq!(ddl_timeout_secs(Duration::from_millis(1500)), 2);
1898        assert_eq!(
1899            ddl_timeout_secs(Duration::from_secs(u32::MAX as u64 + 1)),
1900            u32::MAX
1901        );
1902    }
1903
1904    #[test]
1905    fn test_basic_ser_de_create_table_task() {
1906        let schema = SchemaBuilder::default().build().unwrap();
1907        let table_info = test_table_info(1025, "foo", "bar", "baz", Arc::new(schema));
1908        let task = CreateTableTask::new(CreateTableExpr::default(), Vec::new(), table_info);
1909
1910        let output = serde_json::to_vec(&task).unwrap();
1911
1912        let de = serde_json::from_slice(&output).unwrap();
1913        assert_eq!(task, de);
1914    }
1915
1916    #[test]
1917    fn test_basic_ser_de_alter_table_task() {
1918        let task = AlterTableTask {
1919            alter_table: AlterTableExpr::default(),
1920        };
1921
1922        let output = serde_json::to_vec(&task).unwrap();
1923
1924        let de = serde_json::from_slice(&output).unwrap();
1925        assert_eq!(task, de);
1926    }
1927
1928    #[test]
1929    fn test_undrop_table_task_pb_roundtrip() {
1930        let expected = UndropTableTask { table_id: 1024 };
1931        let request = SubmitDdlTaskRequest::new(DdlTask::UndropTable(expected.clone()));
1932
1933        let pb = PbDdlTaskRequest::try_from(request).unwrap();
1934        let pb_task = pb.task.unwrap();
1935        let de = DdlTask::try_from(pb_task).unwrap();
1936
1937        assert!(matches!(de, DdlTask::UndropTable(task) if task == expected));
1938    }
1939
1940    #[test]
1941    fn test_purge_dropped_table_task_pb_roundtrip() {
1942        let expected = PurgeDroppedTableTask { table_id: 1024 };
1943        let request = SubmitDdlTaskRequest::new(DdlTask::PurgeDroppedTable(expected.clone()));
1944
1945        let pb = PbDdlTaskRequest::try_from(request).unwrap();
1946        let pb_task = pb.task.unwrap();
1947        let de = DdlTask::try_from(pb_task).unwrap();
1948
1949        assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected));
1950    }
1951
1952    #[test]
1953    fn test_undrop_table_task_json_roundtrip() {
1954        let task = UndropTableTask { table_id: 1024 };
1955
1956        let output = serde_json::to_vec(&task).unwrap();
1957
1958        let de = serde_json::from_slice(&output).unwrap();
1959        assert_eq!(task, de);
1960    }
1961
1962    #[test]
1963    fn test_purge_dropped_table_task_json_roundtrip() {
1964        let task = PurgeDroppedTableTask { table_id: 1024 };
1965
1966        let output = serde_json::to_vec(&task).unwrap();
1967
1968        let de = serde_json::from_slice(&output).unwrap();
1969        assert_eq!(task, de);
1970    }
1971
1972    #[test]
1973    fn test_sort_columns() {
1974        // construct RawSchema
1975        let schema = Arc::new(Schema::new(vec![
1976            ColumnSchema::new(
1977                "column3".to_string(),
1978                ConcreteDataType::string_datatype(),
1979                true,
1980            ),
1981            ColumnSchema::new(
1982                "column1".to_string(),
1983                ConcreteDataType::timestamp_millisecond_datatype(),
1984                false,
1985            )
1986            .with_time_index(true),
1987            ColumnSchema::new(
1988                "column2".to_string(),
1989                ConcreteDataType::float64_datatype(),
1990                true,
1991            ),
1992        ]));
1993
1994        // construct RawTableMeta
1995        let meta = TableMeta {
1996            schema,
1997            primary_key_indices: vec![0],
1998            value_indices: vec![2],
1999            engine: METRIC_ENGINE_NAME.to_string(),
2000            next_column_id: 0,
2001            options: Default::default(),
2002            created_on: Default::default(),
2003            updated_on: Default::default(),
2004            partition_key_indices: Default::default(),
2005            column_ids: Default::default(),
2006        };
2007
2008        // construct TableInfo
2009        let raw_table_info = TableInfo {
2010            ident: Default::default(),
2011            meta,
2012            name: Default::default(),
2013            desc: Default::default(),
2014            catalog_name: Default::default(),
2015            schema_name: Default::default(),
2016            table_type: TableType::Base,
2017        };
2018
2019        // construct create table expr
2020        let create_table_expr = CreateTableExpr {
2021            column_defs: vec![
2022                ColumnDef {
2023                    name: "column3".to_string(),
2024                    semantic_type: SemanticType::Tag as i32,
2025                    ..Default::default()
2026                },
2027                ColumnDef {
2028                    name: "column1".to_string(),
2029                    semantic_type: SemanticType::Timestamp as i32,
2030                    ..Default::default()
2031                },
2032                ColumnDef {
2033                    name: "column2".to_string(),
2034                    semantic_type: SemanticType::Field as i32,
2035                    ..Default::default()
2036                },
2037            ],
2038            primary_keys: vec!["column3".to_string()],
2039            ..Default::default()
2040        };
2041
2042        let mut create_table_task =
2043            CreateTableTask::new(create_table_expr, Vec::new(), raw_table_info);
2044
2045        // Call the sort_columns method
2046        create_table_task.sort_columns();
2047
2048        // Assert that the columns are sorted correctly
2049        assert_eq!(
2050            create_table_task.create_table.column_defs[0].name,
2051            "column1".to_string()
2052        );
2053        assert_eq!(
2054            create_table_task.create_table.column_defs[1].name,
2055            "column2".to_string()
2056        );
2057        assert_eq!(
2058            create_table_task.create_table.column_defs[2].name,
2059            "column3".to_string()
2060        );
2061
2062        // Assert that the table_info is updated correctly
2063        assert_eq!(
2064            create_table_task.table_info.meta.schema.timestamp_index(),
2065            Some(0)
2066        );
2067        assert_eq!(
2068            create_table_task.table_info.meta.primary_key_indices,
2069            vec![2]
2070        );
2071        assert_eq!(create_table_task.table_info.meta.value_indices, vec![0, 1]);
2072    }
2073
2074    #[test]
2075    fn test_flow_query_context_conversion_from_query_context() {
2076        use std::collections::HashMap;
2077        let mut extensions = HashMap::new();
2078        extensions.insert("key1".to_string(), "value1".to_string());
2079        extensions.insert("key2".to_string(), "value2".to_string());
2080
2081        let query_ctx = QueryContext {
2082            current_catalog: "test_catalog".to_string(),
2083            current_schema: "test_schema".to_string(),
2084            timezone: "UTC".to_string(),
2085            extensions,
2086            channel: 5,
2087            snapshot_seqs: HashMap::from([(10, 100)]),
2088            sst_min_sequences: HashMap::from([(10, 90)]),
2089        };
2090
2091        let flow_ctx: FlowQueryContext = query_ctx.into();
2092
2093        assert_eq!(flow_ctx.catalog, "test_catalog");
2094        assert_eq!(flow_ctx.schema, "test_schema");
2095        assert_eq!(flow_ctx.timezone, "UTC");
2096        assert_eq!(flow_ctx.channel, 5);
2097        assert_eq!(flow_ctx.snapshot_seqs, HashMap::from([(10, 100)]));
2098        assert_eq!(flow_ctx.sst_min_sequences, HashMap::from([(10, 90)]));
2099    }
2100
2101    #[test]
2102    fn test_flow_query_context_conversion_to_query_context() {
2103        let flow_ctx = FlowQueryContext {
2104            catalog: "prod_catalog".to_string(),
2105            schema: "public".to_string(),
2106            timezone: "America/New_York".to_string(),
2107            extensions: HashMap::from([("k".to_string(), "v".to_string())]),
2108            channel: 7,
2109            snapshot_seqs: HashMap::from([(11, 111)]),
2110            sst_min_sequences: HashMap::from([(11, 101)]),
2111        };
2112
2113        let query_ctx: QueryContext = flow_ctx.clone().into();
2114
2115        assert_eq!(query_ctx.current_catalog, "prod_catalog");
2116        assert_eq!(query_ctx.current_schema, "public");
2117        assert_eq!(query_ctx.timezone, "America/New_York");
2118        assert_eq!(
2119            query_ctx.extensions,
2120            HashMap::from([("k".to_string(), "v".to_string())])
2121        );
2122        assert_eq!(query_ctx.channel, 7);
2123        assert_eq!(query_ctx.snapshot_seqs, HashMap::from([(11, 111)]));
2124        assert_eq!(query_ctx.sst_min_sequences, HashMap::from([(11, 101)]));
2125
2126        // Test roundtrip conversion
2127        let flow_ctx_roundtrip: FlowQueryContext = query_ctx.into();
2128        assert_eq!(flow_ctx, flow_ctx_roundtrip);
2129    }
2130
2131    #[test]
2132    fn test_flow_query_context_serialization() {
2133        let flow_ctx = FlowQueryContext {
2134            catalog: "test_catalog".to_string(),
2135            schema: "test_schema".to_string(),
2136            timezone: "UTC".to_string(),
2137            extensions: HashMap::new(),
2138            channel: 0,
2139            snapshot_seqs: HashMap::new(),
2140            sst_min_sequences: HashMap::new(),
2141        };
2142
2143        let serialized = serde_json::to_string(&flow_ctx).unwrap();
2144        let deserialized: FlowQueryContext = serde_json::from_str(&serialized).unwrap();
2145
2146        assert_eq!(flow_ctx, deserialized);
2147
2148        // Verify JSON structure
2149        let json_value: serde_json::Value = serde_json::from_str(&serialized).unwrap();
2150        assert_eq!(json_value["catalog"], "test_catalog");
2151        assert_eq!(json_value["schema"], "test_schema");
2152        assert_eq!(json_value["timezone"], "UTC");
2153    }
2154
2155    #[test]
2156    fn test_flow_query_context_conversion_to_pb() {
2157        let flow_ctx = FlowQueryContext {
2158            catalog: "pb_catalog".to_string(),
2159            schema: "pb_schema".to_string(),
2160            timezone: "Asia/Tokyo".to_string(),
2161            extensions: HashMap::from([("x".to_string(), "y".to_string())]),
2162            channel: 6,
2163            snapshot_seqs: HashMap::from([(3, 30)]),
2164            sst_min_sequences: HashMap::from([(3, 21)]),
2165        };
2166
2167        let pb_ctx: PbQueryContext = flow_ctx.into();
2168
2169        assert_eq!(pb_ctx.current_catalog, "pb_catalog");
2170        assert_eq!(pb_ctx.current_schema, "pb_schema");
2171        assert_eq!(pb_ctx.timezone, "Asia/Tokyo");
2172        assert_eq!(
2173            pb_ctx.extensions,
2174            HashMap::from([("x".to_string(), "y".to_string())])
2175        );
2176        assert_eq!(pb_ctx.channel, 6);
2177        assert_eq!(
2178            pb_ctx.snapshot_seqs,
2179            Some(api::v1::SnapshotSequences {
2180                snapshot_seqs: HashMap::from([(3, 30)]),
2181                sst_min_sequences: HashMap::from([(3, 21)]),
2182            })
2183        );
2184        assert!(pb_ctx.explain.is_none());
2185    }
2186
2187    #[test]
2188    fn test_pb_query_context_roundtrip_with_snapshot_sequences() {
2189        let pb = PbQueryContext {
2190            current_catalog: "c1".to_string(),
2191            current_schema: "s1".to_string(),
2192            timezone: "UTC".to_string(),
2193            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
2194            channel: 3,
2195            snapshot_seqs: Some(api::v1::SnapshotSequences {
2196                snapshot_seqs: HashMap::from([(1, 100)]),
2197                sst_min_sequences: HashMap::from([(1, 90)]),
2198            }),
2199            explain: None,
2200        };
2201
2202        let query_ctx: QueryContext = pb.clone().into();
2203        let pb_roundtrip: PbQueryContext = query_ctx.into();
2204
2205        assert_eq!(pb_roundtrip.current_catalog, pb.current_catalog);
2206        assert_eq!(pb_roundtrip.current_schema, pb.current_schema);
2207        assert_eq!(pb_roundtrip.timezone, pb.timezone);
2208        assert_eq!(pb_roundtrip.extensions, pb.extensions);
2209        assert_eq!(pb_roundtrip.channel, pb.channel);
2210        assert_eq!(pb_roundtrip.snapshot_seqs, pb.snapshot_seqs);
2211    }
2212
2213    #[test]
2214    fn test_trigger_reason_deserializes_unknown_value() {
2215        let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap();
2216        assert_eq!(TriggerReason::Unknown, reason);
2217    }
2218}