polyglot-sql 0.4.2__tar.gz → 0.4.3__tar.gz

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/Cargo.lock

@@ -605,7 +605,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql"
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
  "criterion",
  "once_cell",
@@ -621,7 +621,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-ffi"
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
  "cbindgen",
  "polyglot-sql",
@@ -631,11 +631,11 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-function-catalogs"
-version = "0.4.2"
+version = "0.4.3"
 
 [[package]]
 name = "polyglot-sql-python"
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
  "polyglot-sql",
  "pyo3",
@@ -646,7 +646,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-wasm"
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
  "console_error_panic_hook",
  "js-sys",

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/Cargo.toml

@@ -6,7 +6,7 @@ exclude = [
 ]
 
 [workspace.package]
-version = "0.4.2"
+version = "0.4.3"
 edition = "2021"
 license = "MIT"
 authors = ["polyglot contributors"]

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polyglot-sql
-Version: 0.4.2
+Version: 0.4.3
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/Cargo.toml

@@ -106,7 +106,7 @@ thiserror = { workspace = true }
 unicode-segmentation = { workspace = true }
 stacker = { version = "0.1", optional = true }
 ts-rs = { version = "12.0", features = ["serde-compat"], optional = true }
-polyglot-sql-function-catalogs = { path = "../polyglot-sql-function-catalogs", version = "0.4.2", optional = true, default-features = false }
+polyglot-sql-function-catalogs = { path = "../polyglot-sql-function-catalogs", version = "0.4.3", optional = true, default-features = false }
 
 [dev-dependencies]
 pretty_assertions = "1.4"

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/dialects/mod.rs

@@ -2918,6 +2918,31 @@ where
             Expression::CreateTask(ct)
         }
 
