From ed459f5a2b73a78b6b37bf01acef31d9c5c8ed55 Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Mon, 1 Jul 2024 03:37:52 -0700
Subject: [PATCH] Fix control flow execution of conditional functions `IF()`,
 `IFNULL()` (#216)

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

diff --git a/internal/formatter.go b/internal/formatter.go
index 4b0dba4..2cbaa84 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 de6c697..a7541d1 100644
--- a/query_test.go
+++ b/query_test.go
@@ -479,6 +479,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)`,