Fix Jinja source parsing boundaries

Author: dataders

src/parser/source_tree.rs

@@ -301,6 +301,18 @@ mod tests {
         assert!(viz.starts_with("VISUALISE FROM mtcars"));
     }
 
+    #[test]
+    fn test_extract_sql_visualise_from_jinja_ref() {
+        let query = "VISUALISE FROM {{ ref('fct_orders') }} DRAW point MAPPING x AS x, y AS y";
+        let tree = SourceTree::new(query).unwrap();
+
+        let sql = tree.extract_sql().unwrap();
+        assert_eq!(sql, "SELECT * FROM {{ ref('fct_orders') }}");
+
+        let viz = tree.extract_visualise().unwrap();
+        assert!(viz.starts_with("VISUALISE FROM {{ ref('fct_orders') }}"));
+    }
+
     #[test]
     fn test_extract_sql_visualise_from_with_cte() {
         let query =
@@ -419,6 +431,15 @@ mod tests {
         assert!(sql.contains("SELECT * FROM 'mtcars.csv'"));
     }
 
+    #[test]
+    fn test_extract_sql_from_first_jinja_ref() {
+        let query = "FROM {{ ref('fct_orders') }} VISUALISE DRAW point MAPPING x AS x, y AS y";
+        let tree = SourceTree::new(query).unwrap();
+
+        let sql = tree.extract_sql().unwrap();
+        assert_eq!(sql, "SELECT * FROM {{ ref('fct_orders') }}");
+    }
+
     #[test]
     fn test_extract_sql_from_first_case_insensitive() {
         let query = "from sales visualise DRAW point MAPPING x AS x, y AS y";

tree-sitter-ggsql/grammar.js

@@ -17,6 +17,10 @@ function caseInsensitive(keyword) {
 module.exports = grammar({
   name: 'ggsql',
 
+  inline: $ => [
+    $.source_ref,
+  ],
+
   conflicts: $ => [
     [$.sql_portion],
   ],
@@ -480,9 +484,16 @@ module.exports = grammar({
       repeat(seq('.', $.identifier))
     )),
 
+    source_ref: $ => choice(
+      $.qualified_name,
+      $.string,
+      $.namespaced_identifier,
+      $.jinja_template
+    ),
+
     table_ref: $ => prec.right(seq(
       choice(
-        field('table', choice($.qualified_name, $.string, $.namespaced_identifier, $.jinja_template)),
+        field('table', $.source_ref),
         $.subquery,
       ),
       optional(seq(
@@ -601,14 +612,14 @@ module.exports = grammar({
         // Option 1: Just FROM (inherit global mappings)
         seq(
           caseInsensitive('FROM'),
-          field('layer_source', choice($.qualified_name, $.string, $.namespaced_identifier))
+          field('layer_source', $.source_ref)
         ),
         // Option 2: Mapping list (uses shared structure), optionally followed by FROM
         seq(
           $.mapping_list,
           optional(seq(
             caseInsensitive('FROM'),
-            field('layer_source', choice($.qualified_name, $.string, $.namespaced_identifier))
+            field('layer_source', $.source_ref)
           ))
         )
       )
@@ -942,9 +953,9 @@ module.exports = grammar({
     // before ggsql executes SQL, but the parser must preserve them while
     // splitting SQL from VISUALISE.
     jinja_template: $ => token(choice(
-      /\{\{[^}]*\}\}/,
-      /\{%[^%]*%\}/,
-      /\{#[^#]*#\}/
+      seq('{{', repeat(choice(/[^}]+/, /}[^}]/)), '}}'),
+      seq('{%', repeat(choice(/[^%]+/, /%[^%]/)), '%}'),
+      seq('{#', repeat(choice(/[^#]+/, /#[^#]/)), '#}')
     )),
 
     // Identifier for use in filter expressions - uses lower precedence so that

tree-sitter-ggsql/test/corpus/basic.txt

@@ -3649,3 +3649,69 @@ DRAW point
     (viz_clause
       (draw_clause
         (geom_type)))))
+
+================================================================================
+SQL source with Jinja var containing dict literal
+================================================================================
+
+SELECT * FROM {{ var('table', {'fallback': 'orders'}) }}
+VISUALISE x AS x, y AS y
+DRAW point
+
+--------------------------------------------------------------------------------
+
+(query
+  (sql_portion
+    (sql_statement
+      (select_statement
+        (select_body
+          (from_clause
+            (table_ref
+              table: (jinja_template)))))))
+  (visualise_statement
+    (visualise_keyword)
+    (global_mapping
+      (mapping_list
+        (mapping_element
+          (explicit_mapping
+            value: (mapping_value
+              (column_reference
+                (identifier
+                  (bare_identifier))))
+            name: (aesthetic_name)))
+        (mapping_element
+          (explicit_mapping
+            value: (mapping_value
+              (column_reference
+                (identifier
+                  (bare_identifier))))
+            name: (aesthetic_name)))))
+    (viz_clause
+      (draw_clause
+        (geom_type)))))
+
+================================================================================
+Layer source with Jinja ref
+================================================================================
+
+VISUALISE
+DRAW point MAPPING x AS x FROM {{ ref('fct_orders') }}
+
+--------------------------------------------------------------------------------
+
+(query
+  (visualise_statement
+    (visualise_keyword)
+    (viz_clause
+      (draw_clause
+        (geom_type)
+        (mapping_clause
+          (mapping_list
+            (mapping_element
+              (explicit_mapping
+                value: (mapping_value
+                  (column_reference
+                    (identifier
+                      (bare_identifier))))
+                name: (aesthetic_name))))
+          layer_source: (jinja_template))))))