Repository: tobymao/sqlglot (2024-09-10 ~ 2024-09-16)

problem_statement: [Snowflake][Parse] Dynamic `PIVOT` query parsed/optimized incorrectly
**Fully reproducible code snippet**

```
example_pivot_q = """
SELECT *
  FROM quarterly_sales
    PIVOT(SUM(amount) FOR quarter IN (ANY ORDER BY quarter))
  ORDER BY empid;
"""
dialect = Snowflake()
pivot_exp = parse_one(example_pivot_q, dialect=dialect)
optimized_exp = optimize(pivot_exp, dialect=dialect)

print(pivot_exp.sql(dialect=dialect))
# returns
# 'SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR IN (ANY ORDER BY quarter)) ORDER BY empid'
print(optimized_exp.sql(dialect=dialect))
# 'SELECT * FROM "QUARTERLY_SALES" AS "QUARTERLY_SALES" PIVOT(SUM("QUARTERLY_SALES"."AMOUNT") FOR IN (ANY ORDER BY "QUARTERLY_SALES"."QUARTER")) AS "_q_0" ORDER BY "_q_0"."EMPID"'
```

**Official Documentation**
https://docs.snowflake.com/en/sql-reference/constructs/pivot#pivot-on-all-distinct-column-values-automatically-with-dynamic-pivot


date: 2024-09-10
patch: diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 746712b6..7ab3416a 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1917,8 +1917,6 @@ class Generator(metaclass=_Generator):
         direction = self.seg("UNPIVOT" if expression.unpivot else "PIVOT")
 
         field = self.sql(expression, "field")
-        if field and isinstance(expression.args.get("field"), exp.PivotAny):
-            field = f"IN ({field})"
 
         include_nulls = expression.args.get("include_nulls")
         if include_nulls is not None:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 9ea81837..f22de203 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -3907,13 +3907,12 @@ class Parser(metaclass=_Parser):
             self.raise_error("Expecting IN (")
 
         if self._match(TokenType.ANY):
-            expr: exp.PivotAny | exp.In = self.expression(exp.PivotAny, this=self._parse_order())
+            exprs: t.List[exp.Expression] = ensure_list(exp.PivotAny(this=self._parse_order()))
         else:
-            aliased_expressions = self._parse_csv(_parse_aliased_expression)
-            expr = self.expression(exp.In, this=value, expressions=aliased_expressions)
+            exprs = self._parse_csv(_parse_aliased_expression)
 
         self._match_r_paren()
-        return expr
+        return self.expression(exp.In, this=value, expressions=exprs)
 
     def _parse_pivot(self) -> t.Optional[exp.Pivot]:
         index = self._index

