Repository: tobymao/sqlglot
Pull request: Export union/intersect/except builders
fixes #540 

Problem statement:
API support for union/intersect/except
Is there a way to construct a union from an existing sqlglot expression?

I tried the following, and hunted around for methods in the codebase but I didn't see anything:

```
In [9]: import sqlglot as sg

In [10]: t = sg.table('t')

In [11]: t.union(t)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In [11], line 1
----> 1 t.union(t)

AttributeError: 'Table' object has no attribute 'union'

In [12]: sel = sg.select('*').from_(t)

In [13]: sel.union(sel)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In [13], line 1
----> 1 sel.union(sel)

AttributeError: 'Select' object has no attribute 'union'
```

Hints from issue discussion:
A quick workaround:

```python
import sqlglot as sg
from sqlglot import expressions as exp

t = sg.table('t')
sel = sg.select('*').from_(t)

union = exp.Union(this=sel, expression=sel.copy())
```

But a `union` builder method would be sweet, yeah.
Ah, great. Thank you. My current workaround was ... not as nice :)

Commit message: Export union/intersect/except builders

Patch to review:
diff --git a/sqlglot/__init__.py b/sqlglot/__init__.py
index 8291ec2b8b..fc5358e0a8 100644
--- a/sqlglot/__init__.py
+++ b/sqlglot/__init__.py
@@ -8,7 +8,9 @@
     and_,
     column,
     condition,
+    except_,
     from_,
+    intersect,
     maybe_parse,
     not_,
     or_,
@@ -16,6 +18,7 @@
     subquery,
 )
 from sqlglot.expressions import table_ as table
+from sqlglot.expressions import union
 from sqlglot.generator import Generator
 from sqlglot.parser import Parser
 from sqlglot.tokens import Tokenizer, TokenType
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index bc6beaa1c6..80ccc660ca 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -428,6 +428,69 @@ def assert_is(self, type_):
         assert isinstance(self, type_)
         return self
 
+    def union(self, expr, distinct=True, dialect=None, **opts):
+        """
+        Builds a UNION expression.
+
+        Example:
+            >>> import sqlglot
+            >>> sqlglot.parse_one("SELECT * FROM foo").union("SELECT * FROM bla").sql()
+            'SELECT * FROM foo UNION SELECT * FROM bla'
+
+        Args:
+            expr (str or Expression): the SQL code string.
+                If an `Expression` instance is passed, it will be used as-is.
+            distinct (bool): set the DISTINCT flag if and only if this is true.
+            dialect (str): the dialect used to parse the input expression.
+            opts (kwargs): other options to use to parse the input expressions.
+        Returns:
+            Union: the Union expression.
+        """
+        expr = maybe_parse(sql_or_expression=expr, dialect=dialect, **opts)
+        return Union(this=self, expression=expr, distinct=distinct)
+
+    def intersect(self, expr, distinct=True, dialect=None, **opts):
+        """
+        Builds an INTERSECT expression.
+
+        Example:
+            >>> import sqlglot
+            >>> sqlglot.parse_one("SELECT * FROM foo").intersect("SELECT * FROM bla").sql()
+            'SELECT * FROM foo INTERSECT SELECT * FROM bla'
+
+        Args:
+            expr (str or Expression): the SQL code string.
+                If an `Expression` instance is passed, it will be used as-is.
+            distinct (bool): set the DISTINCT flag if and only if this is true.
+            dialect (str): the dialect used to parse the input expression.
+            opts (kwargs): other options to use to parse the input expressions.
+        Returns:
+            Intersect: the Intersect expression
+        """
+        expr = maybe_parse(sql_or_expression=expr, dialect=dialect, **opts)
+        return Intersect(this=self, expression=expr, distinct=distinct)
+
+    def except_(self, expr, distinct=True, dialect=None, **opts):
+        """
+        Builds an EXCEPT expression.
+
+        Example:
+            >>> import sqlglot
+            >>> sqlglot.parse_one("SELECT * FROM foo").except_("SELECT * FROM bla").sql()
+            'SELECT * FROM foo EXCEPT SELECT * FROM bla'
+
+        Args:
+            expr (str or Expression): the SQL code string.
+                If an `Expression` instance is passed, it will be used as-is.
+            distinct (bool): set the DISTINCT flag if and only if this is true.
+            dialect (str): the dialect used to parse the input expression.
+            opts (kwargs): other options to use to parse the input expressions.
+        Returns:
+            Except: the Except expression
+        """
+        expr = maybe_parse(sql_or_expression=expr, dialect=dialect, **opts)
+        return Except(this=self, expression=expr, distinct=distinct)
+
 
 class Condition(Expression):
     def and_(self, *expressions, dialect=None, **opts):
@@ -2795,6 +2858,81 @@ def _wrap_operator(expression):
     return expression
 
 
