Repository: tobymao/sqlglot

Issue:
Optimizer fails to coerce BOOLEAN (Snowflake dialect)
Example:
```python
from sqlglot.executor import execute

tables = {
    "public": {
        "venues": [
            {
                "name": "my_pizza_place",
                "my_bool": True,
            }
        ],
    },
}

schema = {
    "public": {
        "venues": {
            "name": "VARCHAR(16777216)",
            "my_bool": "BOOLEAN",
        },
    },
}

q = """
SELECT
    my_bool
FROM
    public.venues
WHERE
  IFF(True, my_bool=True, my_bool in (False, True))
;"""

r = execute(
    q,
    schema=schema,
    tables=tables,
    read="snowflake",
)
```
Fails with

```log
  File "/.../python3.9/site-packages/sqlglot/optimizer/annotate_types.py", line 309, in _maybe_coerce
    return type2 if type2 in self.coerces_to[type1] else type1
KeyError: <Type.BOOLEAN: 'BOOLEAN'>
```

The issues seems to get fixed by adding `BOOLEAN` and an empty set to [this](https://github.com/tobymao/sqlglot/blob/main/sqlglot/optimizer/annotate_types.py#L173) list, but I'm not sure what the consequences of that change would be. Using version `sqlglot-10.2.2`


Discussion:

Resolving patch:
diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py
--- a/sqlglot/optimizer/annotate_types.py
+++ b/sqlglot/optimizer/annotate_types.py
@@ -45,6 +45,8 @@ class TypeAnnotator:
         exp.TryCast: lambda self, expr: self._annotate_with_type(expr, expr.args["to"].this),
         exp.DataType: lambda self, expr: self._annotate_with_type(expr, expr.this),
         exp.Alias: lambda self, expr: self._annotate_unary(expr),
+        exp.Between: lambda self, expr: self._annotate_with_type(expr, exp.DataType.Type.BOOLEAN),
+        exp.In: lambda self, expr: self._annotate_with_type(expr, exp.DataType.Type.BOOLEAN),
         exp.Literal: lambda self, expr: self._annotate_literal(expr),
         exp.Boolean: lambda self, expr: self._annotate_with_type(expr, exp.DataType.Type.BOOLEAN),
         exp.Null: lambda self, expr: self._annotate_with_type(expr, exp.DataType.Type.NULL),
@@ -306,7 +308,7 @@ def _maybe_coerce(self, type1, type2):
         if exp.DataType.Type.UNKNOWN in (type1, type2):
             return exp.DataType.Type.UNKNOWN
 
-        return type2 if type2 in self.coerces_to[type1] else type1
+        return type2 if type2 in self.coerces_to.get(type1, {}) else type1
 
     def _annotate_binary(self, expression):
         self._annotate_args(expression)


Test patch:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -402,21 +402,22 @@ def test_derived_tables_column_annotation(self):
         self.assertEqual(addition.expression.type.this, exp.DataType.Type.FLOAT)
 
     def test_cte_column_annotation(self):
-        schema = {"x": {"cola": "CHAR"}, "y": {"colb": "TEXT"}}
+        schema = {"x": {"cola": "CHAR"}, "y": {"colb": "TEXT", "colc": "BOOLEAN"}}
         sql = """
             WITH tbl AS (
-                SELECT x.cola + 'bla' AS cola, y.colb AS colb
+                SELECT x.cola + 'bla' AS cola, y.colb AS colb, y.colc AS colc
                 FROM (
                     SELECT x.cola AS cola
                     FROM x AS x
                 ) AS x
                 JOIN (
-                    SELECT y.colb AS colb
+                    SELECT y.colb AS colb, y.colc AS colc
                     FROM y AS y
                 ) AS y
             )
             SELECT tbl.cola + tbl.colb + 'foo' AS col
             FROM tbl AS tbl
+            WHERE tbl.colc = True
         """
 
         expression = annotate_types(parse_one(sql), schema=schema)
@@ -433,6 +434,9 @@ def test_cte_column_annotation(self):
         self.assertEqual(inner_addition.left.type.this, exp.DataType.Type.VARCHAR)
         self.assertEqual(inner_addition.right.type.this, exp.DataType.Type.TEXT)
 
+        # WHERE tbl.colc = True
+        self.assertEqual(expression.args["where"].this.type.this, exp.DataType.Type.BOOLEAN)
+
         cte_select = expression.args["with"].expressions[0].this
         self.assertEqual(
             cte_select.expressions[0].type.this, exp.DataType.Type.VARCHAR
@@ -440,6 +444,9 @@ def test_cte_column_annotation(self):
         self.assertEqual(
             cte_select.expressions[1].type.this, exp.DataType.Type.TEXT
         )  # y.colb AS colb
+        self.assertEqual(
+            cte_select.expressions[2].type.this, exp.DataType.Type.BOOLEAN
+        )  # y.colc AS colc
 
         cte_select_addition = cte_select.expressions[0].this  # x.cola + 'bla'
         self.assertEqual(cte_select_addition.type.this, exp.DataType.Type.VARCHAR)
@@ -518,3 +525,10 @@ def test_nullable_annotation(self):
         self.assertEqual(expression.type, nullable)
         self.assertEqual(expression.left.type.this, exp.DataType.Type.NULL)
         self.assertEqual(expression.right.type.this, exp.DataType.Type.BOOLEAN)
+
+    def test_predicate_annotation(self):
+        expression = annotate_types(parse_one("x BETWEEN a AND b"))
+        self.assertEqual(expression.type.this, exp.DataType.Type.BOOLEAN)
+
+        expression = annotate_types(parse_one("x IN (a, b, c, d)"))
+        self.assertEqual(expression.type.this, exp.DataType.Type.BOOLEAN)