Repository: tobymao/sqlglot

Issue:
Bug in transform since 23.0.0
Hello,

I have this bug since the version 23.0.0, if I add a filter in the transform it's not taken in account.

Here's a snippet to reproduce the bug

```python
from sqlglot import parse_one, exp

def t(e):
    if str(e) == "COUNT(1)":
        return exp.Filter(
         this=e,
         expression=exp.Where(
         this=exp.EQ(
            this=exp.Literal(this=1, is_string=False),
            expression=exp.Literal(this=1, is_string=False))))
    return e

print(parse_one("COUNT(1)", dialect="postgres").transform(t).sql(dialect="postgres"))

print(parse_one("SELECT COUNT(1) FROM table", dialect="postgres").transform(t).sql(dialect="postgres"))
```

The output is :

```
COUNT(1) FILTER(WHERE 1 = 1)
SELECT COUNT(1) FROM table
```

I would expect to have

```
COUNT(1) FILTER(WHERE 1 = 1)
SELECT COUNT(1)  FILTER(WHERE 1 = 1) FROM table
```




Discussion:

Patch:
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index 9837025a..470fd7c3 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -38,10 +38,12 @@ def _ts_or_ds_add_sql(self: DuckDB.Generator, expression: exp.TsOrDsAdd) -> str:
     return f"CAST({this} AS {self.sql(expression.return_type)}) + {interval}"
 
 
-def _date_delta_sql(self: DuckDB.Generator, expression: exp.DateAdd | exp.DateSub) -> str:
+def _date_delta_sql(
+    self: DuckDB.Generator, expression: exp.DateAdd | exp.DateSub | exp.TimeAdd
+) -> str:
     this = self.sql(expression, "this")
     unit = self.sql(expression, "unit").strip("'") or "DAY"
-    op = "+" if isinstance(expression, exp.DateAdd) else "-"
+    op = "+" if isinstance(expression, (exp.DateAdd, exp.TimeAdd)) else "-"
     return f"{this} {op} {self.sql(exp.Interval(this=expression.expression, unit=unit))}"
 
 
@@ -427,6 +429,7 @@ class DuckDB(Dialect):
                 "EPOCH", self.func("STRPTIME", e.this, self.format_time(e))
             ),
             exp.Struct: _struct_sql,
+            exp.TimeAdd: _date_delta_sql,
             exp.Timestamp: no_timestamp_sql,
             exp.TimestampDiff: lambda self, e: self.func(
                 "DATE_DIFF", exp.Literal.string(e.unit), e.expression, e.this
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 0ffe4cf0..c13c15ab 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -78,6 +78,17 @@ def _build_datediff(args: t.List) -> exp.DateDiff:
     )
 
 
+def _build_date_time_add(expr_type: t.Type[E]) -> t.Callable[[t.List], E]:
+    def _builder(args: t.List) -> E:
+        return expr_type(
+            this=seq_get(args, 2),
+            expression=seq_get(args, 1),
+            unit=_map_date_part(seq_get(args, 0)),
+        )
+
+    return _builder
+
+
 # https://docs.snowflake.com/en/sql-reference/functions/div0
 def _build_if_from_div0(args: t.List) -> exp.If:
     cond = exp.EQ(this=seq_get(args, 1), expression=exp.Literal.number(0))
@@ -345,11 +356,7 @@ class Snowflake(Dialect):
             "CONVERT_TIMEZONE": _build_convert_timezone,
             "DATE": _build_datetime("DATE", exp.DataType.Type.DATE),
             "DATE_TRUNC": _date_trunc_to_time,
-            "DATEADD": lambda args: exp.DateAdd(
-                this=seq_get(args, 2),
-                expression=seq_get(args, 1),
-                unit=_map_date_part(seq_get(args, 0)),
-            ),
+            "DATEADD": _build_date_time_add(exp.DateAdd),
             "DATEDIFF": _build_datediff,
             "DIV0": _build_if_from_div0,
             "FLATTEN": exp.Explode.from_arg_list,
@@ -367,7 +374,9 @@ class Snowflake(Dialect):
             "REGEXP_SUBSTR": exp.RegexpExtract.from_arg_list,
             "RLIKE": exp.RegexpLike.from_arg_list,
             "SQUARE": lambda args: exp.Pow(this=seq_get(args, 0), expression=exp.Literal.number(2)),
+            "TIMEADD": _build_date_time_add(exp.TimeAdd),
             "TIMEDIFF": _build_datediff,
+            "TIMESTAMPADD": _build_date_time_add(exp.DateAdd),
             "TIMESTAMPDIFF": _build_datediff,
             "TIMESTAMPFROMPARTS": _build_timestamp_from_parts,
             "TIMESTAMP_FROM_PARTS": _build_timestamp_from_parts,
@@ -792,6 +801,7 @@ class Snowflake(Dialect):
             ),
             exp.StrToTime: lambda self, e: self.func("TO_TIMESTAMP", e.this, self.format_time(e)),
             exp.Stuff: rename_func("INSERT"),
+            exp.TimeAdd: date_delta_sql("TIMEADD"),
             exp.TimestampDiff: lambda self, e: self.func(
                 "TIMESTAMPDIFF", e.unit, e.expression, e.this
             ),
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index c24fdff2..57a1044f 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -320,19 +320,39 @@ class Expression(metaclass=_Expression):
             value.index = len(values)
         values.append(value)
 