+def union(left, right, distinct=True, dialect=None, **opts):
+    """
+    Initializes a syntax tree from one UNION expression.
+
+    Example:
+        >>> union("SELECT * FROM foo", "SELECT * FROM bla").sql()
+        'SELECT * FROM foo UNION SELECT * FROM bla'
+
+    Args:
+        left (str or Expression): the SQL code string corresponding to the left-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        right (str or Expression): the SQL code string corresponding to the right-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        distinct (bool): set the DISTINCT flag if and only if this is true.
+        dialect (str): the dialect used to parse the input expression.
+        opts (kwargs): other options to use to parse the input expressions.
+    Returns:
+        Union: the syntax tree for the UNION expression.
+    """
+    left = maybe_parse(sql_or_expression=left, dialect=dialect, **opts)
+    right = maybe_parse(sql_or_expression=right, dialect=dialect, **opts)
+
+    return left.union(right, distinct=distinct)
+
+
+def intersect(left, right, distinct=True, dialect=None, **opts):
+    """
+    Initializes a syntax tree from one INTERSECT expression.
+
+    Example:
+        >>> intersect("SELECT * FROM foo", "SELECT * FROM bla").sql()
+        'SELECT * FROM foo INTERSECT SELECT * FROM bla'
+
+    Args:
+        left (str or Expression): the SQL code string corresponding to the left-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        right (str or Expression): the SQL code string corresponding to the right-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        distinct (bool): set the DISTINCT flag if and only if this is true.
+        dialect (str): the dialect used to parse the input expression.
+        opts (kwargs): other options to use to parse the input expressions.
+    Returns:
+        Intersect: the syntax tree for the INTERSECT expression.
+    """
+    left = maybe_parse(sql_or_expression=left, dialect=dialect, **opts)
+    right = maybe_parse(sql_or_expression=right, dialect=dialect, **opts)
+
+    return left.intersect(right, distinct=distinct)
+
+
+def except_(left, right, distinct=True, dialect=None, **opts):
+    """
+    Initializes a syntax tree from one EXCEPT expression.
+
+    Example:
+        >>> except_("SELECT * FROM foo", "SELECT * FROM bla").sql()
+        'SELECT * FROM foo EXCEPT SELECT * FROM bla'
+
+    Args:
+        left (str or Expression): the SQL code string corresponding to the left-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        right (str or Expression): the SQL code string corresponding to the right-hand side.
+            If an `Expression` instance is passed, it will be used as-is.
+        distinct (bool): set the DISTINCT flag if and only if this is true.
+        dialect (str): the dialect used to parse the input expression.
+        opts (kwargs): other options to use to parse the input expressions.
+    Returns:
+        Except: the syntax tree for the EXCEPT statement.
+    """
+    left = maybe_parse(sql_or_expression=left, dialect=dialect, **opts)
+    right = maybe_parse(sql_or_expression=right, dialect=dialect, **opts)
+
+    return left.except_(right, distinct=distinct)
+
+
 def select(*expressions, dialect=None, **opts):
     """
     Initializes a syntax tree from one or multiple SELECT expressions.
@@ -3001,7 +3139,7 @@ def alias_(expression, alias, table=False, dialect=None, quoted=None, **opts):
             If an Expression instance is passed, this is used as-is.
         alias (str or Identifier): the alias name to use. If the name has
             special characters it is quoted.
-        table (boolean): create a table alias, default false
+        table (bool): create a table alias, default false
         dialect (str): the dialect used to parse the input expression.
         **opts: other options to use to parse the input expressions.
 
diff --git a/tests/test_build.py b/tests/test_build.py
index b5d657c1c4..fe5ac51f85 100644
--- a/tests/test_build.py
+++ b/tests/test_build.py
@@ -1,6 +1,18 @@
 import unittest
 
-from sqlglot import and_, condition, exp, from_, not_, or_, parse_one, select
+from sqlglot import (
+    and_,
+    condition,
+    except_,
+    exp,
+    from_,
+    intersect,
+    not_,
+    or_,
+    parse_one,
+    select,
+    union,
+)
 
 
 class TestBuild(unittest.TestCase):
@@ -320,6 +332,18 @@ def test_build(self):
                 lambda: exp.update("tbl", {"x": 1}, from_="tbl2"),
                 "UPDATE tbl SET x = 1 FROM tbl2",
             ),
+            (
+                lambda: union("SELECT * FROM foo", "SELECT * FROM bla"),
+                "SELECT * FROM foo UNION SELECT * FROM bla",
+            ),
+            (
+                lambda: intersect("SELECT * FROM foo", "SELECT * FROM bla"),
+                "SELECT * FROM foo INTERSECT SELECT * FROM bla",
+            ),
+            (
+                lambda: except_("SELECT * FROM foo", "SELECT * FROM bla"),
+                "SELECT * FROM foo EXCEPT SELECT * FROM bla",
+            ),
         ]:
             with self.subTest(sql):
                 self.assertEqual(expression().sql(dialect[0] if dialect else None), sql)


Code review comments:
diff_hunk: @@ -428,6 +428,69 @@ def assert_is(self, type_):
         assert isinstance(self, type_)
         return self
 
+    def union(self, expr, distinct=True, dialect=None, **opts):
path: sqlglot/expressions.py
text: @user1:
for consistency:
expr -> expression

diff_hunk: @@ -428,6 +428,69 @@ def assert_is(self, type_):
         assert isinstance(self, type_)
         return self
 
+    def union(self, expr, distinct=True, dialect=None, **opts):
+        """
+        Builds a UNION expression.
+
+        Example:
+            >>> import sqlglot
+            >>> sqlglot.parse_one("SELECT * FROM foo").union("SELECT * FROM bla").sql()
+            'SELECT * FROM foo UNION SELECT * FROM bla'
+
+        Args:
+            expr (str or Expression): the SQL code string.
+                If an `Expression` instance is passed, it will be used as-is.
+            distinct (bool): set the DISTINCT flag if and only if this is true.
+            dialect (str): the dialect used to parse the input expression.
+            opts (kwargs): other options to use to parse the input expressions.
+        Returns:
+            Union: the Union expression.
+        """
+        expr = maybe_parse(sql_or_expression=expr, dialect=dialect, **opts)
+        return Union(this=self, expression=expr, distinct=distinct)
path: sqlglot/expressions.py
text: @user1:
I think this can be `union(left=self, right=expr, distinct=distinct, dialect=dialect, **opts)`