Fix control flow execution of conditional functions IF(), IFNULL() (

#46)
Recidiviz · Jun 24, 2024 · 76b0d28 · 76b0d28
1 parent a81591f
commit 76b0d28
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git 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
@@ -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)`,