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

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

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

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

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

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

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

{polyglot_sql-0.4.1 → polyglot_sql-0.4.2}/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.2", optional = true, default-features = false }
 
 [dev-dependencies]
 pretty_assertions = "1.4"

{polyglot_sql-0.4.1 → polyglot_sql-0.4.2}/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.2}/crates/polyglot-sql/src/lineage.rs

@@ -714,7 +714,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 {
@@ -1265,13 +1265,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 +1655,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);
                 }
@@ -2989,6 +3010,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.2}/crates/polyglot-sql/src/openlineage.rs

@@ -1126,6 +1126,45 @@ 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_aggregation_column_lineage() {
         let result =

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

@@ -619,6 +619,39 @@ impl Parser {
         parser.parse()
     }
 
+    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 parse_method_call_expression(
+        &self,
+        this: Expression,
+        method: Identifier,
+        args: Vec<Expression>,
+    ) -> Expression {
+        if matches!(
+            self.config.dialect,
+            Some(crate::dialects::DialectType::BigQuery)
+        ) && Self::is_bigquery_safe_namespace_receiver(&this)
+            && !method.quoted
+        {
+            return Expression::Function(Box::new(Function::new(
+                format!("SAFE.{}", method.name),
+                args,
+            )));
+        }
+
+        Expression::MethodCall(Box::new(MethodCall { this, method, args }))
+    }
+
     /// Parse all remaining statements from the token stream.
     ///
     /// Consumes tokens until the end of input, splitting on semicolons.
@@ -30471,18 +30504,15 @@ impl Parser {
                     self.parse_function_arguments()?
                 };
                 self.expect(TokenType::RParen)?;
-                let method_call = Expression::MethodCall(Box::new(MethodCall {
-                    this: Expression::boxed_column(Column {
-                        name: ident.clone(),
-                        table: None,
-                        join_mark: false,
-                        trailing_comments: Vec::new(),
-                        span: None,
-                        inferred_type: None,
-                    }),
-                    method: col_ident,
-                    args,
-                }));
+                let this = Expression::boxed_column(Column {
+                    name: ident.clone(),
+                    table: None,
+                    join_mark: false,
+                    trailing_comments: Vec::new(),
+                    span: None,
+                    inferred_type: None,
+                });
+                let method_call = self.parse_method_call_expression(this, col_ident, args);
                 return self.maybe_parse_subscript(method_call);
             }
 
@@ -32321,11 +32351,11 @@ impl Parser {
                         self.parse_expression_list()?
                     };
                     self.expect(TokenType::RParen)?;
-                    let method_call = Expression::MethodCall(Box::new(MethodCall {
-                        this: Expression::Identifier(Identifier::new(name)),
-                        method: col_ident,
+                    let method_call = self.parse_method_call_expression(
+                        Expression::Identifier(Identifier::new(name)),
+                        col_ident,
                         args,
-                    }));
+                    );
                     return self.maybe_parse_subscript(method_call);
                 }
 
@@ -37831,11 +37861,11 @@ impl Parser {
                         };
                         self.expect(TokenType::RParen)?;
                         // Create a method call expression (DotAccess with function call)
-                        expr = Expression::MethodCall(Box::new(MethodCall {
-                            this: expr,
-                            method: Identifier::new(field_name),
+                        expr = self.parse_method_call_expression(
+                            expr,
+                            Identifier::new(field_name),
                             args,
-                        }));
+                        );
                     } else {
                         let mut ident = Identifier::new(field_name);
                         if is_quoted {