feat: copy option 'pattern' support variable. (#16525)

* feat: copy option 'pattern' support variable.

* refactor: pattern use type LiteralStringOrVariable.

* fix

* fix

src/query/ast/src/ast/statements/copy.rs

@@ -54,7 +54,7 @@ pub struct CopyIntoTableStmt {
 
     // files to load
     pub files: Option<Vec<String>>,
-    pub pattern: Option<String>,
+    pub pattern: Option<LiteralStringOrVariable>,
     pub force: bool,
 
     // copy options
@@ -110,7 +110,7 @@ impl Display for CopyIntoTableStmt {
         }
 
         if let Some(pattern) = &self.pattern {
-            write!(f, " PATTERN = '{}'", pattern)?;
+            write!(f, " PATTERN = {}", pattern)?;
         }
 
         if !self.file_format.is_empty() {
@@ -438,9 +438,34 @@ impl Display for FileLocation {
     }
 }
 
+/// Used when we want to allow use variable for options etc.
+/// Other expr is not necessary, because
+/// 1. we can always create a variable that can be used directly.
+/// 2. columns can not be referred.
+///
+/// Can extend to all type of Literals if needed later.
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum LiteralStringOrVariable {
+    Literal(String),
+    Variable(String),
+}
+
+impl Display for LiteralStringOrVariable {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            LiteralStringOrVariable::Literal(s) => {
+                write!(f, "'{s}'")
+            }
+            LiteralStringOrVariable::Variable(s) => {
+                write!(f, "${s}")
+            }
+        }
+    }
+}
+
 pub enum CopyIntoTableOption {
     Files(Vec<String>),
-    Pattern(String),
+    Pattern(LiteralStringOrVariable),
     FileFormat(FileFormatOptions),
     ValidationMode(String),
     SizeLimit(usize),

src/query/ast/src/ast/statements/stage.rs

@@ -24,6 +24,7 @@ use crate::ast::write_comma_separated_string_list;
 use crate::ast::write_comma_separated_string_map;
 use crate::ast::CreateOption;
 use crate::ast::FileFormatOptions;
+use crate::ast::LiteralStringOrVariable;
 use crate::ast::UriLocation;
 
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
@@ -80,10 +81,10 @@ impl Display for CreateStageStmt {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum SelectStageOption {
     Files(Vec<String>),
-    Pattern(String),
+    Pattern(LiteralStringOrVariable),
     FileFormat(String),
     Connection(BTreeMap<String, String>),
 }
@@ -103,10 +104,10 @@ impl SelectStageOptions {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Default, Drive, DriveMut)]
+#[derive(Debug, Clone, PartialEq, Default, Drive, DriveMut)]
 pub struct SelectStageOptions {
     pub files: Option<Vec<String>>,
-    pub pattern: Option<String>,
+    pub pattern: Option<LiteralStringOrVariable>,
     pub file_format: Option<String>,
     pub connection: BTreeMap<String, String>,
 }
@@ -150,7 +151,7 @@ impl Display for SelectStageOptions {
         }
 
         if let Some(pattern) = self.pattern.as_ref() {
-            write!(f, " PATTERN => '{}',", pattern)?;
+            write!(f, " PATTERN => {},", pattern)?;
         }
 
         if !self.connection.is_empty() {

src/query/ast/src/parser/copy.rs

@@ -23,6 +23,7 @@ use crate::ast::CopyIntoLocationStmt;
 use crate::ast::CopyIntoTableOption;
 use crate::ast::CopyIntoTableSource;
 use crate::ast::CopyIntoTableStmt;
+use crate::ast::LiteralStringOrVariable;
 use crate::ast::Statement;
 use crate::ast::Statement::CopyIntoLocation;
 use crate::parser::common::comma_separated_list0;
@@ -138,14 +139,21 @@ pub fn copy_into(i: Input) -> IResult<Statement> {
     )(i)
 }
 
+pub fn literal_string_or_variable(i: Input) -> IResult<LiteralStringOrVariable> {
+    alt((
+        map(literal_string, LiteralStringOrVariable::Literal),
+        map(variable_ident, LiteralStringOrVariable::Variable),
+    ))(i)
+}
+
 fn copy_into_table_option(i: Input) -> IResult<CopyIntoTableOption> {
     alt((
         map(
             rule! { FILES ~ "=" ~ "(" ~ #comma_separated_list0(literal_string) ~ ")" },
             |(_, _, _, files, _)| CopyIntoTableOption::Files(files),
         ),
         map(
-            rule! { PATTERN ~ "=" ~ #literal_string },
+            rule! { PATTERN ~ ^"=" ~ ^#literal_string_or_variable },
             |(_, _, pattern)| CopyIntoTableOption::Pattern(pattern),
         ),
         map(rule! { #file_format_clause }, |options| {

src/query/ast/src/parser/expr.rs

@@ -34,6 +34,7 @@ use crate::parser::query::*;
 use crate::parser::token::*;
 use crate::parser::Error;
 use crate::parser::ErrorKind;
+use crate::Span;
 
 pub fn expr(i: Input) -> IResult<Expr> {
     context("expression", subexpr(0))(i)
@@ -643,20 +644,10 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                 span: transform_span(elem.span.tokens),
                 name,
             },
-            ExprElement::VariableAccess(name) => Expr::FunctionCall {
-                span: transform_span(elem.span.tokens),
-                func: FunctionCall {
-                    distinct: false,
-                    name: Identifier::from_name(transform_span(elem.span.tokens), "getvariable"),
-                    args: vec![Expr::Literal {
-                        span: transform_span(elem.span.tokens),
-                        value: Literal::String(name),
-                    }],
-                    params: vec![],
-                    window: None,
-                    lambda: None,
-                },
-            },
+            ExprElement::VariableAccess(name) => {
+                let span = transform_span(elem.span.tokens);
+                make_func_get_variable(span, name)
+            }
             _ => unreachable!(),
         };
         Ok(expr)
@@ -1843,3 +1834,20 @@ pub fn parse_uint(text: &str, radix: u32) -> Result<Literal, ErrorKind> {
         })
     }
 }
+
+pub(crate) fn make_func_get_variable(span: Span, name: String) -> Expr {
+    Expr::FunctionCall {
+        span,
+        func: FunctionCall {
+            distinct: false,
+            name: Identifier::from_name(span, "getvariable"),
+            args: vec![Expr::Literal {
+                span,
+                value: Literal::String(name),
+            }],
+            params: vec![],
+            window: None,
+            lambda: None,
+        },
+    }
+}

src/query/ast/src/parser/stage.rs

@@ -24,6 +24,7 @@ use crate::ast::FileLocation;
 use crate::ast::SelectStageOption;
 use crate::ast::UriLocation;
 use crate::parser::common::*;
+use crate::parser::copy::literal_string_or_variable;
 use crate::parser::expr::*;
 use crate::parser::input::Input;
 use crate::parser::token::*;
@@ -256,7 +257,7 @@ pub fn select_stage_option(i: Input) -> IResult<SelectStageOption> {
             |(_, _, _, files, _)| SelectStageOption::Files(files),
         ),
         map(
-            rule! { PATTERN ~ ^"=>" ~ ^#literal_string },
+            rule! { PATTERN ~ ^"=>" ~ ^#literal_string_or_variable },
             |(_, _, pattern)| SelectStageOption::Pattern(pattern),
         ),
         map(

src/query/ast/tests/it/testdata/stmt.txt

@@ -9279,7 +9279,9 @@ Query(
                         options: SelectStageOptions {
                             files: None,
                             pattern: Some(
-                                "[.]*parquet",
+                                Literal(
+                                    "[.]*parquet",
+                                ),
                             ),
                             file_format: Some(
                                 "tsv",
@@ -16830,7 +16832,9 @@ Query(
                         options: SelectStageOptions {
                             files: None,
                             pattern: Some(
-                                "*.parquet",
+                                Literal(
+                                    "*.parquet",
+                                ),
                             ),
                             file_format: None,
                             connection: {

src/query/sql/src/planner/binder/bind_table_reference/bind_location.rs

@@ -56,9 +56,14 @@ impl Binder {
                     _ => databend_common_base::runtime::block_on(self.ctx.get_file_format(f))?,
                 }
             }
+            let pattern = match &options.pattern {
+                None => None,
+                Some(pattern) => Some(Self::resolve_copy_pattern(self.ctx.clone(), pattern)?),
+            };
+
             let files_info = StageFilesInfo {
                 path,
-                pattern: options.pattern.clone(),
+                pattern,
                 files: options.files.clone(),
             };
             let table_ctx = self.ctx.clone();

src/query/sql/src/planner/binder/copy_into_table.rs

@@ -26,6 +26,7 @@ use databend_common_ast::ast::Hint;
 use databend_common_ast::ast::HintItem;
 use databend_common_ast::ast::Identifier;
 use databend_common_ast::ast::Literal;
+use databend_common_ast::ast::LiteralStringOrVariable;
 use databend_common_ast::ast::Query;
 use databend_common_ast::ast::SelectTarget;
 use databend_common_ast::ast::SetExpr;
@@ -42,6 +43,7 @@ use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::infer_table_schema;
+use databend_common_expression::shrink_scalar;
 use databend_common_expression::types::DataType;
 use databend_common_expression::DataSchema;
 use databend_common_expression::DataSchemaRef;
@@ -111,6 +113,26 @@ impl<'a> Binder {
         }
     }
 
+    pub(crate) fn resolve_copy_pattern(
+        ctx: Arc<dyn TableContext>,
+        pattern: &LiteralStringOrVariable,
+    ) -> Result<String> {
+        match pattern {
+            LiteralStringOrVariable::Literal(s) => Ok(s.clone()),
+            LiteralStringOrVariable::Variable(var_name) => {
+                let var_value = ctx.get_variable(var_name).unwrap_or(Scalar::Null);
+                let var_value = shrink_scalar(var_value);
+                if let Scalar::String(s) = var_value {
+                    Ok(s)
+                } else {
+                    Err(ErrorCode::BadArguments(format!(
+                        "invalid pattern expr: {var_value}"
+                    )))
+                }
+            }
+        }
+    }
+
     async fn bind_copy_into_table_common(
         &mut self,
         bind_context: &mut BindContext,
@@ -136,10 +158,15 @@ impl<'a> Binder {
         let (mut stage_info, path) = resolve_file_location(self.ctx.as_ref(), location).await?;
         self.apply_copy_into_table_options(stmt, &mut stage_info)
             .await?;
+        let pattern = match &stmt.pattern {
+            None => None,
+            Some(pattern) => Some(Self::resolve_copy_pattern(self.ctx.clone(), pattern)?),
+        };
+
         let files_info = StageFilesInfo {
             path,
             files: stmt.files.clone(),
-            pattern: stmt.pattern.clone(),
+            pattern,
         };
         let required_values_schema: DataSchemaRef = Arc::new(
             match &stmt.dst_columns {

tests/sqllogictests/suites/stage/options/pattern_variable.test

@@ -0,0 +1,31 @@
+statement ok
+create or replace table t1(c1 int, c2 string);
+
+statement ok
+set variable pt='it.csv';
+
+query 
+select $1 from @data/csv/ (pattern => $pt, file_format=>'csv')
+----
+1
+2
+
+query 
+copy into t1 from @data/csv/ pattern= $pt file_format=(type=csv)
+----
+csv/it.csv 2 0 NULL NULL
+
+query
+select * from t1
+----
+1 b
+2 d
+
+statement ok
+unset variable pt;
+
+query error 1006
+copy into t1 from @data/csv/ pattern= $p file_format=(type=csv)
+
+query error 1005
+copy into t1 from @data/csv/ pattern=x file_format=(type=csv)