Repository: tobymao/sqlglot

Issue:
EXCLUDE indexes support for PostgreSQL dialect
Thanks for adding INCLUDE indexes in #2855, we've updated and now it works. There are [EXCLUDE indexes in PostgreSQL](https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-EXCLUDE), which would be great to support. Here is the example:

```
Python 3.8.10 (default, Nov 22 2023, 10:22:35) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sqlglot
>>> sqlglot.parse_one("""
... CREATE TABLE ht (
...         vid INTEGER NOT NULL, 
...         fid INTEGER NOT NULL, 
...         nid INTEGER NOT NULL CHECK (nid > vid), 
...         vop CHAR(1) NOT NULL, 
...         nop CHAR(1) NOT NULL, 
...         geom GEOMETRY(POINTZ, 3857), 
...         fld_i INTEGER, 
...         fld_t TEXT, 
...         fld_d DATE, 
...         PRIMARY KEY (vid, fid), 
...         CONSTRAINT ht_vid_nid_fid_idx EXCLUDE USING gist (int4range(vid, nid) WITH &&, int4range(fid, fid, '[]') WITH &&)
... )
... """, read="postgres")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/__init__.py", line 124, in parse_one
    result = dialect.parse(sql, **opts)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/dialects/dialect.py", line 477, in parse
    return self.parser(**opts).parse(self.tokenize(sql), sql)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1099, in parse
    return self._parse(
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1165, in _parse
    expressions.append(parse_method(self))
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1367, in _parse_statement
    return self.STATEMENT_PARSERS[self._prev.token_type](self)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 616, in <lambda>
    TokenType.CREATE: lambda self: self._parse_create(),
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1487, in _parse_create
    this = self._parse_schema(this=table_parts)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 4288, in _parse_schema
    args = self._parse_csv(lambda: self._parse_constraint() or self._parse_field_def())
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 5367, in _parse_csv
    parse_result = parse_method()
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 4288, in <lambda>
    args = self._parse_csv(lambda: self._parse_constraint() or self._parse_field_def())
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 4453, in _parse_constraint
    return self.expression(
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1229, in expression
    return self.validate_expression(instance)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1249, in validate_expression
    self.raise_error(error_message)
  File "/opt/ngw/env/lib/python3.8/site-packages/sqlglot/parser.py", line 1209, in raise_error
    raise error
sqlglot.errors.ParseError: Required keyword: 'expressions' missing for <class 'sqlglot.expressions.Constraint'>. Line 13, Col: 38.
   TEXT, 
        fld_d DATE, 
        PRIMARY KEY (vid, fid), 
        CONSTRAINT ht_vid_nid_fid_idx EXCLUDE USING gist (int4range(vid, nid) WITH &&, int4range(fid, fid, '[]') WITH &&)
)

>>> sqlglot.__version__
'22.2.1'
```


Discussion:

Resolving patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -1580,6 +1580,15 @@ class EncodeColumnConstraint(ColumnConstraintKind):
     pass
 
 
+# https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-EXCLUDE
+class ExcludeColumnConstraint(ColumnConstraintKind):
+    pass
+
+
+class WithOperator(Expression):
+    arg_types = {"this": True, "op": True}
+
+
 class GeneratedAsIdentityColumnConstraint(ColumnConstraintKind):
     # this: True -> ALWAYS, this: False -> BY DEFAULT
     arg_types = {
@@ -1854,14 +1863,22 @@ class Index(Expression):
     arg_types = {
         "this": False,
         "table": False,
-        "using": False,
-        "where": False,
-        "columns": False,
         "unique": False,
         "primary": False,
         "amp": False,  # teradata
+        "params": False,
+    }
+
+
+class IndexParameters(Expression):
+    arg_types = {
+        "using": False,
         "include": False,
-        "partition_by": False,  # teradata
+        "columns": False,
+        "with_storage": False,
+        "partition_by": False,
+        "tablespace": False,
+        "where": False,
     }
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -89,6 +89,7 @@ class Generator(metaclass=_Generator):
         exp.DateFormatColumnConstraint: lambda self, e: f"FORMAT {self.sql(e, 'this')}",
         exp.DefaultColumnConstraint: lambda self, e: f"DEFAULT {self.sql(e, 'this')}",
         exp.EncodeColumnConstraint: lambda self, e: f"ENCODE {self.sql(e, 'this')}",
+        exp.ExcludeColumnConstraint: lambda self, e: f"EXCLUDE {self.sql(e, 'this').lstrip()}",
         exp.ExecuteAsProperty: lambda self, e: self.naked_property(e),
         exp.ExternalProperty: lambda *_: "EXTERNAL",
         exp.GlobalProperty: lambda *_: "GLOBAL",
@@ -140,6 +141,7 @@ class Generator(metaclass=_Generator):
         exp.VarMap: lambda self, e: self.func("MAP", e.args["keys"], e.args["values"]),
         exp.VolatileProperty: lambda *_: "VOLATILE",
         exp.WithJournalTableProperty: lambda self, e: f"WITH JOURNAL TABLE={self.sql(e, 'this')}",
+        exp.WithOperator: lambda self, e: f"{self.sql(e, 'this')} WITH {self.sql(e, 'op')}",
     }
 
     # Whether null ordering is supported in order by
