Repository: tobymao/sqlglot

Issue:
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)


Discussion:

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')",
             },
         )