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

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

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

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

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

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

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

{polyglot_sql-0.4.1 → 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.1", 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.1 → polyglot_sql-0.4.3}/crates/polyglot-sql/src/dialects/bigquery.rs

@@ -1587,6 +1587,38 @@ mod tests {
         assert_eq!(result, expected, "SQL: {}", sql);
     }
 
+    #[test]
+    fn test_safe_namespace_parses_as_function() {
+        let expr = parse_one(
+            "SELECT SAFE.PARSE_JSON(data) AS json_data FROM t",
+            DialectType::BigQuery,
+        )
+        .expect("parse");
+
+        let Expression::Select(select) = expr else {
+            panic!("expected SELECT");
+        };
+        let Expression::Alias(alias) = &select.expressions[0] else {
+            panic!("expected alias");
+        };
+        let Expression::Function(function) = &alias.this else {
+            panic!("expected SAFE namespace call to parse as Function");
+        };
+
+        assert_eq!(function.name, "SAFE.PARSE_JSON");
+        assert_eq!(function.args.len(), 1);
+    }
+
+    #[test]
+    fn test_safe_namespace_identity() {
+        bigquery_identity("SAFE.PARSE_JSON(data)", "SAFE.PARSE_JSON(data)");
+        bigquery_identity(
+            "SAFE.PARSE_DATE('%Y-%m-%d', date_col)",
+            "SAFE.PARSE_DATE('%F', date_col)",
+        );
+        bigquery_identity("SAFE.DIVIDE(a, b)", "SAFE.DIVIDE(a, b)");
+    }
+
     #[test]
     fn test_cast_char_to_string() {
         bigquery_identity("CAST(x AS CHAR)", "CAST(x AS STRING)");

{polyglot_sql-0.4.1 → 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.1 → 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.1 → 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.1 → 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,
         }
@@ -714,7 +740,7 @@ fn to_node_inner(
     }
 
     // 7. Column references — trace each column to its source
-    let col_refs = find_column_refs_in_expr(&select_expr);
+    let col_refs = find_column_refs_in_expr(&select_expr, dialect);
     for col_ref in col_refs {
         let col_name = &col_ref.column;
         if let Some(ref table_id) = col_ref.table {
@@ -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,
         }
@@ -1265,13 +1295,30 @@ struct SimpleColumnRef {
 }
 
 /// Find all column references in an expression (does not recurse into subqueries).
-fn find_column_refs_in_expr(expr: &Expression) -> Vec<SimpleColumnRef> {
+fn find_column_refs_in_expr(
+    expr: &Expression,
+    dialect: Option<DialectType>,
+) -> Vec<SimpleColumnRef> {
     let mut refs = Vec::new();
-    collect_column_refs(expr, &mut refs);
+    collect_column_refs(expr, dialect, &mut refs);
     refs
 }
 
-fn collect_column_refs(expr: &Expression, refs: &mut Vec<SimpleColumnRef>) {
+fn is_bigquery_safe_namespace_receiver(expr: &Expression) -> bool {
+    match expr {
+        Expression::Column(col) => {
+            col.table.is_none() && !col.name.quoted && col.name.name.eq_ignore_ascii_case("SAFE")
+        }
+        Expression::Identifier(id) => !id.quoted && id.name.eq_ignore_ascii_case("SAFE"),
+        _ => false,
+    }
+}
+
+fn collect_column_refs(
+    expr: &Expression,
+    dialect: Option<DialectType>,
+    refs: &mut Vec<SimpleColumnRef>,
+) {
     let mut stack: Vec<&Expression> = vec![expr];
 
     while let Some(current) = stack.pop() {
@@ -1638,7 +1685,11 @@ fn collect_column_refs(expr: &Expression, refs: &mut Vec<SimpleColumnRef>) {
                 stack.push(&d.this);
             }
             Expression::MethodCall(m) => {
-                stack.push(&m.this);
+                if !matches!(dialect, Some(DialectType::BigQuery))
+                    || !is_bigquery_safe_namespace_receiver(&m.this)
+                {
+                    stack.push(&m.this);
+                }
                 for arg in &m.args {
                     stack.push(arg);
                 }
@@ -2465,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(
@@ -2989,6 +3079,71 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_lineage_bigquery_safe_namespace_issue207() {
+        let query = r#"
+WITH import_cte AS (
+  SELECT timestamp, data, operation
+  FROM `project`.`dataset`.`source_table`
+),
+transform_cte AS (
+  SELECT
+    timestamp,
+    SAFE.PARSE_JSON(data) AS json_data
+  FROM import_cte
+)
+SELECT json_data FROM transform_cte
+"#;
+        let expr = parse_one(query, DialectType::BigQuery).expect("parse");
+        let node = lineage("json_data", &expr, Some(DialectType::BigQuery), false)
+            .expect("lineage should resolve SAFE.PARSE_JSON arguments");
+        let names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
+
+        assert!(
+            names.iter().any(|name| name == "source_table.data"),
+            "expected source_table.data in lineage, got {names:?}"
+        );
+        assert!(
+            !names
+                .iter()
+                .any(|name| name.eq_ignore_ascii_case("import_cte.safe")),
+            "did not expect SAFE namespace receiver in lineage, got {names:?}"
+        );
+    }
+
+    #[test]
+    fn test_lineage_bigquery_safe_namespace_method_call_guard() {
+        let expr = parse("SELECT SAFE.PARSE_JSON(data) AS json_data FROM t");
+        let node = lineage("json_data", &expr, Some(DialectType::BigQuery), false)
+            .expect("lineage should resolve SAFE.PARSE_JSON arguments");
+        let names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
+
+        assert!(
+            names.iter().any(|name| name == "t.data"),
+            "expected t.data in lineage, got {names:?}"
+        );
+        assert!(
+            !names.iter().any(|name| name.eq_ignore_ascii_case("t.safe")),
+            "did not expect SAFE namespace receiver in lineage, got {names:?}"
+        );
+    }
+
+    #[test]
+    fn test_lineage_method_call_receiver_control() {
+        let expr = parse("SELECT obj.METHOD(arg) AS out FROM t");
+        let node = lineage("out", &expr, None, false).expect("lineage");
+        let names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
+
+        assert!(
+            names.iter().any(|name| name == "t.obj"),
+            "expected ordinary method receiver to remain in lineage, got {names:?}"
+        );
+        assert!(
+            names.iter().any(|name| name == "t.arg"),
+            "expected method argument in lineage, got {names:?}"
+        );
+    }
+
     #[test]
     fn test_lineage_upper_function() {
         let expr = parse("SELECT UPPER(name) AS upper_name FROM users");

{polyglot_sql-0.4.1 → 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())
@@ -1126,6 +1141,78 @@ mod tests {
             .all(|f| f.transformations[0].subtype == "TRANSFORMATION"));
     }
 
+    #[test]
+    fn omits_bigquery_safe_namespace_from_column_lineage_issue207() {
+        let mut opts = options();
+        opts.dialect = DialectType::BigQuery;
+
+        let result = openlineage_column_lineage(
+            r#"
+WITH import_cte AS (
+  SELECT timestamp, data, operation
+  FROM `project`.`dataset`.`source_table`
+),
+transform_cte AS (
+  SELECT
+    timestamp,
+    SAFE.PARSE_JSON(data) AS json_data
+  FROM import_cte
+)
+SELECT json_data FROM transform_cte
+"#,
+            &opts,
+        )
+        .expect("lineage");
+        let field = result.facet.fields.get("json_data").expect("json_data");
+
+        assert!(
+            field.input_fields.iter().any(|input| input.field == "data"),
+            "expected data input field, got {:?}",
+            field.input_fields
+        );
+        assert!(
+            !field
+                .input_fields
+                .iter()
+                .any(|input| input.field.eq_ignore_ascii_case("safe")),
+            "did not expect SAFE namespace as input field, got {:?}",
+            field.input_fields
+        );
+    }
+
+    #[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 =