@@ -1207,17 +1209,9 @@ def hint_sql(self, expression: exp.Hint) -> str:
 
         return f" /*+ {self.expressions(expression, sep=self.QUERY_HINT_SEP).strip()} */"
 
-    def index_sql(self, expression: exp.Index) -> str:
-        unique = "UNIQUE " if expression.args.get("unique") else ""
-        primary = "PRIMARY " if expression.args.get("primary") else ""
-        amp = "AMP " if expression.args.get("amp") else ""
-        name = self.sql(expression, "this")
-        name = f"{name} " if name else ""
-        table = self.sql(expression, "table")
-        table = f"{self.INDEX_ON} {table}" if table else ""
+    def indexparameters_sql(self, expression: exp.IndexParameters) -> str:
         using = self.sql(expression, "using")
         using = f" USING {using}" if using else ""
-        index = "INDEX " if not table else ""
         columns = self.expressions(expression, key="columns", flat=True)
         columns = f"({columns})" if columns else ""
         partition_by = self.expressions(expression, key="partition_by", flat=True)
@@ -1226,7 +1220,26 @@ def index_sql(self, expression: exp.Index) -> str:
         include = self.expressions(expression, key="include", flat=True)
         if include:
             include = f" INCLUDE ({include})"
-        return f"{unique}{primary}{amp}{index}{name}{table}{using}{columns}{include}{partition_by}{where}"
+        with_storage = self.expressions(expression, key="with_storage", flat=True)
+        with_storage = f" WITH {with_storage}" if with_storage else ""
+        tablespace = self.sql(expression, "tablespace")
+        tablespace = f" USING INDEX TABLESPACE {tablespace}" if tablespace else ""
+
+        return f"{using}{columns}{include}{with_storage}{tablespace}{partition_by}{where}"
+
+    def index_sql(self, expression: exp.Index) -> str:
+        unique = "UNIQUE " if expression.args.get("unique") else ""
+        primary = "PRIMARY " if expression.args.get("primary") else ""
+        amp = "AMP " if expression.args.get("amp") else ""
+        name = self.sql(expression, "this")
+        name = f"{name} " if name else ""
+        table = self.sql(expression, "table")
+        table = f"{self.INDEX_ON} {table}" if table else ""
+
+        index = "INDEX " if not table else ""
+
+        params = self.sql(expression, "params")
+        return f"{unique}{primary}{amp}{index}{name}{table}{params}"
 
     def identifier_sql(self, expression: exp.Identifier) -> str:
         text = expression.name
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -838,6 +838,9 @@ class Parser(metaclass=_Parser):
             exp.DefaultColumnConstraint, this=self._parse_bitwise()
         ),
         "ENCODE": lambda self: self.expression(exp.EncodeColumnConstraint, this=self._parse_var()),
+        "EXCLUDE": lambda self: self.expression(
+            exp.ExcludeColumnConstraint, this=self._parse_index_params()
+        ),
         "FOREIGN KEY": lambda self: self._parse_foreign_key(),
         "FORMAT": lambda self: self.expression(
             exp.DateFormatColumnConstraint, this=self._parse_var_or_string()
@@ -877,7 +880,15 @@ class Parser(metaclass=_Parser):
         "RENAME": lambda self: self._parse_alter_table_rename(),
     }
 
