fix COLLATE in CREATE INDEX

cockroachdb · Oct 27, 2021 · bb2a448 · bb2a448
1 parent eba0d82
commit bb2a448
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/django_cockroachdb/features.py b/django_cockroachdb/features.py
@@ -192,10 +192,6 @@ def django_test_expected_failures(self):
             # unexpected partial unique index in pg_constraint query:
             # https://github.com/cockroachdb/cockroach/issues/61098
             'introspection.tests.IntrospectionTests.test_get_constraints_unique_indexes_orders',
-            # COLLATE expressions in index elements not supported:
-            # https://github.com/cockroachdb/cockroach/issues/71240
-            'schema.tests.SchemaTests.test_func_index_collate',
-            'schema.tests.SchemaTests.test_func_index_collate_f_ordered',
             # ALTER COLUMN ... SET NOT NULL crashes with "validate check
             # constraint: column "crdb_internal_idx_expr" does not exist":
             # https://github.com/cockroachdb/cockroach/issues/72012

diff --git a/django_cockroachdb/functions.py b/django_cockroachdb/functions.py
@@ -5,8 +5,8 @@
 )
 from django.db.models.expressions import When
 from django.db.models.functions import (
-    ACos, ASin, ATan, ATan2, Cast, Ceil, Coalesce, Cos, Cot, Degrees, Exp,
-    Floor, JSONObject, Ln, Log, Radians, Round, Sin, Sqrt, StrIndex, Tan,
+    ACos, ASin, ATan, ATan2, Cast, Ceil, Coalesce, Collate, Cos, Cot, Degrees,
+    Exp, Floor, JSONObject, Ln, Log, Radians, Round, Sin, Sqrt, StrIndex, Tan,
 )
 
 
@@ -24,6 +24,16 @@ def coalesce(self, compiler, connection, **extra_context):
     return self.as_sql(compiler, connection, **extra_context)
 
 
+def collate(self, compiler, connection, **extra_context):
+    return self.as_sql(
+        compiler, connection,
+        # CockroachDB requires parentheses around the expression in
+        # CREATE INDEX: https://github.com/cockroachdb/cockroach/issues/71240
+        template='(%(expressions)s %(function)s %(collation)s)',
+        **extra_context
+    )
+
+
 def float_cast(self, compiler, connection, **extra_context):
     # Most cockroachdb math functions require float arguments instead of
     # decimal or integer.
@@ -50,6 +60,7 @@ def register_functions():
     for func in math_funcs_needing_float_cast:
         func.as_cockroachdb = float_cast
     Coalesce.as_cockroachdb = coalesce
+    Collate.as_cockroachdb = collate
     JSONObject.as_cockroachdb = JSONObject.as_postgresql
     StrIndex.as_cockroachdb = StrIndex.as_postgresql
     When.as_cockroachdb = when