Search code examples
djangopostgresqldatabase-migrationcheck-constraints

How to rewrite this postgres constraint to be more in line with django 2.2?


THis is the raw query that I write for postgres for the check constraint

   ALTER TABLE rea_asplinkage ADD CONSTRAINT asp_sub_project_positive_integer
        CHECK (
            jsonb_typeof(linkage-> 'root' -> 'in_sub_project') is not distinct from 'number'
        and (linkage->'root'->>'in_sub_project')::numeric % 1 = 0
        and (linkage->'root'->>'in_sub_project')::numeric > 0
        );

And the way I create the migration is this way

# Generated by Django 2.2.10 on 2020-05-16 12:59

from django.db import connection, migrations



class Migration(migrations.Migration):

    dependencies = [("rea", "0029_asplinkage")]

    operations = [
        migrations.RunSQL(
            sql="""
            ALTER TABLE rea_asplinkage ADD CONSTRAINT asp_sub_project_positive_integer
            CHECK (
                jsonb_typeof(linkage-> 'root' -> 'in_sub_project') is not distinct from 'number'
            and (linkage->'root'->>'in_sub_project')::numeric % 1 = 0
            and (linkage->'root'->>'in_sub_project')::numeric > 0
            );
            """,
            reverse_sql="""
                ALTER TABLE rea_asplinkage DROP CONSTRAINT "asp_sub_project_positive_integer";
            """,
        )
    ]

And this works.

But this means that my original model does not show the constraint in the class Meta of the ASPLinkage model

class ASPLinkage(TimeStampedModel, SoftDeletableModel, PersonStampedModel, OrganizationOwnedModel):
    linkage = JSONField(default=default_linkage_for_asp)

    objects = OrganizationOwnedSoftDeletableManager()

I have tried ExpressionWrapper and RawSQL in creating the constraints inside the class Meta, but it still doesn't work.

For reference, I have looked at the examples found in https://github.com/django/django/blob/master/tests/constraints/models.py#L12

I have also looked at Separate Database and State migration via https://realpython.com/create-django-index-without-downtime/#when-django-generates-a-new-migration

But I still cannot get it to work

So is this even possible?

Update

Let me write a summary of my question for better readability.

  1. I want to write constraints on a JSONField.
  2. I can do that directly on the Postgres
  3. Therefore I can do it using raw sql in the migration file
  4. But I cannot do the equivalent using Django model meta / CheckConstraint which usually is how anyone does it. See https://docs.djangoproject.com/en/3.0/ref/models/constraints/
  5. So how do I rewrite this raw sql to produce constraint on a jsonfield in postgres but in the Django way?

Solution

  • In order to achieve that on Django 2.2 you'll need to register two new JSONField transforms/lookups since support for conditional expressions was only added in the upcoming 3.1 release.

    You'll first want to register lookups for JSONField key accesses

    from django.db import models
    from django.db.models.lookups import Lookup
    from django.contrib.postgres.fields.jsonb import (
        KeyTransform, KeyTransformTextLookupMixin
    )
    
    @KeyTransform.register_lookup
    class KeyTransformIsInteger(KeyTransformTextLookupMixin, Lookup):
        lookup_name = 'is_int'
        prepare_rhs = False
    
        def as_sql(self, compiler, connection):
            key_expr = KeyTransform(
                self.lhs.key_name, *self.lhs.source_expressions, **self.lhs.extra
            )
            key_sql, key_params = self.process_lhs(
                compiler, connection, lhs=key_expr
            )
            lhs_sql, lhs_params = self.process_lhs(compiler, connection)
            rhs_sql, rhs_params = self.process_rhs(compiler, connection)
            sql = "(jsonb_typeof(%s) = %%s AND mod(%s::numeric, %%s) = %%s) IS %s" % (
                key_sql, lhs_sql, rhs_sql
            )
            params = [
                *key_params, 'number',
                *lhs_params, 1, 0,
                *rhs_params,
            ]
            return sql, params
    
    @KeyTransform.register_lookup
    class KeyTransformIntegerGt(KeyTransformTextLookupMixin, Lookup):
        lookup_name = 'int_gt'
        prepare_rhs = False
    
        def as_sql(self, compiler, connection):
            lhs_sql, lhs_params = self.process_lhs(compiler, connection)
            rhs_sql, rhs_params = self.process_rhs(compiler, connection)
            sql = "%s::int > %s" % (lhs_sql, rhs_sql)
            params = [*lhs_params, *rhs_params]
            return sql, params
    

    Once this is done you should be able to define your constraint like

    CheckConstraint(
        check=Q(
            linkage__root__in_sub_project__is_int=True,
            linkage__root__in_sub_project__int_gt=0,
        ),
        name='asp_sub_project_positive_integer',
    )
    

    Once you're on Django 3.1 you should be able to pass a RawSQL directly to CheckConstraint.check as long as it has an output_field = models.BooleanField().

    RawSQL("""
       jsonb_typeof(linkage-> 'root' -> 'in_sub_project') is not distinct from 'number'
       and (linkage->'root'->>'in_sub_project')::numeric % 1 = 0
       and (linkage->'root'->>'in_sub_project')::numeric > 0
    """,
        output_field=models.BooleanField()
    )