-    SCHEMA_UNNAMED_CONSTRAINTS = {"CHECK", "FOREIGN KEY", "LIKE", "PRIMARY KEY", "UNIQUE", "PERIOD"}
+    SCHEMA_UNNAMED_CONSTRAINTS = {
+        "CHECK",
+        "EXCLUDE",
+        "FOREIGN KEY",
+        "LIKE",
+        "PERIOD",
+        "PRIMARY KEY",
+        "UNIQUE",
+    }
 
     NO_PAREN_FUNCTION_PARSERS = {
         "ANY": lambda self: self.expression(exp.Any, this=self._parse_bitwise()),
@@ -1008,7 +1019,8 @@ class Parser(metaclass=_Parser):
     CLONE_KEYWORDS = {"CLONE", "COPY"}
     HISTORICAL_DATA_KIND = {"TIMESTAMP", "OFFSET", "STATEMENT", "STREAM"}
 
-    OPCLASS_FOLLOW_KEYWORDS = {"ASC", "DESC", "NULLS"}
+    OPCLASS_FOLLOW_KEYWORDS = {"ASC", "DESC", "NULLS", "WITH"}
+
     OPTYPE_FOLLOW_TOKENS = {TokenType.COMMA, TokenType.R_PAREN}
 
     TABLE_INDEX_HINT_TOKENS = {TokenType.FORCE, TokenType.IGNORE, TokenType.USE}
@@ -2841,6 +2853,7 @@ def _parse_join(
 
     def _parse_opclass(self) -> t.Optional[exp.Expression]:
         this = self._parse_conjunction()
+
         if self._match_texts(self.OPCLASS_FOLLOW_KEYWORDS, advance=False):
             return this
 
@@ -2849,6 +2862,37 @@ def _parse_opclass(self) -> t.Optional[exp.Expression]:
 
         return this
 
+    def _parse_index_params(self) -> exp.IndexParameters:
+        using = self._parse_var(any_token=True) if self._match(TokenType.USING) else None
+
+        if self._match(TokenType.L_PAREN, advance=False):
+            columns = self._parse_wrapped_csv(self._parse_with_operator)
+        else:
+            columns = None
+
+        include = self._parse_wrapped_id_vars() if self._match_text_seq("INCLUDE") else None
+        partition_by = self._parse_partition_by()
+        with_storage = (
+            self._parse_csv(self._parse_conjunction) if self._match(TokenType.WITH) else None
+        )
+        tablespace = (
+            self._parse_var(any_token=True)
+            if self._match_text_seq("USING", "INDEX", "TABLESPACE")
+            else None
+        )
+        where = self._parse_where()
+
+        return self.expression(
+            exp.IndexParameters,
+            using=using,
+            columns=columns,
+            include=include,
+            partition_by=partition_by,
+            where=where,
+            with_storage=with_storage,
+            tablespace=tablespace,
+        )
+
     def _parse_index(
         self,
         index: t.Optional[exp.Expression] = None,
@@ -2872,27 +2916,16 @@ def _parse_index(
             index = self._parse_id_var()
             table = None
 
-        using = self._parse_var(any_token=True) if self._match(TokenType.USING) else None
-
-        if self._match(TokenType.L_PAREN, advance=False):
-            columns = self._parse_wrapped_csv(lambda: self._parse_ordered(self._parse_opclass))
-        else:
-            columns = None
-
-        include = self._parse_wrapped_id_vars() if self._match_text_seq("INCLUDE") else None
+        params = self._parse_index_params()
 
         return self.expression(
             exp.Index,
             this=index,
             table=table,
-            using=using,
-            columns=columns,
             unique=unique,
             primary=primary,
             amp=amp,
-            include=include,
-            partition_by=self._parse_partition_by(),
-            where=self._parse_where(),
+            params=params,
         )
 
     def _parse_table_hints(self) -> t.Optional[t.List[exp.Expression]]:
@@ -6094,3 +6127,13 @@ def _parse_truncate_table(self) -> t.Optional[exp.TruncateTable] | exp.Expressio
             option=option,
             partition=partition,
         )
+
+    def _parse_with_operator(self) -> t.Optional[exp.Expression]:
+        this = self._parse_ordered(self._parse_opclass)
+
+        if not self._match(TokenType.WITH):
+            return this
+
+        op = self._parse_var(any_token=True)
+
+        return self.expression(exp.WithOperator, this=this, op=op)


Test patch:
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -130,6 +130,18 @@ def test_postgres(self):
         self.validate_identity(
             "SELECT * FROM foo, LATERAL (SELECT * FROM bar WHERE bar.id = foo.bar_id) AS ss"
         )
+        self.validate_identity(
+            "CREATE TABLE t (vid INT NOT NULL, CONSTRAINT ht_vid_nid_fid_idx EXCLUDE (INT4RANGE(vid, nid) WITH &&, INT4RANGE(fid, fid, '[]') WITH &&))"
+        )
+        self.validate_identity(
+            "CREATE TABLE t (i INT, PRIMARY KEY (i), EXCLUDE USING gist(col varchar_pattern_ops DESC NULLS LAST WITH &&) WITH (sp1 = 1, sp2 = 2))"
+        )
+        self.validate_identity(
+            "CREATE TABLE t (i INT, EXCLUDE USING btree(INT4RANGE(vid, nid, '[]') ASC NULLS FIRST WITH &&) INCLUDE (col1, col2))"
+        )
+        self.validate_identity(
+            "CREATE TABLE t (i INT, EXCLUDE USING gin(col1 WITH &&, col2 WITH ||) USING INDEX TABLESPACE tablespace WHERE (id > 5))"
+        )
         self.validate_identity(
             "SELECT c.oid, n.nspname, c.relname "
             "FROM pg_catalog.pg_class AS c "