Repository: tobymao/sqlglot

Issue:
Fails to parse subquery followed by a join (snowflake)
Trying to parse the following SQL with `read="snowflake"`

```
sqlglot.parse("UPDATE test SET t = 1 FROM (select id from test2) t2 JOIN test3 t3 on t2.id = t3.id;", read="snowflake")
```
I get the following error:
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mdodds/work/sqlglot/sqlglot/__init__.py", line 103, in parse
    return Dialect.get_or_raise(read or dialect).parse(sql, **opts)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mdodds/work/sqlglot/sqlglot/dialects/dialect.py", line 1115, in parse
    return self.parser(**opts).parse(self.tokenize(sql), sql)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mdodds/work/sqlglot/sqlglot/parser.py", line 1669, in parse
    return self._parse(
           ^^^^^^^^^^^^
  File "/Users/mdodds/work/sqlglot/sqlglot/parser.py", line 1741, in _parse
    self.raise_error("Invalid expression / Unexpected token")
  File "/Users/mdodds/work/sqlglot/sqlglot/parser.py", line 1782, in raise_error
    raise error
sqlglot.errors.ParseError: Invalid expression / Unexpected token. Line 1, Col: 57.
  UPDATE test SET t = 1 FROM (select id from test2) t2 JOIN test3 t3 on t2.id = t3.id;
```

This is a valid query on Snowflake.

If i remove the brackets surrounding the subquery:
```
sqlglot.parse("UPDATE test SET t = 1 FROM select id from test2 t2 JOIN test3 t3 on t2.id = t3.id;", read="snowflake")
```

The query parses correctly:
```
[Update(
  this=Table(
    this=Identifier(this=test, quoted=False)),
  expressions=[
    EQ(
      this=Column(
        this=Identifier(this=t, quoted=False)),
      expression=Literal(this=1, is_string=False))],
  from_=From(
    this=Select(
      expressions=[
        Column(
          this=Identifier(this=id, quoted=False))],
      from_=From(
        this=Table(
          this=Identifier(this=test2, quoted=False),
          alias=TableAlias(
            this=Identifier(this=t2, quoted=False)))),
      joins=[
        Join(
          this=Table(
            this=Identifier(this=test3, quoted=False),
            alias=TableAlias(
              this=Identifier(this=t3, quoted=False))),
          on=EQ(
            this=Column(
              this=Identifier(this=id, quoted=False),
              table=Identifier(this=t2, quoted=False)),
            expression=Column(
              this=Identifier(this=id, quoted=False),
              table=Identifier(this=t3, quoted=False))))])))]
```

Discussion:

Fix (patch):
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 441de7819d..cd808191be 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -3184,7 +3184,12 @@ def _parse_update(self) -> exp.Update:
             elif self._match(TokenType.RETURNING, advance=False):
                 kwargs["returning"] = self._parse_returning()
             elif self._match(TokenType.FROM, advance=False):
-                kwargs["from_"] = self._parse_from(joins=True)
+                from_ = self._parse_from(joins=True)
+                table = from_.this if from_ else None
+                if isinstance(table, exp.Subquery) and self._match(TokenType.JOIN, advance=False):
+                    table.set("joins", list(self._parse_joins()) or None)
+
+                kwargs["from_"] = from_
             elif self._match(TokenType.WHERE, advance=False):
                 kwargs["where"] = self._parse_where()
             elif self._match(TokenType.ORDER_BY, advance=False):


Test changes:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 4f6dfe7bf2..690e2b567e 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -561,10 +561,6 @@ def test_snowflake(self):
         self.validate_identity(
             """SELECT TO_TIMESTAMP('2025-01-16T14:45:30.123+0500', 'yyyy-mm-DD"T"hh24:mi:ss.ff3TZHTZM')"""
         )
-        self.validate_identity(
-            "UPDATE sometesttable u FROM (SELECT 5195 AS new_count, '01bee1e5-0000-d31e-0000-e80ef02b9f27' query_id ) b SET qry_hash_count = new_count WHERE u.sample_query_id  = b.query_id",
-            "UPDATE sometesttable AS u SET qry_hash_count = new_count FROM (SELECT 5195 AS new_count, '01bee1e5-0000-d31e-0000-e80ef02b9f27' AS query_id) AS b WHERE u.sample_query_id = b.query_id",
-        )
         self.validate_identity(
             "SELECT * REPLACE (CAST(col AS TEXT) AS scol) FROM t",
             "SELECT * REPLACE (CAST(col AS VARCHAR) AS scol) FROM t",
@@ -4408,3 +4404,15 @@ def test_ceil(self):
                 "duckdb": "SELECT ROUND(CEIL(1.234 * POWER(10, CAST(1.5 AS INT))) / POWER(10, CAST(1.5 AS INT)), CAST(1.5 AS INT))"
             },
         )
+
+    def test_update_statement(self):
+        self.validate_identity("UPDATE test SET t = 1 FROM t1")
+        self.validate_identity("UPDATE test SET t = 1 FROM t2 JOIN t3 ON t2.id = t3.id")
+        self.validate_identity(
+            "UPDATE test SET t = 1 FROM (SELECT id FROM test2) AS t2 JOIN test3 AS t3 ON t2.id = t3.id"
+        )
+
+        self.validate_identity(
+            "UPDATE sometesttable u FROM (SELECT 5195 AS new_count, '01bee1e5-0000-d31e-0000-e80ef02b9f27' query_id ) b SET qry_hash_count = new_count WHERE u.sample_query_id  = b.query_id",
+            "UPDATE sometesttable AS u SET qry_hash_count = new_count FROM (SELECT 5195 AS new_count, '01bee1e5-0000-d31e-0000-e80ef02b9f27' AS query_id) AS b WHERE u.sample_query_id = b.query_id",
+        )