feat(snowflake): Transpile exp.SafeDivide (#4294)

VaggelisD · web-flow · commit 1a26bff61931 · 2024-10-25T18:57:25.000+03:00
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
@@ -1017,10 +1017,10 @@ def no_recursive_cte_sql(self: Generator, expression: exp.With) -> str:
     return self.with_sql(expression)
 
 
-def no_safe_divide_sql(self: Generator, expression: exp.SafeDivide) -> str:
+def no_safe_divide_sql(self: Generator, expression: exp.SafeDivide, if_sql: str = "IF") -> str:
     n = self.sql(expression, "this")
     d = self.sql(expression, "expression")
-    return f"IF(({d}) <> 0, ({n}) / ({d}), NULL)"
+    return f"{if_sql}(({d}) <> 0, ({n}) / ({d}), NULL)"
 
 
 def no_tablesample_sql(self: Generator, expression: exp.TableSample) -> str:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
@@ -22,6 +22,7 @@
     timestrtotime_sql,
     var_map_sql,
     map_date_part,
+    no_safe_divide_sql,
 )
 from sqlglot.helper import flatten, is_float, is_int, seq_get
 from sqlglot.tokens import TokenType
@@ -830,6 +831,7 @@ class Generator(generator.Generator):
                     _unnest_generate_date_array,
                 ]
             ),
+            exp.SafeDivide: lambda self, e: no_safe_divide_sql(self, e, "IFF"),
             exp.SHA: rename_func("SHA1"),
             exp.StarMap: rename_func("OBJECT_CONSTRUCT"),
             exp.StartsWith: rename_func("STARTSWITH"),
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -1533,6 +1533,21 @@ def test_bigquery(self):
 
         self.validate_identity("SELECT * FROM a-b c", "SELECT * FROM a-b AS c")
 
+        self.validate_all(
+            "SAFE_DIVIDE(x, y)",
+            write={
+                "bigquery": "SAFE_DIVIDE(x, y)",
+                "duckdb": "IF((y) <> 0, (x) / (y), NULL)",
+                "presto": "IF((y) <> 0, (x) / (y), NULL)",
+                "trino": "IF((y) <> 0, (x) / (y), NULL)",
+                "hive": "IF((y) <> 0, (x) / (y), NULL)",
+                "spark2": "IF((y) <> 0, (x) / (y), NULL)",
+                "spark": "IF((y) <> 0, (x) / (y), NULL)",
+                "databricks": "IF((y) <> 0, (x) / (y), NULL)",
+                "snowflake": "IFF((y) <> 0, (x) / (y), NULL)",
+            },
+        )
+
     def test_errors(self):
         with self.assertRaises(TokenError):
             transpile("'\\'", read="bigquery")
diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
@@ -619,12 +619,6 @@ def test_duckdb(self):
                 "spark": "ARRAY_SUM(ARRAY(1, 2))",
             },
         )
-        self.validate_all(
-            "IF((y) <> 0, (x) / (y), NULL)",
-            read={
-                "bigquery": "SAFE_DIVIDE(x, y)",
-            },
-        )
         self.validate_all(
             "STRUCT_PACK(x := 1, y := '2')",
             write={
