Repository: tobymao/sqlglot

Issue:
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'
```


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 :)

Resolving patch:
diff --git a/sqlglot/__init__.py b/sqlglot/__init__.py
--- 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
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -508,7 +508,69 @@ def named_selects(self):
         return [select.alias_or_name for select in self.selects]
 
 
-class UDTF(DerivedTable):
+class Unionable:
+    def union(self, expression, 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:
+            expression (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.
+        """
+        return union(left=self, right=expression, distinct=distinct, dialect=dialect, **opts)
+
+    def intersect(self, expression, 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:
+            expression (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
+        """
+        return intersect(left=self, right=expression, distinct=distinct, dialect=dialect, **opts)
+
+    def except_(self, expression, 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:
+            expression (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
+        """
+        return except_(left=self, right=expression, distinct=distinct, dialect=dialect, **opts)
+
+
+class UDTF(DerivedTable, Unionable):
     pass
 
 
@@ -975,7 +1037,7 @@ class Tuple(Expression):
     arg_types = {"expressions": False}
 
 
-class Subqueryable:
+class Subqueryable(Unionable):
     def subquery(self, alias=None, copy=True):
         """
         Convert this expression to an aliased expression that can be used as a Subquery.
@@ -1658,7 +1720,7 @@ def selects(self):
         return self.expressions
 
 
-class Subquery(DerivedTable):
+class Subquery(DerivedTable, Unionable):
     arg_types = {
         "this": True,
         "alias": False,
@@ -2795,6 +2857,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 Union(this=left, expression=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 Intersect(this=left, expression=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 Except(this=left, expression=right, distinct=distinct)
+
+
 def select(*expressions, dialect=None, **opts):
     """
     Initializes a syntax tree from one or multiple SELECT expressions.
@@ -3001,7 +3138,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.
 


Test patch:
diff --git a/tests/test_build.py b/tests/test_build.py
--- 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,38 @@ 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: parse_one("SELECT * FROM foo").union("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: parse_one("SELECT * FROM foo").intersect("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",
+            ),
+            (
+                lambda: parse_one("SELECT * FROM foo").except_("SELECT * FROM bla"),
+                "SELECT * FROM foo EXCEPT SELECT * FROM bla",
+            ),
+            (
+                lambda: parse_one("(SELECT * FROM foo)").union("SELECT * FROM bla"),
+                "(SELECT * FROM foo) UNION SELECT * FROM bla",
+            ),
+            (
+                lambda: parse_one("(SELECT * FROM foo)").union("SELECT * FROM bla", distinct=False),
+                "(SELECT * FROM foo) UNION ALL SELECT * FROM bla",
+            ),
         ]:
             with self.subTest(sql):
                 self.assertEqual(expression().sql(dialect[0] if dialect else None), sql)