common_meta/rpc/

ddl.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature = "enterprise")]
pub mod trigger;

use std::collections::{HashMap, HashSet};
use std::result;
use std::time::Duration;

use api::helper::{from_pb_time_ranges, to_pb_time_ranges};
use api::v1::alter_database_expr::Kind as PbAlterDatabaseKind;
use api::v1::meta::ddl_task_request::Task;
use api::v1::meta::{
    AlterDatabaseTask as PbAlterDatabaseTask, AlterTableTask as PbAlterTableTask,
    AlterTableTasks as PbAlterTableTasks, CommentOnTask as PbCommentOnTask,
    CreateDatabaseTask as PbCreateDatabaseTask, CreateFlowTask as PbCreateFlowTask,
    CreateTableTask as PbCreateTableTask, CreateTableTasks as PbCreateTableTasks,
    CreateViewTask as PbCreateViewTask, DdlTaskRequest as PbDdlTaskRequest,
    DdlTaskResponse as PbDdlTaskResponse, DropDatabaseTask as PbDropDatabaseTask,
    DropFlowTask as PbDropFlowTask, DropTableTask as PbDropTableTask,
    DropTableTasks as PbDropTableTasks, DropViewTask as PbDropViewTask, Partition, ProcedureId,
    PurgeDroppedTableTask as PbPurgeDroppedTableTask, TruncateTableTask as PbTruncateTableTask,
    UndropTableTask as PbUndropTableTask,
};
use api::v1::{
    AlterDatabaseExpr, AlterTableExpr, CommentObjectType as PbCommentObjectType, CommentOnExpr,
    CreateDatabaseExpr, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropDatabaseExpr,
    DropFlowExpr, DropTableExpr, DropViewExpr, EvalInterval, ExpireAfter, Option as PbOption,
    QueryContext as PbQueryContext, TruncateTableExpr,
};
use base64::Engine as _;
use base64::engine::general_purpose;
use common_catalog::{format_full_flow_name, format_full_table_name};
use common_error::ext::BoxedError;
use common_time::{DatabaseTimeToLive, Timestamp};
use prost::Message;
use serde::{Deserialize, Serialize};
use serde_with::{DefaultOnNull, serde_as};
use snafu::{OptionExt, ResultExt};
use table::metadata::{TableId, TableInfo};
use table::requests::validate_database_option;
use table::table_name::TableName;
use table::table_reference::TableReference;

use crate::error::{
    self, ConvertTimeRangesSnafu, ExternalSnafu, InvalidSetDatabaseOptionSnafu,
    InvalidUnsetDatabaseOptionSnafu, Result,
};
use crate::flow_name::FlowName;
use crate::instruction::CacheIdent;
use crate::key::FlowId;
use crate::key::flow::flow_name::FlowNameManager;
use crate::key::table_name::{TableNameKey, TableNameManager};

/// Reserved query-context extension key for the frontend peer address that submitted a DDL request.
pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr";

/// Reserved query-context extension key for the authenticated database creator.
pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator";
/// Internal gRPC metadata key for the authenticated database creator.
pub const CREATE_DATABASE_CREATOR_METADATA_KEY: &str =
    "x-greptime-internal-create-database-creator-bin";

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CreatorGrantIntent {
    pub username: String,
    pub created_at_ns: i64,
}

/// DDL tasks
#[derive(Debug, Clone)]
pub enum DdlTask {
    CreateTable(CreateTableTask),
    DropTable(DropTableTask),
    UndropTable(UndropTableTask),
    PurgeDroppedTable(PurgeDroppedTableTask),
    AlterTable(AlterTableTask),
    TruncateTable(TruncateTableTask),
    CreateLogicalTables(Vec<CreateTableTask>),
    DropLogicalTables(Vec<DropTableTask>),
    AlterLogicalTables(Vec<AlterTableTask>),
    CreateDatabase(CreateDatabaseTask),
    DropDatabase(DropDatabaseTask),
    AlterDatabase(AlterDatabaseTask),
    CreateFlow(CreateFlowTask),
    DropFlow(DropFlowTask),
    #[cfg(feature = "enterprise")]
    DropTrigger(trigger::DropTriggerTask),
    CreateView(CreateViewTask),
    DropView(DropViewTask),
    #[cfg(feature = "enterprise")]
    CreateTrigger(trigger::CreateTriggerTask),
    CommentOn(CommentOnTask),
}

impl DdlTask {
    /// Creates a [`DdlTask`] to create a flow.
    pub fn new_create_flow(expr: CreateFlowTask) -> Self {
        DdlTask::CreateFlow(expr)
    }

    /// Creates a [`DdlTask`] to drop a flow.
    pub fn new_drop_flow(expr: DropFlowTask) -> Self {
        DdlTask::DropFlow(expr)
    }

    /// Creates a [`DdlTask`] to drop a view.
    pub fn new_drop_view(expr: DropViewTask) -> Self {
        DdlTask::DropView(expr)
    }

    /// Creates a [`DdlTask`] to create a table.
    pub fn new_create_table(
        expr: CreateTableExpr,
        partitions: Vec<Partition>,
        table_info: TableInfo,
    ) -> Self {
        DdlTask::CreateTable(CreateTableTask::new(expr, partitions, table_info))
    }

    /// Creates a [`DdlTask`] to create several logical tables.
    pub fn new_create_logical_tables(table_data: Vec<(CreateTableExpr, TableInfo)>) -> Self {
        DdlTask::CreateLogicalTables(
            table_data
                .into_iter()
                .map(|(expr, table_info)| CreateTableTask::new(expr, Vec::new(), table_info))
                .collect(),
        )
    }

    /// Creates a [`DdlTask`] to alter several logical tables.
    pub fn new_alter_logical_tables(table_data: Vec<AlterTableExpr>) -> Self {
        DdlTask::AlterLogicalTables(
            table_data
                .into_iter()
                .map(|alter_table| AlterTableTask { alter_table })
                .collect(),
        )
    }

    /// Creates a [`DdlTask`] to drop a table.
    pub fn new_drop_table(
        catalog: String,
        schema: String,
        table: String,
        table_id: TableId,
        drop_if_exists: bool,
    ) -> Self {
        DdlTask::DropTable(DropTableTask {
            catalog,
            schema,
            table,
            table_id,
            drop_if_exists,
        })
    }

    /// Creates a [`DdlTask`] to undrop a table.
    pub fn new_undrop_table(table_id: TableId) -> Self {
        DdlTask::UndropTable(UndropTableTask { table_id })
    }

    /// Creates a [`DdlTask`] to purge a dropped table.
    pub fn new_purge_dropped_table(table_id: TableId) -> Self {
        DdlTask::PurgeDroppedTable(PurgeDroppedTableTask { table_id })
    }

    /// Creates a [`DdlTask`] to create a database.
    pub fn new_create_database(
        catalog: String,
        schema: String,
        create_if_not_exists: bool,
        options: HashMap<String, String>,
        creator: Option<CreatorGrantIntent>,
    ) -> Self {
        DdlTask::CreateDatabase(CreateDatabaseTask {
            catalog,
            schema,
            create_if_not_exists,
            options,
            creator,
        })
    }

    /// Creates a [`DdlTask`] to drop a database.
    pub fn new_drop_database(catalog: String, schema: String, drop_if_exists: bool) -> Self {
        DdlTask::DropDatabase(DropDatabaseTask {
            catalog,
            schema,
            drop_if_exists,
        })
    }

    /// Creates a [`DdlTask`] to alter a database.
    pub fn new_alter_database(alter_expr: AlterDatabaseExpr) -> Self {
        DdlTask::AlterDatabase(AlterDatabaseTask { alter_expr })
    }

    /// Creates a [`DdlTask`] to alter a table.
    pub fn new_alter_table(alter_table: AlterTableExpr) -> Self {
        DdlTask::AlterTable(AlterTableTask { alter_table })
    }

    /// Creates a [`DdlTask`] to truncate a table.
    pub fn new_truncate_table(
        catalog: String,
        schema: String,
        table: String,
        table_id: TableId,
        time_ranges: Vec<(Timestamp, Timestamp)>,
    ) -> Self {
        DdlTask::TruncateTable(TruncateTableTask {
            catalog,
            schema,
            table,
            table_id,
            time_ranges,
        })
    }