-    def set(self, arg_key: str, value: t.Any) -> None:
+    def set(self, arg_key: str, value: t.Any, index: t.Optional[int] = None) -> None:
         """
         Sets arg_key to value.
 
         Args:
             arg_key: name of the expression arg.
             value: value to set the arg to.
-        """
-        if value is None:
+            index: if the arg is a list, this specifies what position to add the value in it.
+        """
+        if index is not None:
+            expressions = self.args.get(arg_key) or []
+
+            if seq_get(expressions, index) is None:
+                return
+            if value is None:
+                expressions.pop(index)
+                for v in expressions[index:]:
+                    v.index = v.index - 1
+                return
+
+            if isinstance(value, list):
+                expressions.pop(index)
+                expressions[index:index] = value
+            else:
+                expressions[index] = value
+
+            value = expressions
+        elif value is None:
             self.args.pop(arg_key, None)
-        else:
-            self.args[arg_key] = value
-            self._set_parent(arg_key, value)
+            return
+
+        self.args[arg_key] = value
+        self._set_parent(arg_key, value, index)
 
     def _set_parent(self, arg_key: str, value: t.Any, index: t.Optional[int] = None) -> None:
         if hasattr(value, "parent"):
@@ -579,13 +599,13 @@ class Expression(metaclass=_Expression):
         new_node = None
 
         for node in (self.copy() if copy else self).dfs(prune=lambda n: n is not new_node):
+            parent, arg_key, index = node.parent, node.arg_key, node.index
             new_node = fun(node, *args, **kwargs)
 
-            if root:
-                if new_node is not node:
-                    node.replace(new_node)
-            else:
+            if not root:
                 root = new_node
+            elif new_node is not node:
+                parent.set(arg_key, new_node, index)
 
         assert root
         return root.assert_is(Expression)
@@ -617,37 +637,16 @@ class Expression(metaclass=_Expression):
         """
         parent = self.parent
 
-        if not parent:
+        if not parent or parent is expression:
             return expression
 
         key = self.arg_key
         value = parent.args.get(key)
-        exp_is_list = type(expression) is list
 
-        if type(value) is list:
-            index = self.index
-
-            if exp_is_list:
-                value.pop(index)
-                value[index:index] = expression
-                parent._set_parent(key, value)
-            else:
-                if expression is None:
-                    value.pop(index)
-
-                    for v in value[index:]:
-                        v.index = v.index - 1
-                else:
-                    value[index] = expression
-                    parent._set_parent(key, expression, index=index)
-        elif value is not None:
-            if expression is None:
-                parent.args.pop(key)
-            else:
-                if exp_is_list and value.parent:
-                    value.parent.replace(expression)
-                else:
-                    parent.set(key, expression)
+        if type(expression) is list and isinstance(value, Expression) and value.parent:
+            value.parent.replace(expression)
+        else:
+            parent.set(key, expression, self.index)
 
         if expression is not self:
             self.parent = None


Test patch:
diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
index 917230b2..5cfcfea2 100644
--- a/tests/dialects/test_duckdb.py
+++ b/tests/dialects/test_duckdb.py
@@ -15,6 +15,17 @@ class TestDuckDB(Validator):
             "WITH _data AS (SELECT [STRUCT(1 AS a, 2 AS b), STRUCT(2 AS a, 3 AS b)] AS col) SELECT col.b FROM _data, UNNEST(_data.col) AS col WHERE col.a = 1",
         )
 
+        self.validate_all(
+            "SELECT CAST('09:05:03' AS TIME) + INTERVAL 2 HOUR",
+            read={
+                "bigquery": "SELECT TIME_ADD(CAST('09:05:03' AS TIME), INTERVAL 2 HOUR)",
+                "snowflake": "SELECT TIMEADD(HOUR, 2, TO_TIME('09:05:03'))",
+            },
+            write={
+                "duckdb": "SELECT CAST('09:05:03' AS TIME) + INTERVAL '2' HOUR",
+                "snowflake": "SELECT CAST('09:05:03' AS TIME) + INTERVAL '2 HOUR'",
+            },
+        )
         self.validate_all(
             'STRUCT_PACK("a b" := 1)',
             write={
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 00e6169a..4d7d97c5 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -40,6 +40,7 @@ WHERE
   )""",
         )
 
+        self.validate_identity("SELECT TIMEADD(HOUR, 2, CAST('09:05:03' AS TIME))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS MAP(VARCHAR, INT))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS OBJECT(a CHAR NOT NULL))")
         self.validate_identity("SELECT CAST([1, 2, 3] AS ARRAY(INT))")
@@ -956,6 +957,9 @@ WHERE
         )
         self.validate_all(
             "DATEADD(DAY, 5, CAST('2008-12-25' AS DATE))",
+            read={
+                "snowflake": "TIMESTAMPADD(DAY, 5, CAST('2008-12-25' AS DATE))",
+            },
             write={
                 "bigquery": "DATE_ADD(CAST('2008-12-25' AS DATE), INTERVAL 5 DAY)",
                 "snowflake": "DATEADD(DAY, 5, CAST('2008-12-25' AS DATE))",
diff --git a/tests/test_expressions.py b/tests/test_expressions.py
index 547f6356..54bc223c 100644
--- a/tests/test_expressions.py
+++ b/tests/test_expressions.py
@@ -501,6 +501,18 @@ class TestExpressions(unittest.TestCase):
 
         self.assertEqual(expression.transform(fun).sql(), "FUN(a)")
 
+    def test_transform_with_parent_mutation(self):
+        expression = parse_one("SELECT COUNT(1) FROM table")
+
+        def fun(node):
+            if str(node) == "COUNT(1)":
+                # node gets silently mutated here - its parent points to the filter node
+                return exp.Filter(this=node, expression=exp.Where(this=exp.true()))
+            return node
+
+        transformed = expression.transform(fun)
+        self.assertEqual(transformed.sql(), "SELECT COUNT(1) FILTER(WHERE TRUE) FROM table")
+
     def test_transform_multiple_children(self):
         expression = parse_one("SELECT * FROM x")