Feat(postgres): add support for operator classes in CREATE INDEX DDL (#…

…2317)
tobymao · Sep 25, 2023 · f473e88 · f473e88
1 parent cdcc564
commit f473e88
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 3 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1613,6 +1613,11 @@ def output_name(self) -> str:
         return self.name
 
 
+# https://www.postgresql.org/docs/current/indexes-opclass.html
+class Opclass(Expression):
+    arg_types = {"this": True, "expression": True}
+
+
 class Index(Expression):
     arg_types = {
         "this": False,

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -2851,6 +2851,9 @@ def comprehension_sql(self, expression: exp.Comprehension) -> str:
     def columnprefix_sql(self, expression: exp.ColumnPrefix) -> str:
         return f"{self.sql(expression, 'this')}({self.sql(expression, 'expression')})"
 
+    def opclass_sql(self, expression: exp.Opclass) -> str:
+        return f"{self.sql(expression, 'this')} {self.sql(expression, 'expression')}"
+
 
 def cached_generator(
     cache: t.Optional[t.Dict[int, str]] = None

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -850,6 +850,8 @@ class Parser(metaclass=_Parser):
     CLONE_KEYWORDS = {"CLONE", "COPY"}
     CLONE_KINDS = {"TIMESTAMP", "OFFSET", "STATEMENT"}
 
+    OPCLASS_FOLLOW_KEYWORDS = {"ASC", "DESC", "NULLS"}
+
     TABLE_INDEX_HINT_TOKENS = {TokenType.FORCE, TokenType.IGNORE, TokenType.USE}
 
     WINDOW_ALIAS_TOKENS = ID_VAR_TOKENS - {TokenType.ROWS}
@@ -2463,6 +2465,17 @@ def _parse_join(
         comments = [c for token in (method, side, kind) if token for c in token.comments]
         return self.expression(exp.Join, comments=comments, **kwargs)
 
+    def _parse_opclass(self) -> t.Optional[exp.Expression]:
+        this = self._parse_conjunction()
+        if self._match_texts(self.OPCLASS_FOLLOW_KEYWORDS, advance=False):
+            return this
+
+        opclass = self._parse_var(any_token=True)
+        if opclass:
+            return self.expression(exp.Opclass, this=this, expression=opclass)
+
+        return this
+
     def _parse_index(
         self,
         index: t.Optional[exp.Expression] = None,
@@ -2489,7 +2502,7 @@ def _parse_index(
         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_ordered)
+            columns = self._parse_wrapped_csv(lambda: self._parse_ordered(self._parse_opclass))
         else:
             columns = None
 
@@ -2968,8 +2981,8 @@ def _parse_sort(self, exp_class: t.Type[E], token: TokenType) -> t.Optional[E]:
             return None
         return self.expression(exp_class, expressions=self._parse_csv(self._parse_ordered))
 
-    def _parse_ordered(self) -> exp.Ordered:
-        this = self._parse_conjunction()
+    def _parse_ordered(self, parse_method: t.Optional[t.Callable] = None) -> exp.Ordered:
+        this = parse_method() if parse_method else self._parse_conjunction()
 
         asc = self._match(TokenType.ASC)
         desc = self._match(TokenType.DESC) or (asc and False)

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -9,6 +9,9 @@ class TestPostgres(Validator):
     dialect = "postgres"
 
     def test_ddl(self):
+        self.validate_identity(
+            "CREATE INDEX foo ON bar.baz USING btree(col1 varchar_pattern_ops ASC, col2)"
+        )
         self.validate_identity(
             "CREATE TABLE test (x TIMESTAMP WITHOUT TIME ZONE[][])",
             "CREATE TABLE test (x TIMESTAMP[][])",