    /// Creates a [`DdlTask`] to create a view.
    pub fn new_create_view(create_view: CreateViewExpr, view_info: TableInfo) -> Self {
        DdlTask::CreateView(CreateViewTask {
            create_view,
            view_info,
        })
    }

    /// Creates a [`DdlTask`] to comment on a table, column, or flow.
    pub fn new_comment_on(task: CommentOnTask) -> Self {
        DdlTask::CommentOn(task)
    }
}

impl TryFrom<Task> for DdlTask {
    type Error = error::Error;
    fn try_from(task: Task) -> Result<Self> {
        match task {
            Task::CreateTableTask(create_table) => {
                Ok(DdlTask::CreateTable(create_table.try_into()?))
            }
            Task::DropTableTask(drop_table) => Ok(DdlTask::DropTable(drop_table.try_into()?)),
            Task::UndropTableTask(undrop_table) => {
                Ok(DdlTask::UndropTable(undrop_table.try_into()?))
            }
            Task::PurgeDroppedTableTask(purge_dropped_table) => {
                Ok(DdlTask::PurgeDroppedTable(purge_dropped_table.try_into()?))
            }
            Task::AlterTableTask(alter_table) => Ok(DdlTask::AlterTable(alter_table.try_into()?)),
            Task::TruncateTableTask(truncate_table) => {
                Ok(DdlTask::TruncateTable(truncate_table.try_into()?))
            }
            Task::CreateTableTasks(create_tables) => {
                let tasks = create_tables
                    .tasks
                    .into_iter()
                    .map(|task| task.try_into())
                    .collect::<Result<Vec<_>>>()?;

                Ok(DdlTask::CreateLogicalTables(tasks))
            }
            Task::DropTableTasks(drop_tables) => {
                let tasks = drop_tables
                    .tasks
                    .into_iter()
                    .map(|task| task.try_into())
                    .collect::<Result<Vec<_>>>()?;

                Ok(DdlTask::DropLogicalTables(tasks))
            }
            Task::AlterTableTasks(alter_tables) => {
                let tasks = alter_tables
                    .tasks
                    .into_iter()
                    .map(|task| task.try_into())
                    .collect::<Result<Vec<_>>>()?;

                Ok(DdlTask::AlterLogicalTables(tasks))
            }
            Task::CreateDatabaseTask(create_database) => {
                Ok(DdlTask::CreateDatabase(create_database.try_into()?))
            }
            Task::DropDatabaseTask(drop_database) => {
                Ok(DdlTask::DropDatabase(drop_database.try_into()?))
            }
            Task::AlterDatabaseTask(alter_database) => {
                Ok(DdlTask::AlterDatabase(alter_database.try_into()?))
            }
            Task::CreateFlowTask(create_flow) => Ok(DdlTask::CreateFlow(create_flow.try_into()?)),
            Task::DropFlowTask(drop_flow) => Ok(DdlTask::DropFlow(drop_flow.try_into()?)),
            Task::CreateViewTask(create_view) => Ok(DdlTask::CreateView(create_view.try_into()?)),
            Task::DropViewTask(drop_view) => Ok(DdlTask::DropView(drop_view.try_into()?)),
            Task::CreateTriggerTask(create_trigger) => {
                #[cfg(feature = "enterprise")]
                return Ok(DdlTask::CreateTrigger(create_trigger.try_into()?));
                #[cfg(not(feature = "enterprise"))]
                {
                    let _ = create_trigger;
                    crate::error::UnsupportedSnafu {
                        operation: "create trigger",
                    }
                    .fail()
                }
            }
            Task::DropTriggerTask(drop_trigger) => {
                #[cfg(feature = "enterprise")]
                return Ok(DdlTask::DropTrigger(drop_trigger.try_into()?));
                #[cfg(not(feature = "enterprise"))]
                {
                    let _ = drop_trigger;
                    crate::error::UnsupportedSnafu {
                        operation: "drop trigger",
                    }
                    .fail()
                }
            }
            Task::CommentOnTask(comment_on) => Ok(DdlTask::CommentOn(comment_on.try_into()?)),
        }
    }
}

#[derive(Clone)]
pub struct SubmitDdlTaskRequest {
    pub query_context: QueryContext,
    pub wait: bool,
    pub timeout: Duration,
    pub task: DdlTask,
}

impl SubmitDdlTaskRequest {
    /// The default constructor for [`SubmitDdlTaskRequest`].
    pub fn new(query_context: QueryContext, task: DdlTask) -> Self {
        Self {
            query_context,
            wait: Self::default_wait(),
            timeout: Self::default_timeout(),
            task,
        }
    }

    /// The default timeout for a DDL task.
    pub fn default_timeout() -> Duration {
        Duration::from_secs(60)
    }

    /// The default wait for a DDL task.
    pub fn default_wait() -> bool {
        true
    }
}

fn ddl_timeout_secs(timeout: Duration) -> u32 {
    timeout
        .as_nanos()
        .div_ceil(Duration::from_secs(1).as_nanos())
        .try_into()
        .unwrap_or(u32::MAX)
}

impl TryFrom<SubmitDdlTaskRequest> for PbDdlTaskRequest {
    type Error = error::Error;

    fn try_from(request: SubmitDdlTaskRequest) -> Result<Self> {
        let SubmitDdlTaskRequest {
            mut query_context,
            wait,
            timeout,
            task,
        } = request;

        query_context
            .extensions
            .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY);
        if let DdlTask::CreateDatabase(CreateDatabaseTask {
            creator: Some(creator),
            ..
        }) = &task
        {
            query_context.extensions.insert(
                CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(),
                serde_json::to_string(creator).context(error::SerdeJsonSnafu)?,
            );
        }

        let task = match task {
            DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?),
            DdlTask::DropTable(task) => Task::DropTableTask(task.into()),
            DdlTask::UndropTable(task) => Task::UndropTableTask(task.into()),
            DdlTask::PurgeDroppedTable(task) => Task::PurgeDroppedTableTask(task.into()),
            DdlTask::AlterTable(task) => Task::AlterTableTask(task.try_into()?),
            DdlTask::TruncateTable(task) => Task::TruncateTableTask(task.try_into()?),
            DdlTask::CreateLogicalTables(tasks) => {
                let tasks = tasks
                    .into_iter()
                    .map(|task| task.try_into())
                    .collect::<Result<Vec<_>>>()?;

                Task::CreateTableTasks(PbCreateTableTasks { tasks })
            }
            DdlTask::DropLogicalTables(tasks) => {
                let tasks = tasks
                    .into_iter()
                    .map(|task| task.into())
                    .collect::<Vec<_>>();

                Task::DropTableTasks(PbDropTableTasks { tasks })
            }
            DdlTask::AlterLogicalTables(tasks) => {
                let tasks = tasks
                    .into_iter()
                    .map(|task| task.try_into())
                    .collect::<Result<Vec<_>>>()?;

                Task::AlterTableTasks(PbAlterTableTasks { tasks })
            }
            DdlTask::CreateDatabase(task) => Task::CreateDatabaseTask(task.try_into()?),
            DdlTask::DropDatabase(task) => Task::DropDatabaseTask(task.try_into()?),
            DdlTask::AlterDatabase(task) => Task::AlterDatabaseTask(task.try_into()?),
            DdlTask::CreateFlow(task) => Task::CreateFlowTask(task.into()),
            DdlTask::DropFlow(task) => Task::DropFlowTask(task.into()),
            DdlTask::CreateView(task) => Task::CreateViewTask(task.try_into()?),
            DdlTask::DropView(task) => Task::DropViewTask(task.into()),
            #[cfg(feature = "enterprise")]
            DdlTask::CreateTrigger(task) => Task::CreateTriggerTask(task.try_into()?),
            #[cfg(feature = "enterprise")]
            DdlTask::DropTrigger(task) => Task::DropTriggerTask(task.into()),
            DdlTask::CommentOn(task) => Task::CommentOnTask(task.into()),
        };

        Ok(Self {
            header: None,
            query_context: Some(query_context.into()),
            timeout_secs: ddl_timeout_secs(timeout),
            wait,
            task: Some(task),
        })
    }
}

#[derive(Debug, Default)]
pub struct SubmitDdlTaskResponse {
    pub key: Vec<u8>,
    // `table_id`s for `CREATE TABLE` or `CREATE LOGICAL TABLES` task.
    pub table_ids: Vec<TableId>,
}

