operator/
statement.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
15mod admin;
16mod comment;
17mod copy_database;
18mod copy_query_to;
19mod copy_table_from;
20mod copy_table_to;
21mod cursor;
22pub mod ddl;
23mod describe;
24mod dml;
25mod kill;
26mod set;
27mod show;
28mod tql;
29
30use std::collections::HashMap;
31use std::sync::Arc;
32
33use api::v1::RowInsertRequests;
34use catalog::CatalogManagerRef;
35use catalog::kvbackend::KvBackendCatalogManager;
36use catalog::process_manager::ProcessManagerRef;
37use client::RecordBatches;
38use client::error::{ExternalSnafu as ClientExternalSnafu, Result as ClientResult};
39use client::inserter::{InsertOptions, Inserter};
40use common_error::ext::BoxedError;
41use common_meta::cache_invalidator::CacheInvalidatorRef;
42use common_meta::key::flow::{FlowMetadataManager, FlowMetadataManagerRef};
43use common_meta::key::schema_name::SchemaNameKey;
44use common_meta::key::view_info::{ViewInfoManager, ViewInfoManagerRef};
45use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
46use common_meta::kv_backend::KvBackendRef;
47use common_meta::procedure_executor::ProcedureExecutorRef;
48use common_query::Output;
49use common_telemetry::{debug, tracing, warn};
50use common_time::Timestamp;
51use common_time::range::TimestampRange;
52use datafusion_expr::LogicalPlan;
53use datatypes::prelude::ConcreteDataType;
54use datatypes::schema::ColumnSchema;
55use humantime::format_duration;
56use itertools::Itertools;
57use partition::manager::PartitionRuleManagerRef;
58use query::QueryEngineRef;
59use query::parser::QueryStatement;
60use session::context::{Channel, QueryContextBuilder, QueryContextRef};
61use session::table_name::table_idents_to_full_name;
62use set::{set_query_timeout, set_read_preference};
63use snafu::{OptionExt, ResultExt, ensure};
64use sql::ast::ObjectNamePartExt;
65use sql::statements::OptionMap;
66use sql::statements::copy::{
67    CopyDatabase, CopyDatabaseArgument, CopyQueryToArgument, CopyTable, CopyTableArgument,
68};
69use sql::statements::set_variables::SetVariables;
70use sql::statements::show::ShowCreateTableVariant;
71use sql::statements::statement::Statement;
72use sql::util::format_raw_object_name;
73use sqlparser::ast::ObjectName;
74use store_api::mito_engine_options::{APPEND_MODE_KEY, TTL_KEY};
75use table::TableRef;
76use table::requests::{CopyDatabaseRequest, CopyDirection, CopyQueryToRequest, CopyTableRequest};
77use table::table_name::TableName;
78use table::table_reference::TableReference;
79
80use self::set::{
81    set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding,
82};
83use crate::error::{
84    self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu,
85    PlanStatementSnafu, Result, SchemaNotFoundSnafu, SqlCommonSnafu, TableMetadataManagerSnafu,
86    TableNotFoundSnafu, UnexpectedSnafu, UpgradeCatalogManagerRefSnafu,
87};
88use crate::insert::InserterRef;
89use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY};
90use crate::statement::set::set_allow_query_fallback;
91
92/// A configurator that customizes or enhances a [`StatementExecutor`].
93#[async_trait::async_trait]
94pub trait StatementExecutorConfigurator: Send + Sync {
95    async fn configure(
96        &self,
97        executor: StatementExecutor,
98        ctx: ExecutorConfigureContext,
99    ) -> std::result::Result<StatementExecutor, BoxedError>;
100}
101
102pub type StatementExecutorConfiguratorRef = Arc<dyn StatementExecutorConfigurator>;
103
104pub struct ExecutorConfigureContext {
105    pub kv_backend: KvBackendRef,
106}
107
108#[derive(Clone)]
109pub struct StatementExecutor {
110    catalog_manager: CatalogManagerRef,
111    query_engine: QueryEngineRef,
112    procedure_executor: ProcedureExecutorRef,
113    table_metadata_manager: TableMetadataManagerRef,
114    flow_metadata_manager: FlowMetadataManagerRef,
115    view_info_manager: ViewInfoManagerRef,
116    partition_manager: PartitionRuleManagerRef,
117    cache_invalidator: CacheInvalidatorRef,
118    inserter: InserterRef,
119    process_manager: Option<ProcessManagerRef>,
120    origin_frontend_addr: String,
121    #[cfg(feature = "enterprise")]
122    trigger_querier: Option<TriggerQuerierRef>,
123}
124
125pub type StatementExecutorRef = Arc<StatementExecutor>;
126
127/// Trait for querying trigger info, such as `SHOW CREATE TRIGGER` etc.
128#[cfg(feature = "enterprise")]
129#[async_trait::async_trait]
130pub trait TriggerQuerier: Send + Sync {
131    // Query the `SHOW CREATE TRIGGER` statement for the given trigger.
132    async fn show_create_trigger(
133        &self,
134        catalog: &str,
135        trigger: &str,
136        query_ctx: &QueryContextRef,
137    ) -> std::result::Result<Output, BoxedError>;
138
139    fn as_any(&self) -> &dyn std::any::Any;
140}
141
142#[cfg(feature = "enterprise")]
143pub type TriggerQuerierRef = Arc<dyn TriggerQuerier>;
144
145impl StatementExecutor {
146    #[allow(clippy::too_many_arguments)]
147    pub fn new(
148        catalog_manager: CatalogManagerRef,
149        query_engine: QueryEngineRef,
150        procedure_executor: ProcedureExecutorRef,
151        kv_backend: KvBackendRef,
152        cache_invalidator: CacheInvalidatorRef,
153        inserter: InserterRef,
154        partition_manager: PartitionRuleManagerRef,
155        process_manager: Option<ProcessManagerRef>,
156        origin_frontend_addr: String,
157    ) -> Self {
158        Self {
159            catalog_manager,
160            query_engine,
161            procedure_executor,
162            table_metadata_manager: Arc::new(TableMetadataManager::new(kv_backend.clone())),
163            flow_metadata_manager: Arc::new(FlowMetadataManager::new(kv_backend.clone())),
164            view_info_manager: Arc::new(ViewInfoManager::new(kv_backend.clone())),
165            partition_manager,
166            cache_invalidator,
167            inserter,
168            process_manager,
169            origin_frontend_addr,
170            #[cfg(feature = "enterprise")]
171            trigger_querier: None,
172        }
173    }
174
175    #[cfg(feature = "enterprise")]
176    pub fn with_trigger_querier(mut self, querier: TriggerQuerierRef) -> Self {
177        self.trigger_querier = Some(querier);
178        self
179    }
180
181    #[cfg(feature = "testing")]
182    pub async fn execute_stmt(
183        &self,
184        stmt: QueryStatement,
185        query_ctx: QueryContextRef,
186    ) -> Result<Output> {
187        match stmt {
188            QueryStatement::Sql(stmt) => self.execute_sql(stmt, query_ctx).await,
189            QueryStatement::Promql(_, _) => self.plan_exec(stmt, query_ctx).await,
190        }
191    }
192
193    #[tracing::instrument(skip_all)]
194    pub async fn execute_sql(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
195        match stmt {
196            Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
197                self.plan_exec(QueryStatement::Sql(stmt), query_ctx).await
198            }
199
200            Statement::DeclareCursor(declare_cursor) => {
201                self.declare_cursor(declare_cursor, query_ctx).await
202            }
203            Statement::FetchCursor(fetch_cursor) => {
204                self.fetch_cursor(fetch_cursor, query_ctx).await
205            }
206            Statement::CloseCursor(close_cursor) => {
207                self.close_cursor(close_cursor, query_ctx).await
208            }
209
210            Statement::Insert(insert) => self.insert(insert, query_ctx).await,
211
212            Statement::Tql(tql) => self.execute_tql(tql, query_ctx).await,
213
214            Statement::DescribeTable(stmt) => self.describe_table(stmt, query_ctx).await,
215
216            Statement::ShowDatabases(stmt) => self.show_databases(stmt, query_ctx).await,
217
218            Statement::ShowTables(stmt) => self.show_tables(stmt, query_ctx).await,
219
220            Statement::ShowTableStatus(stmt) => self.show_table_status(stmt, query_ctx).await,
221
222            Statement::ShowCollation(kind) => self.show_collation(kind, query_ctx).await,
223
224            Statement::ShowCharset(kind) => self.show_charset(kind, query_ctx).await,
225
226            Statement::ShowViews(stmt) => self.show_views(stmt, query_ctx).await,
227
228            Statement::ShowFlows(stmt) => self.show_flows(stmt, query_ctx).await,
229
230            #[cfg(feature = "enterprise")]
231            Statement::ShowTriggers(stmt) => self.show_triggers(stmt, query_ctx).await,
232
233            Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(stmt)) => {
234                let query_output = self
235                    .plan_exec(QueryStatement::Sql(*stmt.query), query_ctx)
236                    .await?;
237                let req = to_copy_query_request(stmt.arg)?;
238
239                self.copy_query_to(req, query_output)
240                    .await
241                    .map(Output::new_with_affected_rows)
242            }
243
244            Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => {
245                let req = to_copy_table_request(stmt, query_ctx.clone())?;
246                match req.direction {
247                    CopyDirection::Export => self
248                        .copy_table_to(req, query_ctx)
249                        .await
250                        .map(Output::new_with_affected_rows),
251                    CopyDirection::Import => self.copy_table_from(req, query_ctx).await,
252                }
253            }
254
255            Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
256                match copy_database {
257                    CopyDatabase::To(arg) => {
258                        self.copy_database_to(
259                            to_copy_database_request(arg, &query_ctx)?,
260                            query_ctx.clone(),
261                        )
262                        .await
263                    }
264                    CopyDatabase::From(arg) => {
265                        self.copy_database_from(
266                            to_copy_database_request(arg, &query_ctx)?,
267                            query_ctx,
268                        )
269                        .await
270                    }
271                }
272            }
273
274            Statement::CreateTable(stmt) => {
275                let _ = self.create_table(stmt, query_ctx).await?;
276                Ok(Output::new_with_affected_rows(0))
277            }
278            Statement::CreateTableLike(stmt) => {
279                let _ = self.create_table_like(stmt, query_ctx).await?;
280                Ok(Output::new_with_affected_rows(0))
281            }
282            Statement::CreateExternalTable(stmt) => {
283                let _ = self.create_external_table(stmt, query_ctx).await?;
284                Ok(Output::new_with_affected_rows(0))
285            }
286            Statement::CreateFlow(stmt) => self.create_flow(stmt, query_ctx).await,
287            #[cfg(feature = "enterprise")]
288            Statement::CreateTrigger(stmt) => self.create_trigger(stmt, query_ctx).await,
289            Statement::DropFlow(stmt) => {
290                self.drop_flow(
291                    query_ctx.current_catalog().to_string(),
292                    format_raw_object_name(stmt.flow_name()),
293                    stmt.drop_if_exists(),
294                    query_ctx,
295                )
296                .await
297            }
298            #[cfg(feature = "enterprise")]
299            Statement::DropTrigger(stmt) => {
300                self.drop_trigger(
301                    query_ctx.current_catalog().to_string(),
302                    format_raw_object_name(stmt.trigger_name()),
303                    stmt.drop_if_exists(),
304                    query_ctx,
305                )
306                .await
307            }
308            Statement::CreateView(stmt) => {
309                let _ = self.create_view(stmt, query_ctx).await?;
310                Ok(Output::new_with_affected_rows(0))
311            }
312            Statement::DropView(stmt) => {
313                let (catalog_name, schema_name, view_name) =
314                    table_idents_to_full_name(&stmt.view_name, &query_ctx)
315                        .map_err(BoxedError::new)
316                        .context(ExternalSnafu)?;
317
318                self.drop_view(
319                    catalog_name,
320                    schema_name,
321                    view_name,
322                    stmt.drop_if_exists,
323                    query_ctx,
324                )
325                .await
326            }
327            Statement::AlterTable(alter_table) => self.alter_table(alter_table, query_ctx).await,
328
329            Statement::AlterDatabase(alter_database) => {
330                self.alter_database(alter_database, query_ctx).await
331            }
332
333            #[cfg(feature = "enterprise")]
334            Statement::AlterTrigger(alter_trigger) => {
335                self.alter_trigger(alter_trigger, query_ctx).await
336            }
337
338            Statement::DropTable(stmt) => {
339                let mut table_names = Vec::with_capacity(stmt.table_names().len());
340                for table_name_stmt in stmt.table_names() {
341                    let (catalog, schema, table) =
342                        table_idents_to_full_name(table_name_stmt, &query_ctx)
343                            .map_err(BoxedError::new)
344                            .context(ExternalSnafu)?;
345                    table_names.push(TableName::new(catalog, schema, table));
346                }
347                self.drop_tables(&table_names[..], stmt.drop_if_exists(), query_ctx.clone())
348                    .await
349            }
350            Statement::UndropTable(stmt) => {
351                let (catalog, schema, table) =
352                    table_idents_to_full_name(stmt.table_name(), &query_ctx)
353                        .map_err(BoxedError::new)
354                        .context(ExternalSnafu)?;
355                self.undrop_table(TableName::new(catalog, schema, table), query_ctx)
356                    .await
357            }
358            Statement::DropDatabase(stmt) => {
359                self.drop_database(
360                    query_ctx.current_catalog().to_string(),
361                    format_raw_object_name(stmt.name()),
362                    stmt.drop_if_exists(),
363                    query_ctx,
364                )
365                .await
366            }
367            Statement::TruncateTable(stmt) => {
368                let (catalog, schema, table) =
369                    table_idents_to_full_name(stmt.table_name(), &query_ctx)
370                        .map_err(BoxedError::new)
371                        .context(ExternalSnafu)?;
372                let table_name = TableName::new(catalog, schema, table);
373                let time_ranges = self
374                    .convert_truncate_time_ranges(&table_name, stmt.time_ranges(), &query_ctx)
375                    .await?;
376                self.truncate_table(table_name, time_ranges, query_ctx)
377                    .await
378            }
379            Statement::CreateDatabase(stmt) => {
380                self.create_database(
381                    &format_raw_object_name(&stmt.name),
382                    stmt.if_not_exists,
383                    stmt.options.into_map(),
384                    query_ctx,
385                )
386                .await
387            }
388            Statement::ShowCreateDatabase(show) => {
389                let (catalog, database) =
390                    idents_to_full_database_name(&show.database_name, &query_ctx)
391                        .map_err(BoxedError::new)
392                        .context(ExternalSnafu)?;
393                let table_metadata_manager = self
394                    .catalog_manager
395                    .as_any()
396                    .downcast_ref::<KvBackendCatalogManager>()
397                    .map(|manager| manager.table_metadata_manager_ref().clone())
398                    .context(UpgradeCatalogManagerRefSnafu)?;
399                let opts: HashMap<String, String> = table_metadata_manager
400                    .schema_manager()
401                    .get(SchemaNameKey::new(&catalog, &database))
402                    .await
403                    .context(TableMetadataManagerSnafu)?
404                    .context(SchemaNotFoundSnafu {
405                        schema_info: &database,
406                    })?
407                    .into_inner()
408                    .into();
409
410                self.show_create_database(&database, opts.into()).await
411            }
412            Statement::ShowCreateTable(show) => {
413                let (catalog, schema, table) =
414                    table_idents_to_full_name(&show.table_name, &query_ctx)
415                        .map_err(BoxedError::new)
416                        .context(ExternalSnafu)?;
417
418                let table_ref = self
419                    .catalog_manager
420                    .table(&catalog, &schema, &table, Some(&query_ctx))
421                    .await
422                    .context(CatalogSnafu)?
423                    .context(TableNotFoundSnafu { table_name: &table })?;
424                let table_name = TableName::new(catalog, schema, table);
425
426                match show.variant {
427                    ShowCreateTableVariant::Original => {
428                        self.show_create_table(table_name, table_ref, query_ctx)
429                            .await
430                    }
431                    ShowCreateTableVariant::PostgresForeignTable => {
432                        self.show_create_table_for_pg(table_name, table_ref, query_ctx)
433                            .await
434                    }
435                }
436            }
437            Statement::ShowCreateFlow(show) => self.show_create_flow(show, query_ctx).await,
438            Statement::ShowCreateView(show) => self.show_create_view(show, query_ctx).await,
439            #[cfg(feature = "enterprise")]
440            Statement::ShowCreateTrigger(show) => self.show_create_trigger(show, query_ctx).await,
441            Statement::SetVariables(set_var) => self.set_variables(set_var, query_ctx),
442            Statement::ShowVariables(show_variable) => self.show_variable(show_variable, query_ctx),
443            Statement::Comment(stmt) => self.comment(stmt, query_ctx).await,
444            Statement::ShowColumns(show_columns) => {
445                self.show_columns(show_columns, query_ctx).await
446            }
447            Statement::ShowIndex(show_index) => self.show_index(show_index, query_ctx).await,
448            Statement::ShowRegion(show_region) => self.show_region(show_region, query_ctx).await,
449            Statement::ShowStatus(_) => self.show_status(query_ctx).await,
450            Statement::ShowSearchPath(_) => self.show_search_path(query_ctx).await,
451            Statement::Use(db) => self.use_database(db, query_ctx).await,
452            Statement::Admin(admin) => self.execute_admin_command(admin, query_ctx).await,
453            Statement::Kill(kill) => self.execute_kill(query_ctx, kill).await,
454            Statement::ShowProcesslist(show) => self.show_processlist(show, query_ctx).await,
455        }
456    }
457
458    pub async fn use_database(&self, db: String, query_ctx: QueryContextRef) -> Result<Output> {
459        let catalog = query_ctx.current_catalog();
460        ensure!(
461            self.catalog_manager
462                .schema_exists(catalog, db.as_ref(), Some(&query_ctx))
463                .await
464                .context(CatalogSnafu)?,
465            SchemaNotFoundSnafu { schema_info: &db }
466        );
467
468        query_ctx.set_current_schema(&db);
469
470        Ok(Output::new_with_record_batches(RecordBatches::empty()))
471    }
472
473    fn set_variables(&self, set_var: SetVariables, query_ctx: QueryContextRef) -> Result<Output> {
474        let var_name = set_var.variable.to_string().to_uppercase();
475
476        debug!(
477            "Trying to set {}={} for session: {} ",
478            var_name,
479            set_var.value.iter().map(|e| e.to_string()).join(", "),
480            query_ctx.conn_info()
481        );
482
483        match var_name.as_str() {
484            "READ_PREFERENCE" => set_read_preference(set_var.value, query_ctx)?,
485
486            "@@TIME_ZONE" | "@@SESSION.TIME_ZONE" | "TIMEZONE" | "TIME_ZONE" => {
487                set_timezone(set_var.value, query_ctx)?
488            }
489
490            "BYTEA_OUTPUT" => set_bytea_output(set_var.value, query_ctx)?,
491
492            // Same as "bytea_output", we just ignore it here.
493            // Not harmful since it only relates to how date is viewed in client app's output.
494            // The tracked issue is https://github.com/GreptimeTeam/greptimedb/issues/3442.
495            "DATESTYLE" => set_datestyle(set_var.value, query_ctx)?,
496            "INTERVALSTYLE" => set_intervalstyle(set_var.value, query_ctx)?,
497
498            // Allow query to fallback when failed to push down.
499            "ALLOW_QUERY_FALLBACK" => set_allow_query_fallback(set_var.value, query_ctx)?,
500
501            "CLIENT_ENCODING" => validate_client_encoding(set_var)?,
502            "@@SESSION.MAX_EXECUTION_TIME" | "MAX_EXECUTION_TIME" => match query_ctx.channel() {
503                Channel::Mysql => set_query_timeout(set_var.value, query_ctx)?,
504                Channel::Postgres => {
505                    warn!(
506                        "Unsupported set variable {} for channel {:?}",
507                        var_name,
508                        query_ctx.channel()
509                    );
510                    query_ctx.set_warning(format!("Unsupported set variable {}", var_name))
511                }
512                _ => {
513                    return NotSupportedSnafu {
514                        feat: format!("Unsupported set variable {}", var_name),
515                    }
516                    .fail();
517                }
518            },
519            "STATEMENT_TIMEOUT" => match query_ctx.channel() {
520                Channel::Postgres => set_query_timeout(set_var.value, query_ctx)?,
521                Channel::Mysql => {
522                    warn!(
523                        "Unsupported set variable {} for channel {:?}",
524                        var_name,
525                        query_ctx.channel()
526                    );
527                    query_ctx.set_warning(format!("Unsupported set variable {}", var_name));
528                }
529                _ => {
530                    return NotSupportedSnafu {
531                        feat: format!("Unsupported set variable {}", var_name),
532                    }
533                    .fail();
534                }
535            },
536            "SEARCH_PATH" => {
537                if query_ctx.channel() == Channel::Postgres {
538                    let search_path = set_var.search_path().context(NotSupportedSnafu {
539                        feat: "Unsupported search path in set variable statement",
540                    })?;
541                    query_ctx.set_current_schema(search_path);
542                } else {
543                    return NotSupportedSnafu {
544                        feat: format!("Unsupported set variable {}", var_name),
545                    }
546                    .fail();
547                }
548            }
549            _ => {
550                if query_ctx.channel() == Channel::Postgres || query_ctx.channel() == Channel::Mysql
551                {
552                    // For unknown SET statements, we give a warning with success.
553                    // This prevents the SET call from becoming a blocker of MySQL/Postgres clients'
554                    // connection establishment.
555                    warn!(
556                        "Unsupported set variable {} for channel {:?}",
557                        var_name,
558                        query_ctx.channel()
559                    );
560                    query_ctx.set_warning(format!("Unsupported set variable {}", var_name));
561                } else {
562                    return NotSupportedSnafu {
563                        feat: format!("Unsupported set variable {}", var_name),
564                    }
565                    .fail();
566                }
567            }
568        }
569        Ok(Output::new_with_affected_rows(0))
570    }
571
572    #[tracing::instrument(skip_all)]
573    pub async fn plan(
574        &self,
575        stmt: &QueryStatement,
576        query_ctx: QueryContextRef,
577    ) -> Result<LogicalPlan> {
578        self.query_engine
579            .planner()
580            .plan(stmt, query_ctx)
581            .await
582            .context(PlanStatementSnafu)
583    }
584
585    /// Execute [`LogicalPlan`] directly.
586    #[tracing::instrument(skip_all)]
587    pub async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
588        self.query_engine
589            .execute(plan, query_ctx)
590            .await
591            .context(ExecLogicalPlanSnafu)
592    }
593
594    pub fn optimize_logical_plan(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
595        self.query_engine
596            .planner()
597            .optimize(plan)
598            .context(PlanStatementSnafu)
599    }
600
601    #[tracing::instrument(skip_all)]
602    async fn plan_exec(&self, stmt: QueryStatement, query_ctx: QueryContextRef) -> Result<Output> {
603        let plan = self.plan(&stmt, query_ctx.clone()).await?;
604        self.exec_plan(plan, query_ctx).await
605    }
606
607    async fn get_table(&self, table_ref: &TableReference<'_>) -> Result<TableRef> {
608        let TableReference {
609            catalog,
610            schema,
611            table,
612        } = table_ref;
613        self.catalog_manager
614            .table(catalog, schema, table, None)
615            .await
616            .context(CatalogSnafu)?
617            .with_context(|| TableNotFoundSnafu {
618                table_name: table_ref.to_string(),
619            })
620    }
621
622    pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
623        &self.procedure_executor
624    }
625
626    pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
627        &self.cache_invalidator
628    }
629
630    /// Convert truncate time ranges for the given table from sql values to timestamps
631    ///
632    pub async fn convert_truncate_time_ranges(
633        &self,
634        table_name: &TableName,
635        sql_values_time_range: &[(sqlparser::ast::Value, sqlparser::ast::Value)],
636        query_ctx: &QueryContextRef,
637    ) -> Result<Vec<(Timestamp, Timestamp)>> {
638        if sql_values_time_range.is_empty() {
639            return Ok(vec![]);
640        }
641        let table = self.get_table(&table_name.table_ref()).await?;
642        let info = table.table_info();
643        let time_index_dt = info
644            .meta
645            .schema
646            .timestamp_column()
647            .context(UnexpectedSnafu {
648                violated: "Table must have a timestamp column",
649            })?;
650
651        let time_unit = time_index_dt
652            .data_type
653            .as_timestamp()
654            .with_context(|| UnexpectedSnafu {
655                violated: format!(
656                    "Table {}'s time index column must be a timestamp type, found: {:?}",
657                    table_name, time_index_dt
658                ),
659            })?
660            .unit();
661
662        let start_column = ColumnSchema::new(
663            "range_start",
664            ConcreteDataType::timestamp_datatype(time_unit),
665            false,
666        );
667        let end_column = ColumnSchema::new(
668            "range_end",
669            ConcreteDataType::timestamp_datatype(time_unit),
670            false,
671        );
672        let mut time_ranges = Vec::with_capacity(sql_values_time_range.len());
673        for (start, end) in sql_values_time_range {
674            let start = common_sql::convert::sql_value_to_value(
675                &start_column,
676                start,
677                Some(&query_ctx.timezone()),
678                None,
679                false,
680            )
681            .context(SqlCommonSnafu)
682            .and_then(|v| {
683                if let datatypes::value::Value::Timestamp(t) = v {
684                    Ok(t)
685                } else {
686                    error::InvalidSqlSnafu {
687                        err_msg: format!("Expected a timestamp value, found {v:?}"),
688                    }
689                    .fail()
690                }
691            })?;
692
693            let end = common_sql::convert::sql_value_to_value(
694                &end_column,
695                end,
696                Some(&query_ctx.timezone()),
697                None,
698                false,
699            )
700            .context(SqlCommonSnafu)
701            .and_then(|v| {
702                if let datatypes::value::Value::Timestamp(t) = v {
703                    Ok(t)
704                } else {
705                    error::InvalidSqlSnafu {
706                        err_msg: format!("Expected a timestamp value, found {v:?}"),
707                    }
708                    .fail()
709                }
710            })?;
711            time_ranges.push((start, end));
712        }
713        Ok(time_ranges)
714    }
715
716    /// Returns the inserter for the statement executor.
717    pub(crate) fn inserter(&self) -> &InserterRef {
718        &self.inserter
719    }
720}
721
722fn to_copy_query_request(stmt: CopyQueryToArgument) -> Result<CopyQueryToRequest> {
723    let CopyQueryToArgument {
724        with,
725        connection,
726        location,
727    } = stmt;
728
729    Ok(CopyQueryToRequest {
730        location,
731        with: with.into_map(),
732        connection: connection.into_map(),
733    })
734}
735
736// Verifies time related format is valid
737fn verify_time_related_format(with: &OptionMap) -> Result<()> {
738    let time_format = with.get(common_datasource::file_format::TIME_FORMAT);
739    let date_format = with.get(common_datasource::file_format::DATE_FORMAT);
740    let timestamp_format = with.get(common_datasource::file_format::TIMESTAMP_FORMAT);
741    let file_format = with.get(common_datasource::file_format::FORMAT_TYPE);
742
743    if !matches!(file_format, Some(f) if f.eq_ignore_ascii_case("csv")) {
744        ensure!(
745            time_format.is_none() && date_format.is_none() && timestamp_format.is_none(),
746            error::TimestampFormatNotSupportedSnafu {
747                format: "<unknown>".to_string(),
748                file_format: file_format.unwrap_or_default(),
749            }
750        );
751    }
752
753    for (key, format_opt) in [
754        (common_datasource::file_format::TIME_FORMAT, time_format),
755        (common_datasource::file_format::DATE_FORMAT, date_format),
756        (
757            common_datasource::file_format::TIMESTAMP_FORMAT,
758            timestamp_format,
759        ),
760    ] {
761        if let Some(format) = format_opt {
762            chrono::format::strftime::StrftimeItems::new(format)
763                .parse()
764                .map_err(|_| error::InvalidCopyParameterSnafu { key, value: format }.build())?;
765        }
766    }
767
768    Ok(())
769}
770
771fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result<CopyTableRequest> {
772    let direction = match stmt {
773        CopyTable::To(_) => CopyDirection::Export,
774        CopyTable::From(_) => CopyDirection::Import,
775    };
776
777    let CopyTableArgument {
778        location,
779        connection,
780        with,
781        table_name,
782        limit,
783        ..
784    } = match stmt {
785        CopyTable::To(arg) => arg,
786        CopyTable::From(arg) => arg,
787    };
788    let (catalog_name, schema_name, table_name) =
789        table_idents_to_full_name(&table_name, &query_ctx)
790            .map_err(BoxedError::new)
791            .context(ExternalSnafu)?;
792
793    let timestamp_range = timestamp_range_from_option_map(&with, &query_ctx)?;
794
795    verify_time_related_format(&with)?;
796
797    let pattern = with
798        .get(common_datasource::file_format::FILE_PATTERN)
799        .map(|x| x.to_string());
800
801    Ok(CopyTableRequest {
802        catalog_name,
803        schema_name,
804        table_name,
805        location,
806        with: with.into_map(),
807        connection: connection.into_map(),
808        pattern,
809        direction,
810        timestamp_range,
811        limit,
812    })
813}
814
815/// Converts [CopyDatabaseArgument] to [CopyDatabaseRequest].
816/// This function extracts the necessary info including catalog/database name, time range, etc.
817fn to_copy_database_request(
818    arg: CopyDatabaseArgument,
819    query_ctx: &QueryContextRef,
820) -> Result<CopyDatabaseRequest> {
821    let (catalog_name, database_name) = idents_to_full_database_name(&arg.database_name, query_ctx)
822        .map_err(BoxedError::new)
823        .context(ExternalSnafu)?;
824    let time_range = timestamp_range_from_option_map(&arg.with, query_ctx)?;
825
826    Ok(CopyDatabaseRequest {
827        catalog_name,
828        schema_name: database_name,
829        location: arg.location,
830        with: arg.with.into_map(),
831        connection: arg.connection.into_map(),
832        time_range,
833    })
834}
835
836/// Extracts timestamp range from OptionMap with keys `start_time` and `end_time`.
837/// The timestamp ranges should be a valid timestamp string as defined in [Timestamp::from_str].
838/// The timezone used for conversion will respect that inside `query_ctx`.
839fn timestamp_range_from_option_map(
840    options: &OptionMap,
841    query_ctx: &QueryContextRef,
842) -> Result<Option<TimestampRange>> {
843    let start_timestamp = extract_timestamp(options, COPY_DATABASE_TIME_START_KEY, query_ctx)?;
844    let end_timestamp = extract_timestamp(options, COPY_DATABASE_TIME_END_KEY, query_ctx)?;
845    let time_range = match (start_timestamp, end_timestamp) {
846        (Some(start), Some(end)) => Some(TimestampRange::new(start, end).with_context(|| {
847            error::InvalidTimestampRangeSnafu {
848                start: start.to_iso8601_string(),
849                end: end.to_iso8601_string(),
850            }
851        })?),
852        (Some(start), None) => Some(TimestampRange::from_start(start)),
853        (None, Some(end)) => Some(TimestampRange::until_end(end, false)), // exclusive end
854        (None, None) => None,
855    };
856    Ok(time_range)
857}
858
859/// Extracts timestamp from a [HashMap<String, String>] with given key.
860fn extract_timestamp(
861    map: &OptionMap,
862    key: &str,
863    query_ctx: &QueryContextRef,
864) -> Result<Option<Timestamp>> {
865    map.get(key)
866        .map(|v| {
867            Timestamp::from_str(v, Some(&query_ctx.timezone()))
868                .map_err(|_| error::InvalidCopyParameterSnafu { key, value: v }.build())
869        })
870        .transpose()
871}
872
873fn idents_to_full_database_name(
874    obj_name: &ObjectName,
875    query_ctx: &QueryContextRef,
876) -> Result<(String, String)> {
877    match &obj_name.0[..] {
878        [database] => Ok((
879            query_ctx.current_catalog().to_owned(),
880            database.to_string_unquoted(),
881        )),
882        [catalog, database] => Ok((catalog.to_string_unquoted(), database.to_string_unquoted())),
883        _ => InvalidSqlSnafu {
884            err_msg: format!(
885                "expect database name to be <catalog>.<database>, <database>, found: {obj_name}",
886            ),
887        }
888        .fail(),
889    }
890}
891
892/// The [`Inserter`] implementation for the statement executor.
893pub struct InserterImpl {
894    statement_executor: StatementExecutorRef,
895    options: Option<InsertOptions>,
896}
897
898impl InserterImpl {
899    pub fn new(statement_executor: StatementExecutorRef, options: Option<InsertOptions>) -> Self {
900        Self {
901            statement_executor,
902            options,
903        }
904    }
905}
906
907#[async_trait::async_trait]
908impl Inserter for InserterImpl {
909    async fn insert_rows(
910        &self,
911        context: &client::inserter::Context<'_>,
912        requests: RowInsertRequests,
913    ) -> ClientResult<()> {
914        let mut ctx_builder = QueryContextBuilder::default()
915            .current_catalog(context.catalog.to_string())
916            .current_schema(context.schema.to_string());
917        if let Some(options) = self.options.as_ref() {
918            ctx_builder = ctx_builder
919                .set_extension(
920                    TTL_KEY.to_string(),
921                    format_duration(options.ttl).to_string(),
922                )
923                .set_extension(APPEND_MODE_KEY.to_string(), options.append_mode.to_string());
924        }
925        let query_ctx = ctx_builder.build().into();
926
927        self.statement_executor
928            .inserter()
929            .handle_row_inserts(
930                requests,
931                query_ctx,
932                self.statement_executor.as_ref(),
933                false,
934                false,
935            )
936            .await
937            .map_err(BoxedError::new)
938            .context(ClientExternalSnafu)
939            .map(|_| ())
940    }
941
942    fn set_options(&mut self, options: &InsertOptions) {
943        self.options = Some(*options);
944    }
945}
946
947#[cfg(test)]
948mod tests {
949    use std::assert_matches;
950    use std::collections::HashMap;
951
952    use common_time::range::TimestampRange;
953    use common_time::{Timestamp, Timezone};
954    use session::context::QueryContextBuilder;
955    use sql::statements::OptionMap;
956
957    use crate::error;
958    use crate::statement::copy_database::{
959        COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY,
960    };
961    use crate::statement::{timestamp_range_from_option_map, verify_time_related_format};
962
963    fn check_timestamp_range((start, end): (&str, &str)) -> error::Result<Option<TimestampRange>> {
964        let query_ctx = QueryContextBuilder::default()
965            .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
966            .build()
967            .into();
968        let map = OptionMap::from(
969            [
970                (COPY_DATABASE_TIME_START_KEY.to_string(), start.to_string()),
971                (COPY_DATABASE_TIME_END_KEY.to_string(), end.to_string()),
972            ]
973            .into_iter()
974            .collect::<HashMap<_, _>>(),
975        );
976        timestamp_range_from_option_map(&map, &query_ctx)
977    }
978
979    #[test]
980    fn test_timestamp_range_from_option_map() {
981        assert_eq!(
982            Some(
983                TimestampRange::new(
984                    Timestamp::new_second(1649635200),
985                    Timestamp::new_second(1649664000),
986                )
987                .unwrap(),
988            ),
989            check_timestamp_range(("2022-04-11 08:00:00", "2022-04-11 16:00:00"),).unwrap()
990        );
991
992        assert_matches!(
993            check_timestamp_range(("2022-04-11 08:00:00", "2022-04-11 07:00:00")).unwrap_err(),
994            error::Error::InvalidTimestampRange { .. }
995        );
996    }
997
998    #[test]
999    fn test_verify_timestamp_format() {
1000        let map = OptionMap::from(
1001            [
1002                (
1003                    common_datasource::file_format::TIMESTAMP_FORMAT.to_string(),
1004                    "%Y-%m-%d %H:%M:%S".to_string(),
1005                ),
1006                (
1007                    common_datasource::file_format::FORMAT_TYPE.to_string(),
1008                    "csv".to_string(),
1009                ),
1010            ]
1011            .into_iter()
1012            .collect::<HashMap<_, _>>(),
1013        );
1014        assert!(verify_time_related_format(&map).is_ok());
1015
1016        let map = OptionMap::from(
1017            [
1018                (
1019                    common_datasource::file_format::TIMESTAMP_FORMAT.to_string(),
1020                    "%Y-%m-%d %H:%M:%S".to_string(),
1021                ),
1022                (
1023                    common_datasource::file_format::FORMAT_TYPE.to_string(),
1024                    "json".to_string(),
1025                ),
1026            ]
1027            .into_iter()
1028            .collect::<HashMap<_, _>>(),
1029        );
1030
1031        assert_matches!(
1032            verify_time_related_format(&map).unwrap_err(),
1033            error::Error::TimestampFormatNotSupported { .. }
1034        );
1035        let map = OptionMap::from(
1036            [
1037                (
1038                    common_datasource::file_format::TIMESTAMP_FORMAT.to_string(),
1039                    "%111112".to_string(),
1040                ),
1041                (
1042                    common_datasource::file_format::FORMAT_TYPE.to_string(),
1043                    "csv".to_string(),
1044                ),
1045            ]
1046            .into_iter()
1047            .collect::<HashMap<_, _>>(),
1048        );
1049
1050        assert_matches!(
1051            verify_time_related_format(&map).unwrap_err(),
1052            error::Error::InvalidCopyParameter { .. }
1053        );
1054    }
1055}
operator/statement.rs

operator/
statement.rs
