Fix incorrectly captured calls to col

src/snowflake/snowpark/_internal/ast/utils.py

@@ -888,7 +888,7 @@ def build_expr_from_snowpark_column_or_table_fn(
 
 # TODO(SNOW-1491199) - This method is not covered by tests until the end of phase 0. Drop the pragma when it is covered.
 def fill_ast_for_column(  # type: ignore[no-untyped-def] # TODO(SNOW-1491199) # Function is missing a type annotation for one or more arguments
-    expr: proto.Expr, name1: str, name2: Optional[str], fn_name="col"
+    expr: proto.Expr, name1: str, name2: Optional[str], fn_name=None
 ) -> None:  # pragma: no cover
     """
     Fill in expr node to encode Snowpark Column created through col(...) / column(...).
@@ -899,19 +899,24 @@ def fill_ast_for_column(  # type: ignore[no-untyped-def] # TODO(SNOW-1491199) #
         fn_name: alias to use when encoding Snowpark column (should be "col" or "column").
 
     """
-
-    # Handle the special case * (as a SQL column expr).
-    if name2 == "*":
-        ast = with_src_position(expr.sp_column_sql_expr)  # type: ignore[arg-type] # TODO(SNOW-1491199) # Argument 1 to "with_src_position" has incompatible type "SpColumnSqlExpr"; expected "Expr"
-        ast.sql = "*"  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "sql"
-        if name1 is not None:
-            ast.df_alias.value = name1  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "df_alias"
-        return expr  # type: ignore[return-value] # TODO(SNOW-1491199) # No return value expected
-
-    if name1 == "*" and name2 is None:
-        ast = with_src_position(expr.sp_column_sql_expr)  # type: ignore[arg-type] # TODO(SNOW-1491199) # Argument 1 to "with_src_position" has incompatible type "SpColumnSqlExpr"; expected "Expr"
-        ast.sql = "*"  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "sql"
-        return expr  # type: ignore[return-value] # TODO(SNOW-1491199) # No return value expected
+    if fn_name is None:
+        # Handle the special case * (as a SQL column expr).
+        if name2 == "*":
+            ast = with_src_position(expr.sp_column_sql_expr)  # type: ignore[arg-type] # TODO(SNOW-1491199) # Argument 1 to "with_src_position" has incompatible type "SpColumnSqlExpr"; expected "Expr"
+            ast.sql = "*"  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "sql"
+            if name1 is not None:
+                ast.df_alias.value = name1  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "df_alias"
+            return expr  # type: ignore[return-value] # TODO(SNOW-1491199) # No return value expected
+
+        if name1 == "*" and name2 is None:
+            ast = with_src_position(expr.sp_column_sql_expr)  # type: ignore[arg-type] # TODO(SNOW-1491199) # Argument 1 to "with_src_position" has incompatible type "SpColumnSqlExpr"; expected "Expr"
+            ast.sql = "*"  # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "sql"
+            return expr  # type: ignore[return-value] # TODO(SNOW-1491199) # No return value expected
+
+        if name2 is None:
+            ast = with_src_position(expr.string_val)
+            ast.v = name1
+            return expr
 
     # Regular form (without *): build as function ApplyExpr.
     kwargs = (
@@ -930,7 +935,7 @@ def fill_ast_for_column(  # type: ignore[no-untyped-def] # TODO(SNOW-1491199) #
 
 # TODO(SNOW-1491199) - This method is not covered by tests until the end of phase 0. Drop the pragma when it is covered.
 def create_ast_for_column(  # type: ignore[no-untyped-def] # TODO(SNOW-1491199) # Function is missing a type annotation for one or more arguments
-    name1: str, name2: Optional[str], fn_name="col"
+    name1: str, name2: Optional[str], fn_name=None
 ) -> proto.Expr:  # pragma: no cover
     """
     Helper function to create Ast for Snowpark Column. Cf. fill_ast_for_column on parameter details.