+        // Prepare: recurse into the prepared statement body
+        Expression::Prepare(mut prepare) => {
+            prepare.statement = transform_recursive(prepare.statement, transform_fn)?;
+            Expression::Prepare(prepare)
+        }
+
+        // Execute: recurse into procedure/prepared name and argument values
+        Expression::Execute(mut execute) => {
+            execute.this = transform_recursive(execute.this, transform_fn)?;
+            execute.arguments = execute
+                .arguments
+                .into_iter()
+                .map(|argument| transform_recursive(argument, transform_fn))
+                .collect::<Result<Vec<_>>>()?;
+            execute.parameters = execute
+                .parameters
+                .into_iter()
+                .map(|mut parameter| {
+                    parameter.value = transform_recursive(parameter.value, transform_fn)?;
+                    Ok(parameter)
+                })
+                .collect::<Result<Vec<_>>>()?;
+            Expression::Execute(execute)
+        }
+
         // CreateProcedure: recurse into body expressions
         Expression::CreateProcedure(mut cp) => {
             if let Some(body) = cp.body.take() {

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/expressions.rs

@@ -587,6 +587,8 @@ pub enum Expression {
     // Transaction and other commands
     Command(Box<Command>),
     Kill(Box<Kill>),
+    /// PREPARE statement (PostgreSQL/generic prepared statement definition)
+    Prepare(Box<PrepareStatement>),
     /// EXEC/EXECUTE statement (TSQL stored procedure call)
     Execute(Box<ExecuteStatement>),
 
@@ -1197,6 +1199,7 @@ impl Expression {
             | Expression::Describe(_)
             | Expression::Show(_)
             | Expression::Kill(_)
+            | Expression::Prepare(_)
             | Expression::Execute(_)
             | Expression::Declare(_)
             | Expression::Refresh(_)
@@ -2235,6 +2238,7 @@ impl Expression {
             Expression::Command(_) => "command",
             Expression::TryCatch(_) => "try_catch",
             Expression::Kill(_) => "kill",
+            Expression::Prepare(_) => "prepare",
             Expression::Execute(_) => "execute",
             Expression::Raw(_) => "raw",
             Expression::CreateTask(_) => "create_task",
@@ -6008,6 +6012,20 @@ pub struct Command {
     pub this: String,
 }
 
+/// PREPARE statement (PostgreSQL/generic prepared statement definition)
+/// Syntax: PREPARE name [(type, ...)] AS statement
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[cfg_attr(feature = "bindings", derive(TS))]
+pub struct PrepareStatement {
+    /// The prepared statement name.
+    pub name: Identifier,
+    /// Optional PostgreSQL parameter type list.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub parameter_types: Vec<DataType>,
+    /// The statement to execute when the prepared statement is invoked.
+    pub statement: Expression,
+}
+
 /// T-SQL TRY/CATCH block.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[cfg_attr(feature = "bindings", derive(TS))]
@@ -6030,6 +6048,12 @@ pub struct ExecuteStatement {
     /// Named parameters: @param=value pairs
     #[serde(default)]
     pub parameters: Vec<ExecuteParameter>,
+    /// Positional prepared statement arguments, used by PostgreSQL EXECUTE name(...).
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub arguments: Vec<Expression>,
+    /// Whether this statement represents PostgreSQL-style prepared statement execution.
+    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
+    pub prepared: bool,
     /// Trailing clause text (e.g. WITH RESULT SETS ((...)))
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub suffix: Option<String>,

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/generator.rs

@@ -3357,10 +3357,24 @@ impl Generator {
                 self.generate_expression(&kill.this)?;
                 Ok(())
             }
+            Expression::Prepare(prepare) => self.generate_prepare(prepare),
             Expression::Execute(exec) => {
                 self.write_keyword("EXECUTE");
                 self.write_space();
                 self.generate_expression(&exec.this)?;
+                if exec.prepared {
+                    if !exec.arguments.is_empty() {
+                        self.write("(");
+                        for (i, argument) in exec.arguments.iter().enumerate() {
+                            if i > 0 {
+                                self.write(", ");
+                            }
+                            self.generate_expression(argument)?;
+                        }
+                        self.write(")");
+                    }
+                    return Ok(());
+                }
                 for (i, param) in exec.parameters.iter().enumerate() {
                     if i == 0 {
                         self.write_space();
@@ -15616,6 +15630,28 @@ impl Generator {
         Ok(())
     }
 
+    fn generate_prepare(&mut self, prepare: &PrepareStatement) -> Result<()> {
+        self.write_keyword("PREPARE");
+        self.write_space();
+        self.generate_identifier(&prepare.name)?;
+
+        if !prepare.parameter_types.is_empty() {
+            self.write(" (");
+            for (i, data_type) in prepare.parameter_types.iter().enumerate() {
+                if i > 0 {
+                    self.write(", ");
+                }
+                self.generate_data_type(data_type)?;
+            }
+            self.write(")");
+        }
+
+        self.write_space();
+        self.write_keyword("AS");
+        self.write_space();
+        self.generate_expression(&prepare.statement)
+    }
+
     /// Generate a pseudocolumn (Oracle ROWNUM, ROWID, LEVEL, etc.)
     /// Pseudocolumns should NEVER be quoted, as quoting breaks them in Oracle
     fn generate_pseudocolumn(&mut self, pc: &Pseudocolumn) -> Result<()> {

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/lineage.rs

@@ -120,7 +120,8 @@ pub fn lineage(
     trim_selects: bool,
 ) -> Result<LineageNode> {
     // Fast path: skip clone when there are no CTEs to expand
-    let has_with = matches!(sql, Expression::Select(s) if s.with.is_some());
+    let effective_sql = lineage_effective_expression(sql);
+    let has_with = matches!(effective_sql, Expression::Select(s) if s.with.is_some());
     if !has_with {
         return lineage_from_expression(column, sql, dialect, trim_selects);
     }
@@ -181,6 +182,7 @@ fn lineage_from_expression(
     dialect: Option<DialectType>,
     trim_selects: bool,
 ) -> Result<LineageNode> {
+    let sql = lineage_effective_expression(sql);
     let scope = build_scope(sql);
     to_node(
         ColumnRef::Name(column),
@@ -193,6 +195,13 @@ fn lineage_from_expression(
     )
 }
 
+fn lineage_effective_expression(sql: &Expression) -> &Expression {
+    match sql {
+        Expression::Prepare(prepare) => &prepare.statement,
+        _ => sql,
+    }
+}
+
 // ---------------------------------------------------------------------------
 // CTE star expansion
 // ---------------------------------------------------------------------------
@@ -222,6 +231,11 @@ fn normalize_cte_name(ident: &Identifier) -> String {
 /// case-insensitively (lowercased), while quoted names preserve their original case.
 /// This matches sqlglot's `normalize_identifiers` behavior.
 pub fn expand_cte_stars(expr: &mut Expression, schema: Option<&dyn Schema>) {
+    if let Expression::Prepare(prepare) = expr {
+        expand_cte_stars(&mut prepare.statement, schema);
+        return;
+    }
+
     let select = match expr {
         Expression::Select(s) => s,
         _ => return,
@@ -480,6 +494,14 @@ fn get_select_sources(select: &Select) -> Vec<SourceInfo> {
     let mut sources = Vec::new();
 
     fn extract_source(expr: &Expression) -> Option<SourceInfo> {
+        fn virtual_source_info(alias: &Identifier) -> SourceInfo {
+            SourceInfo {
+                alias: alias.name.clone(),
+                normalized: normalize_cte_name(alias),
+                fq_name: alias.name.clone(),
+            }
+        }
+
         match expr {
             Expression::Table(t) => {
                 let normalized = normalize_cte_name(&t.name);
@@ -513,6 +535,10 @@ fn get_select_sources(select: &Select) -> Vec<SourceInfo> {
                     fq_name,
                 })
             }
+            Expression::Unnest(u) => u.alias.as_ref().map(virtual_source_info),
+            Expression::Alias(a) if matches!(&a.this, Expression::Unnest(_)) => {
+                Some(virtual_source_info(&a.alias))
+            }
             Expression::Paren(p) => extract_source(&p.this),
             _ => None,
         }
@@ -971,6 +997,10 @@ fn source_names_from_from_join(scope: &Scope) -> Vec<String> {
             Expression::Subquery(subquery) => {
                 subquery.alias.as_ref().map(|alias| alias.name.clone())
             }
+            Expression::Unnest(unnest) => unnest.alias.as_ref().map(|alias| alias.name.clone()),
+            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
+                Some(alias.alias.name.clone())
+            }
             Expression::Paren(paren) => source_name(&paren.this),
             _ => None,
         }
@@ -2486,6 +2516,45 @@ mod tests {
         assert_eq!(node.name, "name");
     }
 
+    #[test]
+    fn test_lineage_bigquery_unnest_alias_source_issue_209() {
+        let expr = parse_one(
+            r#"
+SELECT date_val AS week_start
+FROM UNNEST(GENERATE_DATE_ARRAY('2024-01-01', '2024-12-31', INTERVAL 1 WEEK)) AS date_val
+"#,
+            DialectType::BigQuery,
+        )
+        .expect("parse");
+
+        let node = lineage("week_start", &expr, Some(DialectType::BigQuery), false)
+            .expect("lineage should resolve UNNEST alias as a source");
+        let child = node
+            .downstream
+            .first()
+            .expect("week_start should have downstream lineage");
+
+        assert_eq!(child.name, "date_val.date_val");
+        assert_eq!(child.source_name, "date_val");
+
+        let Expression::Column(column) = &child.expression else {
+            panic!(
+                "expected downstream column expression, got {:?}",
+                child.expression
+            );
+        };
+        assert_eq!(column.name.name, "date_val");
+        assert_eq!(
+            column.table.as_ref().map(|table| table.name.as_str()),
+            Some("date_val")
+        );
+        assert!(
+            matches!(&child.source, Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) && alias.alias.name == "date_val"),
+            "expected UNNEST source expression, got {:?}",
+            child.source
+        );
+    }
+
     #[test]
     fn test_lineage_with_schema_snowflake_datediff_date_part_issue_61() {
         let expr = parse_one(

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/openlineage.rs

@@ -363,6 +363,7 @@ fn analyze_statement(
     warnings: &mut Vec<OpenLineageWarning>,
 ) -> Result<StatementAnalysis> {
     match expr {
+        Expression::Prepare(prepare) => analyze_statement(&prepare.statement, options, warnings),
         Expression::Select(select) => {
             let output = if let Some(into) = &select.into {
                 dataset_from_expression(&into.this, options)?
@@ -914,6 +915,7 @@ fn collect_source_table(expr: &Expression, result: &mut Vec<(String, String)>) {
 
 fn leftmost_select(expr: &Expression) -> Option<&Select> {
     match expr {
+        Expression::Prepare(prepare) => leftmost_select(&prepare.statement),
         Expression::Select(select) => Some(select),
         Expression::Union(union) => leftmost_select(&union.left),
         Expression::Intersect(intersect) => leftmost_select(&intersect.left),
@@ -1103,6 +1105,19 @@ mod tests {
         assert_eq!(field.input_fields[0].transformations[0].subtype, "IDENTITY");
     }
 
+    #[test]
+    fn emits_column_lineage_for_prepared_statement_body() {
+        let result = openlineage_column_lineage(
+            "PREPARE leak AS SELECT id FROM sensitive_table WHERE id = $1",
+            &options(),
+        )
+        .expect("lineage");
+        let field = result.facet.fields.get("id").expect("field id");
+        assert_eq!(field.input_fields.len(), 1);
+        assert_eq!(field.input_fields[0].name, "sensitive_table");
+        assert_eq!(field.input_fields[0].field, "id");
+    }
+
     #[test]
     fn resolves_input_dataset_behind_table_alias() {
         let result = openlineage_column_lineage("SELECT o.total FROM orders o", &options())
@@ -1165,6 +1180,39 @@ SELECT json_data FROM transform_cte
         );
     }
 
+    #[test]
+    fn emits_bigquery_unnest_alias_column_lineage_issue209() {
+        let mut opts = options();
+        opts.dialect = DialectType::BigQuery;
+        opts.dataset_namespace = Some("bigquery://warehouse".to_string());
+        opts.output_dataset = Some(OpenLineageDatasetId::new(
+            "bigquery://warehouse",
+            "calendar",
+        ));
+
+        let result = openlineage_column_lineage(
+            r#"
+SELECT date_val AS week_start
+FROM UNNEST(GENERATE_DATE_ARRAY('2024-01-01', '2024-12-31', INTERVAL 1 WEEK)) AS date_val
+"#,
+            &opts,
+        )
+        .expect("lineage");
+        let field = result.facet.fields.get("week_start").expect("week_start");
+
+        assert_eq!(field.input_fields.len(), 1);
+        assert_eq!(field.input_fields[0].name, "date_val");
+        assert_eq!(field.input_fields[0].field, "date_val");
+        assert!(
+            result
+                .warnings
+                .iter()
+                .all(|warning| warning.code != "W_EMPTY_FIELD_LINEAGE"),
+            "did not expect empty-lineage warning, got {:?}",
+            result.warnings
+        );
+    }
+
     #[test]
     fn emits_aggregation_column_lineage() {
         let result =

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/parser.rs

@@ -1136,6 +1136,7 @@ impl Parser {
                 self.fallback_to_command(statement_start)
             }
             TokenType::Kill => self.parse_kill(),
+            TokenType::Prepare => self.parse_prepare(),
             TokenType::Execute => {
                 // ClickHouse: EXECUTE AS username statement → parse as command
                 if matches!(
@@ -23658,6 +23659,48 @@ impl Parser {
         Ok(Expression::Kill(Box::new(Kill { this, kind })))
     }
 
+    /// Parse PREPARE statement (PostgreSQL/generic prepared statement definition)
+    /// PREPARE name [(type, ...)] AS statement
+    fn parse_prepare(&mut self) -> Result<Expression> {
+        self.expect(TokenType::Prepare)?;
+
+        let name = self.expect_identifier_or_keyword_with_quoted()?;
+
+        let mut parameter_types = Vec::new();
+        if self.match_token(TokenType::LParen) {
+            if !self.check(TokenType::RParen) {
+                loop {
+                    parameter_types.push(self.parse_data_type()?);
+                    if !self.match_token(TokenType::Comma) {
+                        break;
+                    }
+                }
+            }
+            self.expect(TokenType::RParen)?;
+        }
+
+        self.expect(TokenType::As)?;
+        let statement = self.parse_statement()?;
+
+        Ok(Expression::Prepare(Box::new(PrepareStatement {
+            name,
+            parameter_types,
+            statement,
+        })))
+    }
+
+    fn is_postgres_prepared_execute_dialect(&self) -> bool {
+        matches!(
+            self.config.dialect,
+            None | Some(crate::dialects::DialectType::Generic)
+                | Some(crate::dialects::DialectType::PostgreSQL)
+                | Some(crate::dialects::DialectType::Redshift)
+                | Some(crate::dialects::DialectType::Materialize)
+                | Some(crate::dialects::DialectType::RisingWave)
+                | Some(crate::dialects::DialectType::CockroachDB)
+        )
+    }
+
     /// Parse EXEC/EXECUTE statement (TSQL stored procedure call)
     /// EXEC [schema.]procedure_name [@param=value, ...]
     fn parse_execute(&mut self) -> Result<Expression> {
@@ -23680,6 +23723,28 @@ impl Parser {
             Expression::Table(Box::new(proc_name))
         };
 
+        if self.is_postgres_prepared_execute_dialect() && self.check(TokenType::LParen) {
+            let arguments = self.parse_wrapped_csv_expressions()?;
+            return Ok(Expression::Execute(Box::new(ExecuteStatement {
+                this,
+                parameters: Vec::new(),
+                arguments,
+                prepared: true,
+                suffix: None,
+            })));
+        }
+        if self.is_postgres_prepared_execute_dialect()
+            && (self.check(TokenType::Semicolon) || self.is_at_end())
+        {
+            return Ok(Expression::Execute(Box::new(ExecuteStatement {
+                this,
+                parameters: Vec::new(),
+                arguments: Vec::new(),
+                prepared: true,
+                suffix: None,
+            })));
+        }
+
         // Parse optional parameters: @param=value [OUTPUT], ...
         let mut parameters = Vec::new();
 
@@ -23743,6 +23808,8 @@ impl Parser {
         Ok(Expression::Execute(Box::new(ExecuteStatement {
             this,
             parameters,
+            arguments: Vec::new(),
+            prepared: false,
             suffix,
         })))
     }

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/scope.rs

@@ -533,6 +533,9 @@ fn collect_columns(expr: &Expression, columns: &mut Vec<ColumnRef>) {
         Expression::Exists(_) | Expression::Subquery(_) => {
             // These create their own scopes - don't collect from here
         }
+        Expression::Prepare(prepare) => {
+            collect_columns(&prepare.statement, columns);
+        }
         _ => {
             // For other expressions, we might need to add more cases
         }
@@ -551,6 +554,9 @@ pub fn build_scope(expression: &Expression) -> Scope {
 
 fn build_scope_impl(expression: &Expression, current_scope: &mut Scope) {
     match expression {
+        Expression::Prepare(prepare) => {
+            build_scope_impl(&prepare.statement, current_scope);
+        }
         Expression::Select(select) => {
             // Process CTEs first
             if let Some(with) = &select.with {
@@ -674,6 +680,14 @@ fn add_table_to_scope(expr: &Expression, scope: &mut Scope) {
             scope.add_source(name.clone(), expr.clone(), true);
             scope.derived_table_scopes.push(derived_scope);
         }
+        Expression::Unnest(unnest) => {
+            if let Some(alias) = &unnest.alias {
+                scope.add_source(alias.name.clone(), expr.clone(), false);
+            }
+        }
+        Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
+            scope.add_source(alias.alias.name.clone(), expr.clone(), false);
+        }
         Expression::Paren(paren) => {
             add_table_to_scope(&paren.this, scope);
         }
@@ -848,6 +862,9 @@ impl<'a> WalkInScopeIter<'a> {
         let mut children = Vec::new();
 
         match expr {
+            Expression::Prepare(prepare) => {
+                children.push(&prepare.statement);
+            }
             Expression::Select(select) => {
                 // Walk SELECT expressions
                 for e in &select.expressions {
@@ -1117,6 +1134,7 @@ pub fn traverse_scope(expression: &Expression) -> Vec<Scope> {
         | Expression::Union(_)
         | Expression::Intersect(_)
         | Expression::Except(_)
+        | Expression::Prepare(_)
         | Expression::CreateTable(_) => {
             let root = build_scope(expression);
             root.traverse().into_iter().cloned().collect()

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/tokens.rs

@@ -377,6 +377,7 @@ pub enum TokenType {
     Command,
     Comment,
     Commit,
+    Prepare,
     Preserve,
     Connect,
     ConnectBy,
@@ -898,6 +899,7 @@ impl TokenType {
                 | TokenType::Qualify
                 | TokenType::Returning
                 | TokenType::Language
+                | TokenType::Prepare
                 | TokenType::Preserve
                 | TokenType::Savepoint
                 | TokenType::Rollback
@@ -1171,6 +1173,7 @@ static DEFAULT_KEYWORDS: LazyLock<HashMap<String, TokenType>> = LazyLock::new(||
     keywords.insert("COMMIT".to_string(), TokenType::Commit);
     keywords.insert("BEGIN".to_string(), TokenType::Begin);
     keywords.insert("DESCRIBE".to_string(), TokenType::Describe);
+    keywords.insert("PREPARE".to_string(), TokenType::Prepare);
     keywords.insert("PRESERVE".to_string(), TokenType::Preserve);
     keywords.insert("TRANSACTION".to_string(), TokenType::Transaction);
     keywords.insert("SAVEPOINT".to_string(), TokenType::Savepoint);

{polyglot_sql-0.4.2 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/traversal.rs

@@ -701,6 +701,18 @@ fn iter_children(expr: &Expression) -> Vec<(&'static str, &Expression)> {
         Expression::CreateTask(ct) => {
             children.push(("body", &ct.body));
         }
+        Expression::Prepare(prepare) => {
+            children.push(("statement", &prepare.statement));
+        }
+        Expression::Execute(exec) => {
+            children.push(("this", &exec.this));
+            for argument in &exec.arguments {
+                children.push(("argument", argument));
+            }
+            for parameter in &exec.parameters {
+                children.push(("parameter", &parameter.value));
+            }
+        }
         Expression::Analyze(a) => {
             if let Some(this) = &a.this {
                 children.push(("this", this));