test_patch: diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index b82f9928..9c036c76 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -107,10 +107,10 @@ WHERE
             "SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR quarter IN (SELECT DISTINCT quarter FROM ad_campaign_types_by_quarter WHERE television = TRUE ORDER BY quarter)) ORDER BY empid"
         )
         self.validate_identity(
-            "SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR IN (ANY ORDER BY quarter)) ORDER BY empid"
+            "SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR quarter IN (ANY ORDER BY quarter)) ORDER BY empid"
         )
         self.validate_identity(
-            "SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR IN (ANY)) ORDER BY empid"
+            "SELECT * FROM quarterly_sales PIVOT(SUM(amount) FOR quarter IN (ANY)) ORDER BY empid"
         )
         self.validate_identity(
             "MERGE INTO my_db AS ids USING (SELECT new_id FROM my_model WHERE NOT col IS NULL) AS new_ids ON ids.type = new_ids.type AND ids.source = new_ids.source WHEN NOT MATCHED THEN INSERT VALUES (new_ids.new_id)"


----

problem_statement: Using singe quote as part of JSON object key leads to invalid SQL generation
**Before you file an issue**
- Make sure you specify the "read" dialect eg. `parse_one(sql, read="spark")` - ok
- Make sure you specify the "write" dialect eg. `ast.sql(dialect="duckdb")` - ok
- Check if the issue still exists on main - checked on sqlglot-25.20.2.dev5

**Overview**

When querying inside JSON object in Snowflake and referenced object's key contains `'` character the statements seems to be parsed correctly but when generating SQL the output is not valid (at least for snowflake and postgres dialects). Replacing single `'` in key name with escaped one: `''` resolves at least this case.

Object key can be referenced using `:` notation (e.g. `json_doc:"Customer's department"`) or using `get_path` function (`GET_PATH(PARSE_JSON(doc), '"Customer''s department"')`) - the result is the same.

**Fully reproducible code snippet**

```python
import sqlglot
sql_text = """select parse_json(doc):"Customer's department"::string as "dept" from json_tbl"""
print(sqlglot.parse_one(sql_text, dialect="snowflake").sql(dialect="snowflake"))
# SELECT CAST(GET_PATH(PARSE_JSON(doc), '["Customer's department"]') AS TEXT) AS "dept" FROM json_tbl

print(sqlglot.parse_one(sql_text, dialect="snowflake").sql(dialect="postgres"))
# SELECT CAST(JSON_EXTRACT_PATH(CAST(doc AS JSON), 'Customer's department') AS TEXT) AS "dept" FROM json_tbl
```
Expecting output for Snowflake to be 
```sql
SELECT CAST(GET_PATH(PARSE_JSON(doc), '["Customer''s department"]') AS TEXT) AS "dept" FROM json_tbl;
-- or
SELECT CAST(GET_PATH(PARSE_JSON(doc), '"Customer''s department"') AS TEXT) AS "dept" FROM json_tbl;
-- or, if not using GET_PATH function the original name without escaped quote can be used
SELECT CAST(PARSE_JSON(doc):"Customer's department" AS TEXT) AS "dept" FROM json_tbl;
```
And for PostgreSQL
```sql
SELECT CAST(JSON_EXTRACT_PATH(CAST(doc AS JSON), 'Customer''s department') AS TEXT) AS "dept" FROM json_tbl;
```

To test query with postgres:
```sql
create table json_tbl(doc jsonb);
insert into json_tbl values('{"Customer''s department": "Department A"}');

-- invalid query:
SELECT CAST(JSON_EXTRACT_PATH(CAST(doc AS JSON), 'Customer's department') AS TEXT) AS "dept" FROM json_tbl;
/*
ERROR:  syntax error at or near "s"
LINE 1: ...ST(JSON_EXTRACT_PATH(CAST(doc AS JSON), 'Customer's departme...
                                                             ^ 

SQL state: 42601
Character: 60
*/
-- valid query:
SELECT CAST(JSON_EXTRACT_PATH(CAST(doc AS JSON), 'Customer''s department') AS TEXT) AS "dept" FROM json_tbl;
```


**Official Documentation**
Please include links to official SQL documentation related to your issue.
- [Snowflake: Double quoted identifiers](https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers)
- [Snowflake: Single-quoted string constants](https://docs.snowflake.com/en/sql-reference/data-types-text#label-single-quoted-string-constants)

date: 2024-09-11
patch: diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py
index 00f63532..0715eb7e 100644
--- a/sqlglot/dialects/databricks.py
+++ b/sqlglot/dialects/databricks.py
@@ -115,3 +115,7 @@ class Databricks(Spark):
         ) -> str:
             expression.set("this", True)  # trigger ALWAYS in super class
             return super().generatedasidentitycolumnconstraint_sql(expression)
+
+        def jsonpath_sql(self, expression: exp.JSONPath) -> str:
+            expression.set("escape", None)
+            return super().jsonpath_sql(expression)
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 78a93593..21e52c91 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -1595,6 +1595,8 @@ def json_extract_segments(
         if not isinstance(path, exp.JSONPath):
             return rename_func(name)(self, expression)
 
+        escape = path.args.get("escape")
+
         segments = []
         for segment in path.expressions:
             path = self.sql(segment)
@@ -1602,6 +1604,9 @@ def json_extract_segments(
                 if isinstance(segment, exp.JSONPathPart) and (
                     quoted_index or not isinstance(segment, exp.JSONPathSubscript)
                 ):
+                    if escape:
+                        path = self.escape_str(path)
+
                     path = f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
 
                 segments.append(path)
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index a7bc2ca4..5355aed1 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -5523,7 +5523,7 @@ class JSON(Expression):
 
 
 class JSONPath(Expression):
-    arg_types = {"expressions": True}
+    arg_types = {"expressions": True, "escape": False}
 
     @property
     def output_name(self) -> str:
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 7ab3416a..26c20c80 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -2737,6 +2737,10 @@ class Generator(metaclass=_Generator):
 
     def jsonpath_sql(self, expression: exp.JSONPath) -> str:
         path = self.expressions(expression, sep="", flat=True).lstrip(".")
+
+        if expression.args.get("escape"):
+            path = self.escape_str(path)
+
         if self.QUOTE_JSON_PATH:
             path = f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index f22de203..5a9b9180 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -4960,6 +4960,7 @@ class Parser(metaclass=_Parser):
     ) -> t.Optional[exp.Expression]:
         casts = []
         json_path = []
+        escape = None
 
         while self._match(TokenType.COLON):
             start_index = self._index
@@ -4986,25 +4987,25 @@ class Parser(metaclass=_Parser):
                 end_token = self._prev
 
             if path:
-                path_sql = self._find_sql(self._tokens[start_index], end_token)
-
                 # Escape single quotes from Snowflake's colon extraction (e.g. col:"a'b") as
                 # it'll roundtrip to a string literal in GET_PATH
-                if path_sql[0] == '"':
-                    quote_end = self.dialect.QUOTE_END
-                    path_sql = path_sql.replace(
-                        quote_end, self.dialect.tokenizer_class.STRING_ESCAPES[0] + quote_end
-                    )
+                if isinstance(path, exp.Identifier) and path.quoted:
+                    escape = True
 
-                json_path.append(path_sql)
+                json_path.append(self._find_sql(self._tokens[start_index], end_token))
 
         # The VARIANT extract in Snowflake/Databricks is parsed as a JSONExtract; Snowflake uses the json_path in GET_PATH() while
         # Databricks transforms it back to the colon/dot notation
         if json_path:
+            json_path_expr = self.dialect.to_json_path(exp.Literal.string(".".join(json_path)))
+
+            if json_path_expr:
+                json_path_expr.set("escape", escape)
+
             this = self.expression(
                 exp.JSONExtract,
                 this=this,
-                expression=self.dialect.to_json_path(exp.Literal.string(".".join(json_path))),
+                expression=json_path_expr,
                 variant_extract=True,
             )
 

test_patch: diff --git a/tests/dialects/test_databricks.py b/tests/dialects/test_databricks.py
index a46794bb..c2965e83 100644
--- a/tests/dialects/test_databricks.py
+++ b/tests/dialects/test_databricks.py
@@ -142,6 +142,13 @@ class TestDatabricks(Validator):
             """SELECT raw:`zip code`, raw:`fb:testid`, raw:store['bicycle'], raw:store["zip code"]""",
             """SELECT raw:["zip code"], raw:["fb:testid"], raw:store.bicycle, raw:store["zip code"]""",
         )
+        self.validate_all(
+            "SELECT col:`fr'uit`",
+            write={
+                "databricks": """SELECT col:["fr'uit"]""",
+                "postgres": "SELECT JSON_EXTRACT_PATH(col, 'fr''uit')",
+            },
+        )
 
     def test_datediff(self):
         self.validate_all(
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 9c036c76..40413c82 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -2143,7 +2143,7 @@ SINGLE = TRUE""",
             """,
             write={
                 "snowflake": """SELECT GET_PATH(col, '["customer\\'s department"]')""",
-                "postgres": "SELECT JSON_EXTRACT_PATH(col, 'customer\\'s department')",
+                "postgres": "SELECT JSON_EXTRACT_PATH(col, 'customer''s department')",
             },
         )
 


----

problem_statement: Optimizer Misinterprets Union Identifiers as Columns in DuckDB Queries
When testing the creation of a union with DuckDB, it appears that sqlglot mistakenly treats the union identifier as a column rather than recognizing it as a general identifier when the statement is qualified in the optimizer.

DuckDB uses identifiers to tag unions, as described in [official documentation](https://duckdb.org/docs/sql/data_types/union)

Here is how to reproduce the exception:

```python
import sqlglot
import sqlglot.optimizer

statement = sqlglot.parse_one("select union_value(k1 := 1)", dialect="duckdb")
sqlglot.optimizer.optimize(statement).sql(dialect='duckdb')
```

In this case, `k1` is simply an identifier—just like a key in a struct. It would be helpful if it could be handled in a case-insensitive manner, similar to other struct keys.

The exception raised is as follows:

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/rusty/Development/flight-cloud/.venv/lib/python3.12/site-packages/sqlglot/optimizer/optimizer.py", line 92, in optimize
    optimized = rule(optimized, **rule_kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/rusty/Development/flight-cloud/.venv/lib/python3.12/site-packages/sqlglot/optimizer/qualify.py", line 99, in qualify
    validate_qualify_columns_func(expression)
  File "/Users/rusty/Development/flight-cloud/.venv/lib/python3.12/site-packages/sqlglot/optimizer/qualify_columns.py", line 106, in validate_qualify_columns
    raise OptimizeError(f"Column '{column}' could not be resolved{for_table}")
sqlglot.errors.OptimizeError: Column 'k1' could not be resolved
```

In contrast, DuckDB handles this case correctly, as shown below:

```sql
v1.1.0 fa5c2fe15f
Enter ".help" for usage hints.
Connected to a transient in-memory database.
Use ".open FILENAME" to reopen on a persistent database.
D select union_value(k1 := 1);
┌──────────────────────┐
│ union_value(k1 := 1) │
│  union(k1 integer)   │
├──────────────────────┤
│ 1                    │
└──────────────────────┘
```

Thank you for looking into this.

This behavior in DuckDB demonstrates that identifiers like `k1` should be recognized as keys in a struct rather than treated as columns. Handling this more gracefully in `sqlglot` would align with DuckDB’s behavior.
date: 2024-09-16
patch: diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index a47ca8eb..a2180357 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -4343,6 +4343,9 @@ class Parser(metaclass=_Parser):
             )
 
         while self._match_set(self.ASSIGNMENT):
+            if isinstance(this, exp.Column) and len(this.parts) == 1:
+                this = this.this
+
             this = self.expression(
                 self.ASSIGNMENT[self._prev.token_type],
                 this=this,

test_patch: diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
index 2887dd0f..ad503ae9 100644
--- a/tests/dialects/test_duckdb.py
+++ b/tests/dialects/test_duckdb.py
@@ -295,6 +295,9 @@ class TestDuckDB(Validator):
         self.validate_identity("SUMMARIZE tbl").assert_is(exp.Summarize)
         self.validate_identity("SUMMARIZE SELECT * FROM tbl").assert_is(exp.Summarize)
         self.validate_identity("CREATE TABLE tbl_summary AS SELECT * FROM (SUMMARIZE tbl)")
+        self.validate_identity("UNION_VALUE(k1 := 1)").find(exp.PropertyEQ).this.assert_is(
+            exp.Identifier
+        )
         self.validate_identity(
             "SELECT species, island, COUNT(*) FROM t GROUP BY GROUPING SETS (species), GROUPING SETS (island)"
         )