Task LAV-1995: parse external functions (apache#2446)

Author: sabir-akhadov-localstack

src/ast/ddl.rs

@@ -4102,6 +4102,101 @@ impl fmt::Display for FunctionReturnType {
     }
 }
 
+/// A `HEADERS` entry in a Snowflake external function definition.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ExternalFunctionHeader {
+    /// Header-name string literal.
+    pub name: ValueWithSpan,
+    /// Header-value string literal.
+    pub value: ValueWithSpan,
+}
+
+impl fmt::Display for ExternalFunctionHeader {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{} = {}", self.name, self.value)
+    }
+}
+
+/// Compression mode for a Snowflake external function request.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ExternalFunctionCompression {
+    /// No compression.
+    None,
+    /// Gzip compression.
+    Gzip,
+    /// Deflate compression.
+    Deflate,
+    /// Compression selected automatically.
+    Auto,
+}
+
+impl fmt::Display for ExternalFunctionCompression {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::None => write!(f, "NONE"),
+            Self::Gzip => write!(f, "GZIP"),
+            Self::Deflate => write!(f, "DEFLATE"),
+            Self::Auto => write!(f, "AUTO"),
+        }
+    }
+}
+
+/// Snowflake-specific parameters for `CREATE EXTERNAL FUNCTION`.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ExternalFunctionParams {
+    /// Whether the return type has a `NOT NULL` annotation.
+    pub return_not_null: bool,
+    /// API integration used to invoke the remote endpoint.
+    pub api_integration: ObjectName,
+    /// Custom HTTP headers.
+    pub headers: Option<Vec<ExternalFunctionHeader>>,
+    /// Context function names forwarded as headers.
+    pub context_headers: Option<Vec<Ident>>,
+    /// Maximum rows per request batch.
+    pub max_batch_rows: Option<u64>,
+    /// Request compression mode.
+    pub compression: Option<ExternalFunctionCompression>,
+    /// Optional request translator function.
+    pub request_translator: Option<ObjectName>,
+    /// Optional response translator function.
+    pub response_translator: Option<ObjectName>,
+}
+
+impl fmt::Display for ExternalFunctionParams {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "API_INTEGRATION = {}", self.api_integration)?;
+        if let Some(headers) = &self.headers {
+            write!(f, " HEADERS = ({})", display_comma_separated(headers))?;
+        }
+        if let Some(context_headers) = &self.context_headers {
+            write!(
+                f,
+                " CONTEXT_HEADERS = ({})",
+                display_comma_separated(context_headers)
+            )?;
+        }
+        if let Some(max_batch_rows) = self.max_batch_rows {
+            write!(f, " MAX_BATCH_ROWS = {max_batch_rows}")?;
+        }
+        if let Some(compression) = &self.compression {
+            write!(f, " COMPRESSION = {compression}")?;
+        }
+        if let Some(request_translator) = &self.request_translator {
+            write!(f, " REQUEST_TRANSLATOR = {request_translator}")?;
+        }
+        if let Some(response_translator) = &self.response_translator {
+            write!(f, " RESPONSE_TRANSLATOR = {response_translator}")?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -4185,16 +4280,23 @@ pub struct CreateFunction {
     /// ```
     /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#create_a_remote_function)
     pub remote_connection: Option<ObjectName>,
+    /// Snowflake external-function parameters.
+    pub external_params: Option<ExternalFunctionParams>,
 }
 
 impl fmt::Display for CreateFunction {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
-            "CREATE {or_alter}{or_replace}{temp}{secure}{data_metric}FUNCTION {if_not_exists}{name}",
+            "CREATE {or_alter}{or_replace}{temp}{secure}{external}{data_metric}FUNCTION {if_not_exists}{name}",
             name = self.name,
             temp = if self.temporary { "TEMPORARY " } else { "" },
             secure = if self.secure { "SECURE " } else { "" },
+            external = if self.external_params.is_some() {
+                "EXTERNAL "
+            } else {
+                ""
+            },
             data_metric = if self.data_metric { "DATA METRIC " } else { "" },
             or_alter = if self.or_alter { "OR ALTER " } else { "" },
             or_replace = if self.or_replace { "OR REPLACE " } else { "" },
@@ -4210,6 +4312,32 @@ impl fmt::Display for CreateFunction {
         if let Some(return_type) = &self.return_type {
             write!(f, " RETURNS {return_type}")?;
         }
+        if let Some(external) = &self.external_params {
+            if external.return_not_null {
+                write!(f, " NOT NULL")?;
+            }
+            if let Some(called_on_null) = &self.called_on_null {
+                write!(f, " {called_on_null}")?;
+            }
+            if let Some(behavior) = &self.behavior {
+                write!(f, " {behavior}")?;
+            }
+            if let Some(options) = &self.options {
+                for option in options {
+                    write!(f, " {option}")?;
+                }
+            }
+            write!(f, " {external}")?;
+            if let Some(CreateFunctionBody::AsBeforeOptions { body, link_symbol }) =
+                &self.function_body
+            {
+                write!(f, " AS {body}")?;
+                if let Some(link_symbol) = link_symbol {
+                    write!(f, ", {link_symbol}")?;
+                }
+            }
+            return Ok(());
+        }
         if let Some(determinism_specifier) = &self.determinism_specifier {
             write!(f, " {determinism_specifier}")?;
         }

src/ast/mod.rs

@@ -77,7 +77,8 @@ pub use self::ddl::{
     CreateOperatorClass, CreateOperatorFamily, CreatePolicy, CreatePolicyCommand, CreatePolicyType,
     CreateTable, CreateTrigger, CreateView, Deduplicate, DeferrableInitial, DistStyle,
     DropBehavior, DropExtension, DropFunction, DropOperator, DropOperatorClass, DropOperatorFamily,
-    DropOperatorSignature, DropPolicy, DropTrigger, ExternalTablePartitionColumn, ForValues,
+    DropOperatorSignature, DropPolicy, DropTrigger, ExternalFunctionCompression,
+    ExternalFunctionHeader, ExternalFunctionParams, ExternalTablePartitionColumn, ForValues,
     FunctionReturnType, GeneratedAs, GeneratedExpressionMode, IdentityParameters, IdentityProperty,
     IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder, IndexColumn,
     IndexOption, IndexType, KeyOrIndexDisplay, Msck, NullsDistinctOption, OperatorArgTypes,

src/ast/spans.rs

@@ -18,8 +18,9 @@
 use crate::{
     ast::{
         ddl::AlterSchema, query::SelectItemQualifiedWildcardKind, AlterSchemaOperation, AlterTable,
-        ColumnOptions, CreateOperator, CreateOperatorClass, CreateOperatorFamily, CreateView,
-        ExportData, Owner, TypedString,
+        ColumnOptions, CreateFunction, CreateFunctionBody, CreateOperator, CreateOperatorClass,
+        CreateOperatorFamily, CreateView, ExportData, ExternalFunctionHeader,
+        ExternalFunctionParams, Owner, TypedString,
     },
     tokenizer::TokenWithSpan,
 };
@@ -469,7 +470,7 @@ impl Spanned for Statement {
             Statement::Rollback { .. } => Span::empty(),
             Statement::CreateSchema { .. } => Span::empty(),
             Statement::CreateDatabase { .. } => Span::empty(),
-            Statement::CreateFunction { .. } => Span::empty(),
+            Statement::CreateFunction(create_function) => create_function.span(),
             Statement::CreateDomain { .. } => Span::empty(),
             Statement::CreateTrigger { .. } => Span::empty(),
             Statement::DropTrigger { .. } => Span::empty(),
@@ -1233,6 +1234,81 @@ impl Spanned for SqlOption {
     }
 }
 
+impl Spanned for ExternalFunctionHeader {
+    fn span(&self) -> Span {
+        self.name.span().union(&self.value.span())
+    }
+}
+
+impl Spanned for ExternalFunctionParams {
+    fn span(&self) -> Span {
+        union_spans(
+            iter::once(self.api_integration.span())
+                .chain(
+                    self.headers
+                        .iter()
+                        .flatten()
+                        .map(Spanned::span),
+                )
+                .chain(
+                    self.context_headers
+                        .iter()
+                        .flatten()
+                        .map(|header| header.span),
+                )
+                .chain(
+                    self.request_translator
+                        .iter()
+                        .map(|translator| translator.span()),
+                )
+                .chain(
+                    self.response_translator
+                        .iter()
+                        .map(|translator| translator.span()),
+                ),
+        )
+    }
+}
+
+impl Spanned for CreateFunctionBody {
+    fn span(&self) -> Span {
+        match self {
+            CreateFunctionBody::AsBeforeOptions { body, link_symbol } => {
+                body.span().union_opt(&link_symbol.as_ref().map(Spanned::span))
+            }
+            CreateFunctionBody::AsAfterOptions(body)
+            | CreateFunctionBody::Return(body)
+            | CreateFunctionBody::AsReturnExpr(body) => body.span(),
+            CreateFunctionBody::AsBeginEnd(body) => body.span(),
+            CreateFunctionBody::AsReturnSelect(body) => body.span(),
+        }
+    }
+}
+
+impl Spanned for CreateFunction {
+    fn span(&self) -> Span {
+        union_spans(
+            iter::once(self.name.span())
+                .chain(
+                    self.args
+                        .iter()
+                        .flatten()
+                        .flat_map(|argument| {
+                            [
+                                argument.name.as_ref().map(|name| name.span),
+                                argument.default_expr.as_ref().map(Spanned::span),
+                            ]
+                            .into_iter()
+                            .flatten()
+                        }),
+                )
+                .chain(self.options.iter().flatten().map(Spanned::span))
+                .chain(self.function_body.iter().map(Spanned::span))
+                .chain(self.external_params.iter().map(Spanned::span)),
+        )
+    }
+}
+
 /// # partial span
 ///
 /// Missing spans: