From c2b2aee76e37d371bfe237ff143efd819b8effaa Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Wed, 19 Jun 2024 15:47:01 -0700
Subject: [PATCH] Fix control flow execution of conditional functions `IF()`,
 `IFNULL()`

---
 internal/formatter.go | 14 ++++++++++++++
 query_test.go         | 15 +++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/internal/formatter.go b/internal/formatter.go
index 25a0df8..3a1960d 100644
--- a/internal/formatter.go
+++ b/internal/formatter.go
@@ -251,6 +251,20 @@ func (n *FunctionCallNode) FormatSQL(ctx context.Context) (string, error) {
 		return "", err
 	}
 	switch funcName {
+	case "zetasqlite_ifnull":
+		return fmt.Sprintf(
+			"CASE WHEN %s IS NULL THEN %s ELSE %s END",
+			args[0],
+			args[1],
+			args[0],
+		), nil
+	case "zetasqlite_if":
+		return fmt.Sprintf(
+			"CASE WHEN %s THEN %s ELSE %s END",
+			args[0],
+			args[1],
+			args[2],
+		), nil
 	case "zetasqlite_case_no_value":
 		var whenStmts []string
 		for i := 0; i < len(args)-1; i += 2 {
diff --git a/query_test.go b/query_test.go
index be86437..da020d9 100644
--- a/query_test.go
+++ b/query_test.go
@@ -478,6 +478,21 @@ FROM UNNEST([1, 2, 3, 4]) AS val`,
 			query:        `SELECT IF("a" = "a", "true", "false")`,
 			expectedRows: [][]interface{}{{"true"}},
 		},
+		{
+			name:         "if with case that causes errors",
+			query:        `SELECT IF(FALSE, ERROR("error case!"), "false")`,
+			expectedRows: [][]interface{}{{"false"}},
+		},
+		{
+			name:         "ifnull with case that causes errors",
+			query:        `SELECT IFNULL("STRING", ERROR("error case!"))`,
+			expectedRows: [][]interface{}{{"STRING"}},
+		},
+		{
+			name:         "case with case that causes errors",
+			query:        `SELECT CASE WHEN FALSE THEN ERROR("error case!") ELSE "false" END`,
+			expectedRows: [][]interface{}{{"false"}},
+		},
 		{
 			name:         "ifnull",
 			query:        `SELECT IFNULL(10, 0)`,