impl TryFrom<PbDdlTaskResponse> for SubmitDdlTaskResponse {
    type Error = error::Error;

    fn try_from(resp: PbDdlTaskResponse) -> Result<Self> {
        let table_ids = resp.table_ids.into_iter().map(|t| t.id).collect();
        Ok(Self {
            key: resp.pid.map(|pid| pid.key).unwrap_or_default(),
            table_ids,
        })
    }
}

impl From<SubmitDdlTaskResponse> for PbDdlTaskResponse {
    fn from(val: SubmitDdlTaskResponse) -> Self {
        Self {
            pid: Some(ProcedureId { key: val.key }),
            table_ids: val
                .table_ids
                .into_iter()
                .map(|id| api::v1::TableId { id })
                .collect(),
            ..Default::default()
        }
    }
}

/// A `CREATE VIEW` task.
#[derive(Debug, PartialEq, Clone)]
pub struct CreateViewTask {
    pub create_view: CreateViewExpr,
    pub view_info: TableInfo,
}

impl CreateViewTask {
    /// Returns the [`TableReference`] of view.
    pub fn table_ref(&self) -> TableReference<'_> {
        TableReference {
            catalog: &self.create_view.catalog_name,
            schema: &self.create_view.schema_name,
            table: &self.create_view.view_name,
        }
    }

    /// Returns the encoded logical plan
    pub fn raw_logical_plan(&self) -> &Vec<u8> {
        &self.create_view.logical_plan
    }

    /// Returns the view definition in SQL
    pub fn view_definition(&self) -> &str {
        &self.create_view.definition
    }

    /// Returns the resolved table names in view's logical plan
    pub fn table_names(&self) -> HashSet<TableName> {
        self.create_view
            .table_names
            .iter()
            .map(|t| t.clone().into())
            .collect()
    }

    /// Returns the view's columns
    pub fn columns(&self) -> &Vec<String> {
        &self.create_view.columns
    }

    /// Returns the original logical plan's columns
    pub fn plan_columns(&self) -> &Vec<String> {
        &self.create_view.plan_columns
    }
}

impl TryFrom<PbCreateViewTask> for CreateViewTask {
    type Error = error::Error;

    fn try_from(pb: PbCreateViewTask) -> Result<Self> {
        let view_info = serde_json::from_slice(&pb.view_info).context(error::SerdeJsonSnafu)?;

        Ok(CreateViewTask {
            create_view: pb.create_view.context(error::InvalidProtoMsgSnafu {
                err_msg: "expected create view",
            })?,
            view_info,
        })
    }
}

impl TryFrom<CreateViewTask> for PbCreateViewTask {
    type Error = error::Error;

    fn try_from(task: CreateViewTask) -> Result<PbCreateViewTask> {
        Ok(PbCreateViewTask {
            create_view: Some(task.create_view),
            view_info: serde_json::to_vec(&task.view_info).context(error::SerdeJsonSnafu)?,
        })
    }
}

impl Serialize for CreateViewTask {
    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let view_info = serde_json::to_vec(&self.view_info)
            .map_err(|err| serde::ser::Error::custom(err.to_string()))?;

        let pb = PbCreateViewTask {
            create_view: Some(self.create_view.clone()),
            view_info,
        };
        let buf = pb.encode_to_vec();
        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
        serializer.serialize_str(&encoded)
    }
}

impl<'de> Deserialize<'de> for CreateViewTask {
    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let encoded = String::deserialize(deserializer)?;
        let buf = general_purpose::STANDARD_NO_PAD
            .decode(encoded)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
        let expr: PbCreateViewTask = PbCreateViewTask::decode(&*buf)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        let expr = CreateViewTask::try_from(expr)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        Ok(expr)
    }
}

/// A `DROP VIEW` task.
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct DropViewTask {
    pub catalog: String,
    pub schema: String,
    pub view: String,
    pub view_id: TableId,
    pub drop_if_exists: bool,
}

impl DropViewTask {
    /// Returns the [`TableReference`] of view.
    pub fn table_ref(&self) -> TableReference<'_> {
        TableReference {
            catalog: &self.catalog,
            schema: &self.schema,
            table: &self.view,
        }
    }
}

impl TryFrom<PbDropViewTask> for DropViewTask {
    type Error = error::Error;

    fn try_from(pb: PbDropViewTask) -> Result<Self> {
        let expr = pb.drop_view.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected drop view",
        })?;

        Ok(DropViewTask {
            catalog: expr.catalog_name,
            schema: expr.schema_name,
            view: expr.view_name,
            view_id: expr
                .view_id
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected view_id",
                })?
                .id,
            drop_if_exists: expr.drop_if_exists,
        })
    }
}

impl From<DropViewTask> for PbDropViewTask {
    fn from(task: DropViewTask) -> Self {
        PbDropViewTask {
            drop_view: Some(DropViewExpr {
                catalog_name: task.catalog,
                schema_name: task.schema,
                view_name: task.view,
                view_id: Some(api::v1::TableId { id: task.view_id }),
                drop_if_exists: task.drop_if_exists,
            }),
        }
    }
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct DropTableTask {
    pub catalog: String,
    pub schema: String,
    pub table: String,
    pub table_id: TableId,
    #[serde(default)]
    pub drop_if_exists: bool,
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct UndropTableTask {
    pub table_id: TableId,
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct PurgeDroppedTableTask {
    pub table_id: TableId,
}

impl TryFrom<PbUndropTableTask> for UndropTableTask {
    type Error = error::Error;

    fn try_from(pb: PbUndropTableTask) -> Result<Self> {
        Ok(Self {
            table_id: pb
                .table_id
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected table_id",
                })?
                .id,
        })
    }
}

impl From<UndropTableTask> for PbUndropTableTask {
    fn from(task: UndropTableTask) -> Self {
        Self {
            table_id: Some(api::v1::TableId { id: task.table_id }),
        }
    }
}

impl TryFrom<PbPurgeDroppedTableTask> for PurgeDroppedTableTask {
    type Error = error::Error;

    fn try_from(pb: PbPurgeDroppedTableTask) -> Result<Self> {
        Ok(Self {
            table_id: pb
                .table_id
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected table_id",
                })?
                .id,
        })
    }
}

impl From<PurgeDroppedTableTask> for PbPurgeDroppedTableTask {
    fn from(task: PurgeDroppedTableTask) -> Self {
        Self {
            table_id: Some(api::v1::TableId { id: task.table_id }),
        }
    }
}

impl DropTableTask {
    pub fn table_ref(&self) -> TableReference<'_> {
        TableReference {
            catalog: &self.catalog,
            schema: &self.schema,
            table: &self.table,
        }
    }

    pub fn table_name(&self) -> TableName {
        TableName {
            catalog_name: self.catalog.clone(),
            schema_name: self.schema.clone(),
            table_name: self.table.clone(),
        }
    }
}

impl TryFrom<PbDropTableTask> for DropTableTask {
    type Error = error::Error;

    fn try_from(pb: PbDropTableTask) -> Result<Self> {
        let drop_table = pb.drop_table.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected drop table",
        })?;

        Ok(Self {
            catalog: drop_table.catalog_name,
            schema: drop_table.schema_name,
            table: drop_table.table_name,
            table_id: drop_table
                .table_id
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected table_id",
                })?
                .id,
            drop_if_exists: drop_table.drop_if_exists,
        })
    }
}

