feat: Firebolt sqlglot dialect (#31825)

tests/unit_tests/sql/parse_tests.py

@@ -301,7 +301,7 @@ def test_format_no_dialect() -> None:
     Test format with an engine that has no corresponding dialect.
     """
     assert (
-        SQLScript("SELECT col FROM t WHERE col NOT IN (1, 2)", "firebolt").format()
+        SQLScript("SELECT col FROM t WHERE col NOT IN (1, 2)", "dremio").format()
         == "SELECT col\nFROM t\nWHERE col NOT IN (1,\n                  2)"
     )
 
@@ -311,7 +311,7 @@ def test_split_no_dialect() -> None:
     Test the statement split when the engine has no corresponding dialect.
     """
     sql = "SELECT col FROM t WHERE col NOT IN (1, 2); SELECT * FROM t; SELECT foo"
-    statements = SQLScript(sql, "firebolt").statements
+    statements = SQLScript(sql, "dremio").statements
     assert len(statements) == 3
     assert statements[0]._sql == "SELECT col FROM t WHERE col NOT IN (1, 2)"
     assert statements[1]._sql == "SELECT * FROM t"
@@ -1112,4 +1112,37 @@ WHERE anon_1.a > 1
   AND anon_1.b = 2"""
 
     assert SQLStatement(sql, "sqlite").optimize().format() == optimized
-    assert SQLStatement(sql, "firebolt").optimize().format() == not_optimized
+    assert SQLStatement(sql, "dremio").optimize().format() == not_optimized
+
+
+def test_firebolt() -> None:
+    """
+    Test that Firebolt 3rd party dialect is registered correctly.
+
+    We need a custom dialect for Firebolt because it parses `NOT col IN (1, 2)` as
+    `(NOT col) IN (1, 2)` instead of `NOT (col IN (1, 2))`, which will fail when `col`
+    is not a boolean.
+
+    Note that `NOT col = 1` works as expected in Firebolt, parsing as `NOT (col = 1)`.
+    """
+    sql = "SELECT col NOT IN (1, 2) FROM tbl"
+    assert (
+        SQLStatement(sql, "firebolt").format()
+        == """
+SELECT
+  NOT (
+    col IN (1, 2)
+  )
+FROM tbl
+    """.strip()
+    )
+
+    sql = "SELECT NOT col = 1 FROM tbl"
+    assert (
+        SQLStatement(sql, "firebolt").format()
+        == """
+SELECT
+  NOT col = 1
+FROM tbl
+    """.strip()
+    )