Fix(snowflake): handle DIV0 case where divident is null (#3975)

tobymao · Aug 26, 2024 · d37a5bb · d37a5bb
1 parent f7e4e4a
commit d37a5bb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
@@ -93,7 +93,9 @@ def _builder(args: t.List) -> E:
 
 # 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))
+ cond = exp.EQ(this=seq_get(args, 1), expression=exp.Literal.number(0)).and_(
+ exp.Is(this=seq_get(args, 0), expression=exp.null()).not_()
+ )
  true = exp.Literal.number(0)
  false = exp.Div(this=seq_get(args, 0), expression=seq_get(args, 1))
  return exp.If(this=cond, true=true, false=false)

diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -598,12 +598,12 @@ def test_snowflake(self):
  self.validate_all(
  "DIV0(foo, bar)",
  write={
- "snowflake": "IFF(bar = 0, 0, foo / bar)",
- "sqlite": "IIF(bar = 0, 0, CAST(foo AS REAL) / bar)",
- "presto": "IF(bar = 0, 0, CAST(foo AS DOUBLE) / bar)",
- "spark": "IF(bar = 0, 0, foo / bar)",
- "hive": "IF(bar = 0, 0, foo / bar)",
- "duckdb": "CASE WHEN bar = 0 THEN 0 ELSE foo / bar END",
+ "snowflake": "IFF(bar = 0 AND NOT foo IS NULL, 0, foo / bar)",
+ "sqlite": "IIF(bar = 0 AND NOT foo IS NULL, 0, CAST(foo AS REAL) / bar)",
+ "presto": "IF(bar = 0 AND NOT foo IS NULL, 0, CAST(foo AS DOUBLE) / bar)",
+ "spark": "IF(bar = 0 AND NOT foo IS NULL, 0, foo / bar)",
+ "hive": "IF(bar = 0 AND NOT foo IS NULL, 0, foo / bar)",
+ "duckdb": "CASE WHEN bar = 0 AND NOT foo IS NULL THEN 0 ELSE foo / bar END",
  },
  )
  self.validate_all(