impl From<DropTableTask> for PbDropTableTask {
    fn from(task: DropTableTask) -> Self {
        PbDropTableTask {
            drop_table: Some(DropTableExpr {
                catalog_name: task.catalog,
                schema_name: task.schema,
                table_name: task.table,
                table_id: Some(api::v1::TableId { id: task.table_id }),
                drop_if_exists: task.drop_if_exists,
            }),
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct CreateTableTask {
    pub create_table: CreateTableExpr,
    pub partitions: Vec<Partition>,
    pub table_info: TableInfo,
}

impl TryFrom<PbCreateTableTask> for CreateTableTask {
    type Error = error::Error;

    fn try_from(pb: PbCreateTableTask) -> Result<Self> {
        let table_info = serde_json::from_slice(&pb.table_info).context(error::SerdeJsonSnafu)?;

        Ok(CreateTableTask::new(
            pb.create_table.context(error::InvalidProtoMsgSnafu {
                err_msg: "expected create table",
            })?,
            pb.partitions,
            table_info,
        ))
    }
}

impl TryFrom<CreateTableTask> for PbCreateTableTask {
    type Error = error::Error;

    fn try_from(task: CreateTableTask) -> Result<Self> {
        Ok(PbCreateTableTask {
            table_info: serde_json::to_vec(&task.table_info).context(error::SerdeJsonSnafu)?,
            create_table: Some(task.create_table),
            partitions: task.partitions,
        })
    }
}

impl CreateTableTask {
    pub fn new(
        expr: CreateTableExpr,
        partitions: Vec<Partition>,
        table_info: TableInfo,
    ) -> CreateTableTask {
        CreateTableTask {
            create_table: expr,
            partitions,
            table_info,
        }
    }

    pub fn table_name(&self) -> TableName {
        let table = &self.create_table;

        TableName {
            catalog_name: table.catalog_name.clone(),
            schema_name: table.schema_name.clone(),
            table_name: table.table_name.clone(),
        }
    }

    pub fn table_ref(&self) -> TableReference<'_> {
        let table = &self.create_table;

        TableReference {
            catalog: &table.catalog_name,
            schema: &table.schema_name,
            table: &table.table_name,
        }
    }

    /// Sets the `table_info`'s table_id.
    pub fn set_table_id(&mut self, table_id: TableId) {
        self.table_info.ident.table_id = table_id;
    }

    /// Sort the columns in [CreateTableExpr] and [TableInfo].
    ///
    /// This function won't do any check or verification. Caller should
    /// ensure this task is valid.
    pub fn sort_columns(&mut self) {
        // sort create table expr
        // sort column_defs by name
        self.create_table
            .column_defs
            .sort_unstable_by(|a, b| a.name.cmp(&b.name));

        self.table_info.sort_columns();
    }
}

impl Serialize for CreateTableTask {
    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let table_info = serde_json::to_vec(&self.table_info)
            .map_err(|err| serde::ser::Error::custom(err.to_string()))?;

        let pb = PbCreateTableTask {
            create_table: Some(self.create_table.clone()),
            partitions: self.partitions.clone(),
            table_info,
        };
        let buf = pb.encode_to_vec();
        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
        serializer.serialize_str(&encoded)
    }
}

impl<'de> Deserialize<'de> for CreateTableTask {
    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let encoded = String::deserialize(deserializer)?;
        let buf = general_purpose::STANDARD_NO_PAD
            .decode(encoded)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
        let expr: PbCreateTableTask = PbCreateTableTask::decode(&*buf)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        let expr = CreateTableTask::try_from(expr)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        Ok(expr)
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct AlterTableTask {
    // TODO(CookiePieWw): Replace proto struct with user-defined struct
    pub alter_table: AlterTableExpr,
}

impl AlterTableTask {
    pub fn validate(&self) -> Result<()> {
        self.alter_table
            .kind
            .as_ref()
            .context(error::UnexpectedSnafu {
                err_msg: "'kind' is absent",
            })?;
        Ok(())
    }

    pub fn table_ref(&self) -> TableReference<'_> {
        TableReference {
            catalog: &self.alter_table.catalog_name,
            schema: &self.alter_table.schema_name,
            table: &self.alter_table.table_name,
        }
    }

    pub fn table_name(&self) -> TableName {
        let table = &self.alter_table;

        TableName {
            catalog_name: table.catalog_name.clone(),
            schema_name: table.schema_name.clone(),
            table_name: table.table_name.clone(),
        }
    }
}

impl TryFrom<PbAlterTableTask> for AlterTableTask {
    type Error = error::Error;

    fn try_from(pb: PbAlterTableTask) -> Result<Self> {
        let alter_table = pb.alter_table.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected alter_table",
        })?;

        Ok(AlterTableTask { alter_table })
    }
}

impl TryFrom<AlterTableTask> for PbAlterTableTask {
    type Error = error::Error;

    fn try_from(task: AlterTableTask) -> Result<Self> {
        Ok(PbAlterTableTask {
            alter_table: Some(task.alter_table),
        })
    }
}

impl Serialize for AlterTableTask {
    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let pb = PbAlterTableTask {
            alter_table: Some(self.alter_table.clone()),
        };
        let buf = pb.encode_to_vec();
        let encoded = general_purpose::STANDARD_NO_PAD.encode(buf);
        serializer.serialize_str(&encoded)
    }
}

impl<'de> Deserialize<'de> for AlterTableTask {
    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let encoded = String::deserialize(deserializer)?;
        let buf = general_purpose::STANDARD_NO_PAD
            .decode(encoded)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;
        let expr: PbAlterTableTask = PbAlterTableTask::decode(&*buf)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        let expr = AlterTableTask::try_from(expr)
            .map_err(|err| serde::de::Error::custom(err.to_string()))?;

        Ok(expr)
    }
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct TruncateTableTask {
    pub catalog: String,
    pub schema: String,
    pub table: String,
    pub table_id: TableId,
    pub time_ranges: Vec<(Timestamp, Timestamp)>,
}

impl TruncateTableTask {
    pub fn table_ref(&self) -> TableReference<'_> {
        TableReference {
            catalog: &self.catalog,
            schema: &self.schema,
            table: &self.table,
        }
    }

    pub fn table_name(&self) -> TableName {
        TableName {
            catalog_name: self.catalog.clone(),
            schema_name: self.schema.clone(),
            table_name: self.table.clone(),
        }
    }
}

impl TryFrom<PbTruncateTableTask> for TruncateTableTask {
    type Error = error::Error;

    fn try_from(pb: PbTruncateTableTask) -> Result<Self> {
        let truncate_table = pb.truncate_table.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected truncate table",
        })?;

        Ok(Self {
            catalog: truncate_table.catalog_name,
            schema: truncate_table.schema_name,
            table: truncate_table.table_name,
            table_id: truncate_table
                .table_id
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected table_id",
                })?
                .id,
            time_ranges: truncate_table
                .time_ranges
                .map(from_pb_time_ranges)
                .transpose()
                .map_err(BoxedError::new)
                .context(ExternalSnafu)?
                .unwrap_or_default(),
        })
    }
}

impl TryFrom<TruncateTableTask> for PbTruncateTableTask {
    type Error = error::Error;

    fn try_from(task: TruncateTableTask) -> Result<Self> {
        Ok(PbTruncateTableTask {
            truncate_table: Some(TruncateTableExpr {
                catalog_name: task.catalog,
                schema_name: task.schema,
                table_name: task.table,
                table_id: Some(api::v1::TableId { id: task.table_id }),
                time_ranges: Some(
                    to_pb_time_ranges(&task.time_ranges).context(ConvertTimeRangesSnafu)?,
                ),
            }),
        })
    }
}

#[serde_as]
#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct CreateDatabaseTask {
    pub catalog: String,
    pub schema: String,
    pub create_if_not_exists: bool,
    #[serde_as(deserialize_as = "DefaultOnNull")]
    pub options: HashMap<String, String>,
    #[serde(default)]
    pub creator: Option<CreatorGrantIntent>,
}

impl TryFrom<PbCreateDatabaseTask> for CreateDatabaseTask {
    type Error = error::Error;

    fn try_from(pb: PbCreateDatabaseTask) -> Result<Self> {
        let CreateDatabaseExpr {
            catalog_name,
            schema_name,
            create_if_not_exists,
            options,
        } = pb.create_database.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected create database",
        })?;

        Ok(CreateDatabaseTask {
            catalog: catalog_name,
            schema: schema_name,
            create_if_not_exists,
            options,
            creator: None,
        })
    }
}

impl TryFrom<CreateDatabaseTask> for PbCreateDatabaseTask {
    type Error = error::Error;

    fn try_from(
        CreateDatabaseTask {
            catalog,
            schema,
            create_if_not_exists,
            options,
            creator: _,
        }: CreateDatabaseTask,
    ) -> Result<Self> {
        Ok(PbCreateDatabaseTask {
            create_database: Some(CreateDatabaseExpr {
                catalog_name: catalog,
                schema_name: schema,
                create_if_not_exists,
                options,
            }),
        })
    }
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub struct DropDatabaseTask {
    pub catalog: String,
    pub schema: String,
    pub drop_if_exists: bool,
}

impl TryFrom<PbDropDatabaseTask> for DropDatabaseTask {
    type Error = error::Error;

