Minor upgrades to Snowflake parameter processing

core/src/main/antlr4/com/databricks/labs/remorph/parsers/snowflake/SnowflakeParser.g4

@@ -2977,7 +2977,7 @@ id
     | ID2
     | DOUBLE_QUOTE_ID
     | DOUBLE_QUOTE_BLANK
-    | AMP ID           // Snowflake variables from CLI or injection
+    | AMP LCB? ID RCB? // Snowflake variables from CLI or injection - we rely on valid input
     | nonReservedWords //id is used for object name. Snowflake is very permissive
     ;
 

core/src/main/scala/com/databricks/labs/remorph/generators/sql/ExpressionGenerator.scala

@@ -457,7 +457,7 @@ class ExpressionGenerator(val callMapper: ir.CallMapper = new ir.CallMapper())
     arg.name_parts.mkString(".")
   }
 
-  private def variable(ctx: GeneratorContext, v: ir.Variable): String = s"$$${v.name}"
+  private def variable(ctx: GeneratorContext, v: ir.Variable): String = s"$${${v.name}}"
 
   private def concat(ctx: GeneratorContext, c: ir.Concat): String = {
     val args = c.children.map(expression(ctx, _))

core/src/main/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeExpressionBuilder.scala

@@ -36,7 +36,9 @@ class SnowflakeExpressionBuilder()
       // To rework that is quite a big job. So, for now, we translate &id to $id here. It is not wrong for the id rule
       // to hold the AMP ID alt, but ideally it would produce an ir.Variable and we would process that at generation
       // time instead of concatenating into strings :(
-      ir.Id("$" + v.ID().getText)
+      ir.Id(s"$$${v.ID().getText}")
+    case d if d.ID2() != null =>
+      ir.Id(s"$$${d.ID2().getText.drop(1)}")
     case id => ir.Id(id.getText)
   }
 
@@ -168,7 +170,7 @@ class SnowflakeExpressionBuilder()
             .stringPart()
             .asScala
             .map {
-              case p if p.VAR_SIMPLE() != null => s"$${${p.VAR_SIMPLE().getText.drop(1)}}" // &var => ${var}
+              case p if p.VAR_SIMPLE() != null => s"$${${p.VAR_SIMPLE().getText.drop(1)}}" // &var => ${var} (soon)
               case p if p.VAR_COMPLEX() != null => s"$$${p.VAR_COMPLEX().getText.drop(1)}" // &{var} => ${var}
               case p if p.STRING_AMPAMP() != null => "&" // && => &
               case p if p.STRING_CONTENT() != null => p.STRING_CONTENT().getText

core/src/test/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeAstBuilderSpec.scala

@@ -488,6 +488,7 @@ class SnowflakeAstBuilderSpec extends AnyWordSpec with SnowflakeParserTestCommon
       "select * from a where b = &ids" in {
         singleQueryExample(
           "select * from a where b = &ids",
+          // Note when we truly process &vars we should get Variable, not Id
           Project(Filter(namedTable("a"), Equals(Id("b"), Id("$ids"))), Seq(Star())))
       }
     }

core/src/test/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeExpressionBuilderSpec.scala

@@ -509,12 +509,18 @@ class SnowflakeExpressionBuilderSpec
     }
   }
 
+  // Note that when we truly handle &vars, we will get ir.Variable here and not Id
+  // and the & parts will not be changed to ${} until we get to the final SQL generation
+  // but we are in a half way house transition state
   "variable substitution" should {
     "&abc" in {
       exampleExpr("&abc", _.expr(), Id("$abc"))
     }
     "&ab_c.bc_d" in {
       exampleExpr("&ab_c.bc_d", _.expr(), Dot(Id("$ab_c"), Id("bc_d")))
     }
+    "&{ab_c}.&bc_d" in {
+      exampleExpr("&{ab_c}.&bc_d", _.expr(), Dot(Id("$ab_c"), Id("$bc_d")))
+    }
   }
 }

tests/resources/functional/snowflake/test_tokens_parameter/test_tokens_parameter_1.sql

@@ -2,4 +2,4 @@
 select emp_id from abc.emp where emp_id = &ids;
 
 -- databricks sql:
-select emp_id from abc.emp where emp_id = $ids;
+select emp_id from abc.emp where emp_id = $ids;

tests/resources/functional/snowflake/test_tokens_parameter/test_tokens_parameter_2.sql

@@ -2,4 +2,4 @@
 select emp_id from abc.emp where emp_id = $ids;
 
 -- databricks sql:
-select emp_id from abc.emp where emp_id = $ids;
+select emp_id from abc.emp where emp_id = $ids;

tests/resources/functional/snowflake/test_tokens_parameter/test_tokens_parameter_4.sql

@@ -0,0 +1,11 @@
+-- Note that we cannot enable the WHERE clause here as teh python
+-- tests will fail because SqlGlot cannot translate variable references
+-- within strings. Enable the WHERE clause when we are not limited by Python
+
+-- WHERE EMP_ID = '&empNo' =>  WHERE EMP_ID = '${empNo}'
+
+-- snowflake sql:
+select count(*) from &TEST_USER.EMP_TBL;
+
+-- databricks sql:
+select count(*) from $TEST_USER.EMP_TBL;