feat: support setting the timezone

fakesnow/fakes.py

@@ -145,6 +145,7 @@ def execute(
             .transform(transforms.values_columns)
             .transform(transforms.to_date)
             .transform(transforms.object_construct)
+            .transform(transforms.set_timezone)
         )
 
         sql = transformed.sql(dialect="duckdb")

fakesnow/transforms.py

@@ -373,6 +373,32 @@ def set_schema(expression: exp.Expression, current_database: str | None) -> exp.
     return expression
 
 
+def set_timezone(expression: exp.Expression) -> exp.Expression:
+    """Set timezone
+
+    Example:
+        >>> import sqlglot
+        >>> sqlglot.parse_one("ALTER SESSION SET TIMEZONE='UTC').transform(set_timezone).sql()
+        "SET TIMEZONE='UTC'"
+    Args:
+        expression (exp.Expression): the expression that will be transformed.
+
+    Returns:
+        exp.Expression: The transformed expression.
+    """
+    if (
+        isinstance(expression, exp.Command)
+        and isinstance(expression.this, str)
+        and expression.this.upper() == "ALTER"
+        and isinstance(expression.expression, str)
+        and (set_index := expression.expression.upper().find("SET")) != -1
+        and "TIMEZONE" in expression.expression.upper()
+    ):
+        return exp.Command(this="SET", expression=exp.Literal.string(expression.expression[set_index + 3 :]))
+
+    return expression
+
+
 def tag(expression: exp.Expression) -> exp.Expression:
     """Handle tags. Transfer tags into upserts of the tag table.
 

tests/test_transforms.py

@@ -14,6 +14,7 @@
     regex,
     semi_structured_types,
     set_schema,
+    set_timezone,
     tag,
     to_date,
     upper_case_unquoted_identifiers,
@@ -84,6 +85,10 @@ def test_object_construct() -> None:
     )
 
 
+def test_set_timezone() -> None:
+    assert sqlglot.parse_one("ALTER SESSION SET TIMEZONE='UTC'").transform(set_timezone).sql() == "SET TIMEZONE='UTC'"
+
+
 def test_parse_json() -> None:
     assert (
         sqlglot.parse_one("""insert into table1 (name) select parse_json('{"first":"foo", "last":"bar"}')""")