    fn try_from(pb: PbDropDatabaseTask) -> Result<Self> {
        let DropDatabaseExpr {
            catalog_name,
            schema_name,
            drop_if_exists,
        } = pb.drop_database.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected drop database",
        })?;

        Ok(DropDatabaseTask {
            catalog: catalog_name,
            schema: schema_name,
            drop_if_exists,
        })
    }
}

impl TryFrom<DropDatabaseTask> for PbDropDatabaseTask {
    type Error = error::Error;

    fn try_from(
        DropDatabaseTask {
            catalog,
            schema,
            drop_if_exists,
        }: DropDatabaseTask,
    ) -> Result<Self> {
        Ok(PbDropDatabaseTask {
            drop_database: Some(DropDatabaseExpr {
                catalog_name: catalog,
                schema_name: schema,
                drop_if_exists,
            }),
        })
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct AlterDatabaseTask {
    pub alter_expr: AlterDatabaseExpr,
}

impl TryFrom<AlterDatabaseTask> for PbAlterDatabaseTask {
    type Error = error::Error;

    fn try_from(task: AlterDatabaseTask) -> Result<Self> {
        Ok(PbAlterDatabaseTask {
            task: Some(task.alter_expr),
        })
    }
}

impl TryFrom<PbAlterDatabaseTask> for AlterDatabaseTask {
    type Error = error::Error;

    fn try_from(pb: PbAlterDatabaseTask) -> Result<Self> {
        let alter_expr = pb.task.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected alter database",
        })?;

        Ok(AlterDatabaseTask { alter_expr })
    }
}

impl TryFrom<PbAlterDatabaseKind> for AlterDatabaseKind {
    type Error = error::Error;

    fn try_from(pb: PbAlterDatabaseKind) -> Result<Self> {
        match pb {
            PbAlterDatabaseKind::SetDatabaseOptions(options) => {
                Ok(AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(
                    options
                        .set_database_options
                        .into_iter()
                        .map(SetDatabaseOption::try_from)
                        .collect::<Result<Vec<_>>>()?,
                )))
            }
            PbAlterDatabaseKind::UnsetDatabaseOptions(options) => Ok(
                AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(
                    options
                        .keys
                        .iter()
                        .map(|key| UnsetDatabaseOption::try_from(key.as_str()))
                        .collect::<Result<Vec<_>>>()?,
                )),
            ),
        }
    }
}

const TTL_KEY: &str = "ttl";

impl TryFrom<PbOption> for SetDatabaseOption {
    type Error = error::Error;

    fn try_from(PbOption { key, value }: PbOption) -> Result<Self> {
        let key_lower = key.to_ascii_lowercase();
        match key_lower.as_str() {
            TTL_KEY => {
                let ttl = DatabaseTimeToLive::from_humantime_or_str(&value)
                    .map_err(|_| InvalidSetDatabaseOptionSnafu { key, value }.build())?;

                Ok(SetDatabaseOption::Ttl(ttl))
            }
            _ => {
                if validate_database_option(&key_lower) {
                    Ok(SetDatabaseOption::Other(key_lower, value))
                } else {
                    InvalidSetDatabaseOptionSnafu { key, value }.fail()
                }
            }
        }
    }
}

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub enum SetDatabaseOption {
    Ttl(DatabaseTimeToLive),
    Other(String, String),
}

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub enum UnsetDatabaseOption {
    Ttl,
    Other(String),
}

impl TryFrom<&str> for UnsetDatabaseOption {
    type Error = error::Error;

    fn try_from(key: &str) -> Result<Self> {
        let key_lower = key.to_ascii_lowercase();
        match key_lower.as_str() {
            TTL_KEY => Ok(UnsetDatabaseOption::Ttl),
            _ => {
                if validate_database_option(&key_lower) {
                    Ok(UnsetDatabaseOption::Other(key_lower))
                } else {
                    InvalidUnsetDatabaseOptionSnafu { key }.fail()
                }
            }
        }
    }
}

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct SetDatabaseOptions(pub Vec<SetDatabaseOption>);

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct UnsetDatabaseOptions(pub Vec<UnsetDatabaseOption>);

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub enum AlterDatabaseKind {
    SetDatabaseOptions(SetDatabaseOptions),
    UnsetDatabaseOptions(UnsetDatabaseOptions),
}

impl AlterDatabaseTask {
    pub fn catalog(&self) -> &str {
        &self.alter_expr.catalog_name
    }

    pub fn schema(&self) -> &str {
        &self.alter_expr.catalog_name
    }
}

/// Create flow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateFlowTask {
    pub catalog_name: String,
    pub flow_name: String,
    pub source_table_names: Vec<TableName>,
    pub sink_table_name: TableName,
    pub or_replace: bool,
    pub create_if_not_exists: bool,
    /// Duration in seconds. Data older than this duration will not be used.
    pub expire_after: Option<i64>,
    pub eval_interval_secs: Option<i64>,
    pub comment: String,
    pub sql: String,
    pub flow_options: HashMap<String, String>,
    /// Typed schedule configuration resolved during `on_prepare`.
    /// Not populated from proto; set by the procedure layer after
    /// defaults are resolved.
    #[serde(default)]
    pub eval_schedule: Option<crate::key::flow::flow_info::FlowScheduleConfig>,
}

impl TryFrom<PbCreateFlowTask> for CreateFlowTask {
    type Error = error::Error;

    fn try_from(pb: PbCreateFlowTask) -> Result<Self> {
        let CreateFlowExpr {
            catalog_name,
            flow_name,
            source_table_names,
            sink_table_name,
            or_replace,
            create_if_not_exists,
            expire_after,
            eval_interval,
            comment,
            sql,
            flow_options,
        } = pb.create_flow.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected create_flow",
        })?;

        Ok(CreateFlowTask {
            catalog_name,
            flow_name,
            source_table_names: source_table_names.into_iter().map(Into::into).collect(),
            sink_table_name: sink_table_name
                .context(error::InvalidProtoMsgSnafu {
                    err_msg: "expected sink_table_name",
                })?
                .into(),
            or_replace,
            create_if_not_exists,
            expire_after: expire_after.map(|e| e.value),
            eval_interval_secs: eval_interval.map(|e| e.seconds),
            comment,
            sql,
            flow_options,
            eval_schedule: None,
        })
    }
}

impl From<CreateFlowTask> for PbCreateFlowTask {
    fn from(
        CreateFlowTask {
            catalog_name,
            flow_name,
            source_table_names,
            sink_table_name,
            or_replace,
            create_if_not_exists,
            expire_after,
            eval_interval_secs: eval_interval,
            comment,
            sql,
            flow_options,
            ..
        }: CreateFlowTask,
    ) -> Self {
        PbCreateFlowTask {
            create_flow: Some(CreateFlowExpr {
                catalog_name,
                flow_name,
                source_table_names: source_table_names.into_iter().map(Into::into).collect(),
                sink_table_name: Some(sink_table_name.into()),
                or_replace,
                create_if_not_exists,
                expire_after: expire_after.map(|value| ExpireAfter { value }),
                eval_interval: eval_interval.map(|seconds| EvalInterval { seconds }),
                comment,
                sql,
                flow_options,
            }),
        }
    }
}

/// Drop flow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DropFlowTask {
    pub catalog_name: String,
    pub flow_name: String,
    pub flow_id: FlowId,
    pub drop_if_exists: bool,
}

impl TryFrom<PbDropFlowTask> for DropFlowTask {
    type Error = error::Error;

    fn try_from(pb: PbDropFlowTask) -> Result<Self> {
        let DropFlowExpr {
            catalog_name,
            flow_name,
            flow_id,
            drop_if_exists,
        } = pb.drop_flow.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected drop_flow",
        })?;
        let flow_id = flow_id
            .context(error::InvalidProtoMsgSnafu {
                err_msg: "expected flow_id",
            })?
            .id;
        Ok(DropFlowTask {
            catalog_name,
            flow_name,
            flow_id,
            drop_if_exists,
        })
    }
}

impl From<DropFlowTask> for PbDropFlowTask {
    fn from(
        DropFlowTask {
            catalog_name,
            flow_name,
            flow_id,
            drop_if_exists,
        }: DropFlowTask,
    ) -> Self {
        PbDropFlowTask {
            drop_flow: Some(DropFlowExpr {
                catalog_name,
                flow_name,
                flow_id: Some(api::v1::FlowId { id: flow_id }),
                drop_if_exists,
            }),
        }
    }
}

/// Represents the ID of the object being commented on (Table or Flow).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum CommentObjectId {
    Table(TableId),
    Flow(FlowId),
}

/// Comment on table, column, or flow
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct CommentOnTask {
    pub catalog_name: String,
    pub schema_name: String,
    pub object_type: CommentObjectType,
    pub object_name: String,
    /// Column name (only for Column comments)
    pub column_name: Option<String>,
    /// Object ID (Table or Flow) for validation and cache invalidation
    pub object_id: Option<CommentObjectId>,
    pub comment: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum CommentObjectType {
    Table,
    Column,
    Flow,
}

impl CommentOnTask {
    pub fn table_id(&self) -> Option<TableId> {
        match self.object_id.as_ref() {
            Some(CommentObjectId::Table(table_id)) => Some(*table_id),
            _ => None,
        }
    }

    pub fn flow_id(&self) -> Option<FlowId> {
        match self.object_id.as_ref() {
            Some(CommentObjectId::Flow(flow_id)) => Some(*flow_id),
            _ => None,
        }
    }

    fn set_table_id(&mut self, table_id: TableId) {
        self.object_id = Some(CommentObjectId::Table(table_id));
    }

    fn set_flow_id(&mut self, flow_id: FlowId) {
        self.object_id = Some(CommentObjectId::Flow(flow_id));
    }

    /// Returns the cache identifiers for the object being commented on.
    pub fn cache_idents(&self) -> Vec<CacheIdent> {
        match self.object_type {
            CommentObjectType::Table | CommentObjectType::Column => {
                let mut cache_idents = Vec::with_capacity(2);
                if let Some(CommentObjectId::Table(table_id)) = self.object_id.as_ref() {
                    cache_idents.push(CacheIdent::TableId(*table_id));
                }
                cache_idents.push(CacheIdent::TableName(TableName {
                    catalog_name: self.catalog_name.clone(),
                    schema_name: self.schema_name.clone(),
                    table_name: self.object_name.clone(),
                }));
                cache_idents
            }
            CommentObjectType::Flow => {
                let mut cache_idents = Vec::with_capacity(2);
                if let Some(CommentObjectId::Flow(flow_id)) = self.object_id.as_ref() {
                    cache_idents.push(CacheIdent::FlowId(*flow_id));
                }
                cache_idents.push(CacheIdent::FlowName(FlowName {
                    catalog_name: self.catalog_name.clone(),
                    flow_name: self.object_name.clone(),
                }));
                cache_idents
            }
        }
    }

    /// Enriches the `object_id` field of the `CommentOnTask`
    /// by looking up the corresponding table or flow ID using the provided managers.
    pub async fn enrich_object_id(
        &mut self,
        table_name_manager: &TableNameManager,
        flow_name_manager: &FlowNameManager,
    ) -> Result<()> {
        match self.object_type {
            CommentObjectType::Table | CommentObjectType::Column => {
                let table_id = table_name_manager
                    .get(TableNameKey::new(
                        &self.catalog_name,
                        &self.schema_name,
                        &self.object_name,
                    ))
                    .await?
                    .with_context(|| error::TableNotFoundSnafu {
                        table_name: format_full_table_name(
                            &self.catalog_name,
                            &self.schema_name,
                            &self.object_name,
                        ),
                    })?
                    .table_id();

                self.set_table_id(table_id);
            }
            CommentObjectType::Flow => {
                let flow_id = flow_name_manager
                    .get(&self.catalog_name, &self.object_name)
                    .await?
                    .with_context(|| error::FlowNotFoundSnafu {
                        flow_name: format_full_flow_name(&self.catalog_name, &self.object_name),
                    })?
                    .flow_id();

                self.set_flow_id(flow_id);
            }
        }

        Ok(())
    }
}

// Proto conversions for CommentObjectType
impl From<CommentObjectType> for PbCommentObjectType {
    fn from(object_type: CommentObjectType) -> Self {
        match object_type {
            CommentObjectType::Table => PbCommentObjectType::Table,
            CommentObjectType::Column => PbCommentObjectType::Column,
            CommentObjectType::Flow => PbCommentObjectType::Flow,
        }
    }
}

impl TryFrom<i32> for CommentObjectType {
    type Error = error::Error;

    fn try_from(value: i32) -> Result<Self> {
        match value {
            0 => Ok(CommentObjectType::Table),
            1 => Ok(CommentObjectType::Column),
            2 => Ok(CommentObjectType::Flow),
            _ => error::InvalidProtoMsgSnafu {
                err_msg: format!(
                    "Invalid CommentObjectType value: {}. Valid values are: 0 (Table), 1 (Column), 2 (Flow)",
                    value
                ),
            }
            .fail(),
        }
    }
}

// Proto conversions for CommentOnTask
impl TryFrom<PbCommentOnTask> for CommentOnTask {
    type Error = error::Error;

    fn try_from(pb: PbCommentOnTask) -> Result<Self> {
        let comment_on = pb.comment_on.context(error::InvalidProtoMsgSnafu {
            err_msg: "expected comment_on",
        })?;

        Ok(CommentOnTask {
            catalog_name: comment_on.catalog_name,
            schema_name: comment_on.schema_name,
            object_type: comment_on.object_type.try_into()?,
            object_name: comment_on.object_name,
            column_name: if comment_on.column_name.is_empty() {
                None
            } else {
                Some(comment_on.column_name)
            },
            comment: if comment_on.comment.is_empty() {
                None
            } else {
                Some(comment_on.comment)
            },
            object_id: None,
        })
    }
}

impl From<CommentOnTask> for PbCommentOnTask {
    fn from(task: CommentOnTask) -> Self {
        let pb_object_type: PbCommentObjectType = task.object_type.into();
        PbCommentOnTask {
            comment_on: Some(CommentOnExpr {
                catalog_name: task.catalog_name,
                schema_name: task.schema_name,
                object_type: pb_object_type as i32,
                object_name: task.object_name,
                column_name: task.column_name.unwrap_or_default(),
                comment: task.comment.unwrap_or_default(),
            }),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct QueryContext {
    pub current_catalog: String,
    pub current_schema: String,
    pub timezone: String,
    pub extensions: HashMap<String, String>,
    pub channel: u8,
    /// Maps region id -> snapshot upper bound sequence for that region.
    #[serde(default)]
    pub snapshot_seqs: HashMap<u64, u64>,
    /// Maps region id -> minimal SST sequence allowed for that region.
    #[serde(default)]
    pub sst_min_sequences: HashMap<u64, u64>,
}

impl QueryContext {
    /// Get the current catalog
    pub fn current_catalog(&self) -> &str {
        &self.current_catalog
    }

    /// Get the current schema
    pub fn current_schema(&self) -> &str {
        &self.current_schema
    }

    /// Get the timezone
    pub fn timezone(&self) -> &str {
        &self.timezone
    }

    /// Get the extensions
    pub fn extensions(&self) -> &HashMap<String, String> {
        &self.extensions
    }

    /// Get the channel
    pub fn channel(&self) -> u8 {
        self.channel
    }

    pub fn snapshot_seqs(&self) -> &HashMap<u64, u64> {
        &self.snapshot_seqs
    }

    pub fn sst_min_sequences(&self) -> &HashMap<u64, u64> {
        &self.sst_min_sequences
    }
}

/// Lightweight query context for flow operations containing only essential fields.
/// This is a subset of QueryContext that includes only the fields actually needed
/// for flow creation and execution.
#[derive(Debug, Clone, Serialize, PartialEq)]
pub struct FlowQueryContext {
    /// Current catalog name used for flow metadata and execution.
    pub catalog: String,
    /// Current schema name used for table resolution during flow execution.
    pub schema: String,
    /// Timezone used for timestamp evaluation in the flow.
    pub timezone: String,
    /// Query extensions carried into flow execution.
    #[serde(default)]
    pub extensions: HashMap<String, String>,
    /// Request channel propagated from the original query context.
    #[serde(default)]
    pub channel: u8,
    /// Per-region snapshot upper bounds bound during query planning/execution.
    #[serde(default)]
    pub snapshot_seqs: HashMap<u64, u64>,
    /// Per-region lower SST scan bounds carried with the flow context.
    #[serde(default)]
    pub sst_min_sequences: HashMap<u64, u64>,
}

impl<'de> Deserialize<'de> for FlowQueryContext {
    fn deserialize<D>(deserializer: D) -> result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        // Support both QueryContext format and FlowQueryContext format
        #[derive(Deserialize)]
        #[serde(untagged)]
        enum ContextCompat {
            Flow(FlowQueryContextHelper),
            Full(QueryContext),
        }

        #[derive(Deserialize)]
        struct FlowQueryContextHelper {
            catalog: String,
            schema: String,
            timezone: String,
            #[serde(default)]
            extensions: HashMap<String, String>,
            #[serde(default)]
            channel: u8,
            #[serde(default)]
            snapshot_seqs: HashMap<u64, u64>,
            #[serde(default)]
            sst_min_sequences: HashMap<u64, u64>,
        }

        match ContextCompat::deserialize(deserializer)? {
            ContextCompat::Flow(helper) => Ok(FlowQueryContext {
                catalog: helper.catalog,
                schema: helper.schema,
                timezone: helper.timezone,
                extensions: helper.extensions,
                channel: helper.channel,
                snapshot_seqs: helper.snapshot_seqs,
                sst_min_sequences: helper.sst_min_sequences,
            }),
            ContextCompat::Full(full_ctx) => Ok(full_ctx.into()),
        }
    }
}

impl From<PbQueryContext> for QueryContext {
    fn from(pb_ctx: PbQueryContext) -> Self {
        let (snapshot_seqs, sst_min_sequences) = pb_ctx
            .snapshot_seqs
            .map(|seqs| (seqs.snapshot_seqs, seqs.sst_min_sequences))
            .unwrap_or_default();

        Self {
            current_catalog: pb_ctx.current_catalog,
            current_schema: pb_ctx.current_schema,
            timezone: pb_ctx.timezone,
            extensions: pb_ctx.extensions,
            channel: pb_ctx.channel as u8,
            snapshot_seqs,
            sst_min_sequences,
        }
    }
}

impl From<QueryContext> for PbQueryContext {
    fn from(
        QueryContext {
            current_catalog,
            current_schema,
            timezone,
            extensions,
            channel,
            snapshot_seqs,
            sst_min_sequences,
        }: QueryContext,
    ) -> Self {
        PbQueryContext {
            current_catalog,
            current_schema,
            timezone,
            extensions,
            channel: channel as u32,
            snapshot_seqs: (!snapshot_seqs.is_empty() || !sst_min_sequences.is_empty()).then_some(
                api::v1::SnapshotSequences {
                    snapshot_seqs,
                    sst_min_sequences,
                },
            ),
            explain: None,
        }
    }
}

impl From<QueryContext> for FlowQueryContext {
    fn from(ctx: QueryContext) -> Self {
        Self {
            catalog: ctx.current_catalog,
            schema: ctx.current_schema,
            timezone: ctx.timezone,
            extensions: ctx.extensions,
            channel: ctx.channel,
            snapshot_seqs: ctx.snapshot_seqs,
            sst_min_sequences: ctx.sst_min_sequences,
        }
    }
}

impl From<FlowQueryContext> for QueryContext {
    fn from(flow_ctx: FlowQueryContext) -> Self {
        Self {
            current_catalog: flow_ctx.catalog,
            current_schema: flow_ctx.schema,
            timezone: flow_ctx.timezone,
            extensions: flow_ctx.extensions,
            channel: flow_ctx.channel,
            snapshot_seqs: flow_ctx.snapshot_seqs,
            sst_min_sequences: flow_ctx.sst_min_sequences,
        }
    }
}

impl From<FlowQueryContext> for PbQueryContext {
    fn from(flow_ctx: FlowQueryContext) -> Self {
        let query_ctx: QueryContext = flow_ctx.into();
        query_ctx.into()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use api::v1::{AlterTableExpr, ColumnDef, CreateTableExpr, SemanticType};
    use datatypes::schema::{ColumnSchema, Schema, SchemaBuilder};
    use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
    use store_api::storage::ConcreteDataType;
    use table::metadata::{TableInfo, TableMeta, TableType};
    use table::test_util::table_info::test_table_info;

    use super::{AlterTableTask, CreateTableTask, *};

    #[test]
    fn test_ddl_timeout_secs() {
        assert_eq!(ddl_timeout_secs(Duration::ZERO), 0);
        assert_eq!(ddl_timeout_secs(Duration::from_nanos(1)), 1);
        assert_eq!(ddl_timeout_secs(Duration::from_secs(1)), 1);
        assert_eq!(ddl_timeout_secs(Duration::from_millis(1500)), 2);
        assert_eq!(
            ddl_timeout_secs(Duration::from_secs(u32::MAX as u64 + 1)),
            u32::MAX
        );
    }

    #[test]
    fn test_basic_ser_de_create_table_task() {
        let schema = SchemaBuilder::default().build().unwrap();
        let table_info = test_table_info(1025, "foo", "bar", "baz", Arc::new(schema));
        let task = CreateTableTask::new(CreateTableExpr::default(), Vec::new(), table_info);

        let output = serde_json::to_vec(&task).unwrap();

        let de = serde_json::from_slice(&output).unwrap();
        assert_eq!(task, de);
    }

    #[test]
    fn test_basic_ser_de_alter_table_task() {
        let task = AlterTableTask {
            alter_table: AlterTableExpr::default(),
        };

        let output = serde_json::to_vec(&task).unwrap();

        let de = serde_json::from_slice(&output).unwrap();
        assert_eq!(task, de);
    }

    #[test]
    fn test_undrop_table_task_pb_roundtrip() {
        let expected = UndropTableTask { table_id: 1024 };
        let request = SubmitDdlTaskRequest::new(
            QueryContext::default(),
            DdlTask::UndropTable(expected.clone()),
        );

        let pb = PbDdlTaskRequest::try_from(request).unwrap();
        let pb_task = pb.task.unwrap();
        let de = DdlTask::try_from(pb_task).unwrap();

        assert!(matches!(de, DdlTask::UndropTable(task) if task == expected));
    }

    #[test]
    fn test_purge_dropped_table_task_pb_roundtrip() {
        let expected = PurgeDroppedTableTask { table_id: 1024 };
        let request = SubmitDdlTaskRequest::new(
            QueryContext::default(),
            DdlTask::PurgeDroppedTable(expected.clone()),
        );

        let pb = PbDdlTaskRequest::try_from(request).unwrap();
        let pb_task = pb.task.unwrap();
        let de = DdlTask::try_from(pb_task).unwrap();

        assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected));
    }

    #[test]
    fn test_create_database_creator_transport_sanitizes_extension() {
        let request = |creator| {
            let mut query_context = QueryContext::default();
            query_context.extensions.insert(
                CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(),
                "spoofed".to_string(),
            );
            SubmitDdlTaskRequest::new(
                query_context,
                DdlTask::new_create_database(
                    "greptime".to_string(),
                    "test".to_string(),
                    false,
                    HashMap::new(),
                    creator,
                ),
            )
        };
        let creator = CreatorGrantIntent {
            username: "alice".to_string(),
            created_at_ns: 42,
        };

        let mut pb = PbDdlTaskRequest::try_from(request(Some(creator.clone()))).unwrap();
        let encoded = pb
            .query_context
            .as_ref()
            .unwrap()
            .extensions
            .get(CREATE_DATABASE_CREATOR_EXTENSION_KEY)
            .unwrap();
        assert_eq!(creator, serde_json::from_str(encoded).unwrap());
        assert!(matches!(
            DdlTask::try_from(pb.task.take().unwrap()).unwrap(),
            DdlTask::CreateDatabase(CreateDatabaseTask { creator: None, .. })
        ));

        let pb = PbDdlTaskRequest::try_from(request(None)).unwrap();
        assert!(
            !pb.query_context
                .unwrap()
                .extensions
                .contains_key(CREATE_DATABASE_CREATOR_EXTENSION_KEY)
        );
    }

    #[test]
    fn test_undrop_table_task_json_roundtrip() {
        let task = UndropTableTask { table_id: 1024 };

        let output = serde_json::to_vec(&task).unwrap();

        let de = serde_json::from_slice(&output).unwrap();
        assert_eq!(task, de);
    }

    #[test]
    fn test_purge_dropped_table_task_json_roundtrip() {
        let task = PurgeDroppedTableTask { table_id: 1024 };

        let output = serde_json::to_vec(&task).unwrap();

        let de = serde_json::from_slice(&output).unwrap();
        assert_eq!(task, de);
    }

    #[test]
    fn test_sort_columns() {
        // construct RawSchema
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                "column3".to_string(),
                ConcreteDataType::string_datatype(),
                true,
            ),
            ColumnSchema::new(
                "column1".to_string(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
            ColumnSchema::new(
                "column2".to_string(),
                ConcreteDataType::float64_datatype(),
                true,
            ),
        ]));

        // construct RawTableMeta
        let meta = TableMeta {
            schema,
            primary_key_indices: vec![0],
            value_indices: vec![2],
            engine: METRIC_ENGINE_NAME.to_string(),
            next_column_id: 0,
            options: Default::default(),
            created_on: Default::default(),
            updated_on: Default::default(),
            partition_key_indices: Default::default(),
            column_ids: Default::default(),
        };

        // construct TableInfo
        let raw_table_info = TableInfo {
            ident: Default::default(),
            meta,
            name: Default::default(),
            desc: Default::default(),
            catalog_name: Default::default(),
            schema_name: Default::default(),
            table_type: TableType::Base,
        };

        // construct create table expr
        let create_table_expr = CreateTableExpr {
            column_defs: vec![
                ColumnDef {
                    name: "column3".to_string(),
                    semantic_type: SemanticType::Tag as i32,
                    ..Default::default()
                },
                ColumnDef {
                    name: "column1".to_string(),
                    semantic_type: SemanticType::Timestamp as i32,
                    ..Default::default()
                },
                ColumnDef {
                    name: "column2".to_string(),
                    semantic_type: SemanticType::Field as i32,
                    ..Default::default()
                },
            ],
            primary_keys: vec!["column3".to_string()],
            ..Default::default()
        };

        let mut create_table_task =
            CreateTableTask::new(create_table_expr, Vec::new(), raw_table_info);

        // Call the sort_columns method
        create_table_task.sort_columns();

        // Assert that the columns are sorted correctly
        assert_eq!(
            create_table_task.create_table.column_defs[0].name,
            "column1".to_string()
        );
        assert_eq!(
            create_table_task.create_table.column_defs[1].name,
            "column2".to_string()
        );
        assert_eq!(
            create_table_task.create_table.column_defs[2].name,
            "column3".to_string()
        );

        // Assert that the table_info is updated correctly
        assert_eq!(
            create_table_task.table_info.meta.schema.timestamp_index(),
            Some(0)
        );
        assert_eq!(
            create_table_task.table_info.meta.primary_key_indices,
            vec![2]
        );
        assert_eq!(create_table_task.table_info.meta.value_indices, vec![0, 1]);
    }

    #[test]
    fn test_flow_query_context_conversion_from_query_context() {
        use std::collections::HashMap;
        let mut extensions = HashMap::new();
        extensions.insert("key1".to_string(), "value1".to_string());
        extensions.insert("key2".to_string(), "value2".to_string());

        let query_ctx = QueryContext {
            current_catalog: "test_catalog".to_string(),
            current_schema: "test_schema".to_string(),
            timezone: "UTC".to_string(),
            extensions,
            channel: 5,
            snapshot_seqs: HashMap::from([(10, 100)]),
            sst_min_sequences: HashMap::from([(10, 90)]),
        };

        let flow_ctx: FlowQueryContext = query_ctx.into();

        assert_eq!(flow_ctx.catalog, "test_catalog");
        assert_eq!(flow_ctx.schema, "test_schema");
        assert_eq!(flow_ctx.timezone, "UTC");
        assert_eq!(flow_ctx.channel, 5);
        assert_eq!(flow_ctx.snapshot_seqs, HashMap::from([(10, 100)]));
        assert_eq!(flow_ctx.sst_min_sequences, HashMap::from([(10, 90)]));
    }

    #[test]
    fn test_flow_query_context_conversion_to_query_context() {
        let flow_ctx = FlowQueryContext {
            catalog: "prod_catalog".to_string(),
            schema: "public".to_string(),
            timezone: "America/New_York".to_string(),
            extensions: HashMap::from([("k".to_string(), "v".to_string())]),
            channel: 7,
            snapshot_seqs: HashMap::from([(11, 111)]),
            sst_min_sequences: HashMap::from([(11, 101)]),
        };

        let query_ctx: QueryContext = flow_ctx.clone().into();

        assert_eq!(query_ctx.current_catalog, "prod_catalog");
        assert_eq!(query_ctx.current_schema, "public");
        assert_eq!(query_ctx.timezone, "America/New_York");
        assert_eq!(
            query_ctx.extensions,
            HashMap::from([("k".to_string(), "v".to_string())])
        );
        assert_eq!(query_ctx.channel, 7);
        assert_eq!(query_ctx.snapshot_seqs, HashMap::from([(11, 111)]));
        assert_eq!(query_ctx.sst_min_sequences, HashMap::from([(11, 101)]));

        // Test roundtrip conversion
        let flow_ctx_roundtrip: FlowQueryContext = query_ctx.into();
        assert_eq!(flow_ctx, flow_ctx_roundtrip);
    }

    #[test]
    fn test_flow_query_context_serialization() {
        let flow_ctx = FlowQueryContext {
            catalog: "test_catalog".to_string(),
            schema: "test_schema".to_string(),
            timezone: "UTC".to_string(),
            extensions: HashMap::new(),
            channel: 0,
            snapshot_seqs: HashMap::new(),
            sst_min_sequences: HashMap::new(),
        };

        let serialized = serde_json::to_string(&flow_ctx).unwrap();
        let deserialized: FlowQueryContext = serde_json::from_str(&serialized).unwrap();

        assert_eq!(flow_ctx, deserialized);

        // Verify JSON structure
        let json_value: serde_json::Value = serde_json::from_str(&serialized).unwrap();
        assert_eq!(json_value["catalog"], "test_catalog");
        assert_eq!(json_value["schema"], "test_schema");
        assert_eq!(json_value["timezone"], "UTC");
    }

    #[test]
    fn test_flow_query_context_conversion_to_pb() {
        let flow_ctx = FlowQueryContext {
            catalog: "pb_catalog".to_string(),
            schema: "pb_schema".to_string(),
            timezone: "Asia/Tokyo".to_string(),
            extensions: HashMap::from([("x".to_string(), "y".to_string())]),
            channel: 6,
            snapshot_seqs: HashMap::from([(3, 30)]),
            sst_min_sequences: HashMap::from([(3, 21)]),
        };

        let pb_ctx: PbQueryContext = flow_ctx.into();

        assert_eq!(pb_ctx.current_catalog, "pb_catalog");
        assert_eq!(pb_ctx.current_schema, "pb_schema");
        assert_eq!(pb_ctx.timezone, "Asia/Tokyo");
        assert_eq!(
            pb_ctx.extensions,
            HashMap::from([("x".to_string(), "y".to_string())])
        );
        assert_eq!(pb_ctx.channel, 6);
        assert_eq!(
            pb_ctx.snapshot_seqs,
            Some(api::v1::SnapshotSequences {
                snapshot_seqs: HashMap::from([(3, 30)]),
                sst_min_sequences: HashMap::from([(3, 21)]),
            })
        );
        assert!(pb_ctx.explain.is_none());
    }

    #[test]
    fn test_pb_query_context_roundtrip_with_snapshot_sequences() {
        let pb = PbQueryContext {
            current_catalog: "c1".to_string(),
            current_schema: "s1".to_string(),
            timezone: "UTC".to_string(),
            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
            channel: 3,
            snapshot_seqs: Some(api::v1::SnapshotSequences {
                snapshot_seqs: HashMap::from([(1, 100)]),
                sst_min_sequences: HashMap::from([(1, 90)]),
            }),
            explain: None,
        };

        let query_ctx: QueryContext = pb.clone().into();
        let pb_roundtrip: PbQueryContext = query_ctx.into();

        assert_eq!(pb_roundtrip.current_catalog, pb.current_catalog);
        assert_eq!(pb_roundtrip.current_schema, pb.current_schema);
        assert_eq!(pb_roundtrip.timezone, pb.timezone);
        assert_eq!(pb_roundtrip.extensions, pb.extensions);
        assert_eq!(pb_roundtrip.channel, pb.channel);
        assert_eq!(pb_roundtrip.snapshot_seqs, pb.snapshot_seqs);
    }
}