Skip to content

Commit

Permalink
Added ability to define custom Funcs as constraints to colocate the f…
Browse files Browse the repository at this point in the history
…uction's definition
  • Loading branch information
David Sanders committed Jul 12, 2024
1 parent 1dc5d70 commit a9f24ed
Show file tree
Hide file tree
Showing 9 changed files with 175 additions and 0 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,3 +27,4 @@ Various tricks with Django - some silly, some quite useful.
20. [CloneDbTestCase](./clone_db_testcase)
21. [Additional Code Formatters](./isort_migrations)
22. [Bulk Create Form](./bulk_create_form)
23. [Custom Funcs as Constraints](./func_as_constraint)
73 changes: 73 additions & 0 deletions func_as_constraint/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
Functions as Constraints
========================

July 2024


We've seen that we can use the constraints API to add artifacts to the database. An interesting application
is to colocate the definition for custom functions defined as `Func`:


```python
class HelloWorld(Func):
function = "hello_world"
output_field = models.CharField()

create_function = """\
CREATE OR REPLACE FUNCTION hello_world()
RETURNS varchar
AS $$
BEGIN
RETURN 'Hello World!';
END;
$$ LANGUAGE plpgsql;
"""


class PlaceholderModel(models.Model):
class Meta:
constraints = [
func_as_constraint(HelloWorld),
]
```

Using a simple function we can create a `RawSQL` constraint defined in [../abusing_constraints](Having Fun with Constraints)

```python
def func_as_constraint(func_class):
sql = textwrap.dedent(func_class.create_function).strip()
return RawSQL(
name=func_class.function,
sql=sql,
...
)
```

A next step might be to create a reverse for RawSQL that drops the function:

```python
def func_as_constraint(func_class):
sql = textwrap.dedent(func_class.create_function).strip()
return RawSQL(
name=func_class.function,
sql=sql,
reverse_sql=f"DROP FUNCTION IF EXISTS {func_class.function}",
)
```

We can even go so far as to add a comment to the function if the function class has a docstring:

```python
def func_as_constraint(func_class):
sql = textwrap.dedent(func_class.create_function).strip()
if func_class.__doc__:
comment = psycopg_any_sql.quote(textwrap.dedent(func_class.__doc__).strip())
if sql[-1] != ";":
sql += ";"
sql += f"COMMENT ON FUNCTION {func_class.function} IS {comment};"
return RawSQL(
name=func_class.function,
sql=sql,
reverse_sql=f"DROP FUNCTION IF EXISTS {func_class.function}",
)
```
Empty file added func_as_constraint/__init__.py
Empty file.
6 changes: 6 additions & 0 deletions func_as_constraint/apps.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
from django.apps import AppConfig


class FuncAsConstraintConfig(AppConfig):
default_auto_field = "django.db.models.BigAutoField"
name = "func_as_constraint"
35 changes: 35 additions & 0 deletions func_as_constraint/migrations/0001_initial.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
from django.db import migrations
from django.db import models

import abusing_constraints.constraints


class Migration(migrations.Migration):
initial = True

dependencies = []

operations = [
migrations.CreateModel(
name="PlaceholderModel",
fields=[
(
"id",
models.BigAutoField(
auto_created=True,
primary_key=True,
serialize=False,
verbose_name="ID",
),
),
],
),
migrations.AddConstraint(
model_name="placeholdermodel",
constraint=abusing_constraints.constraints.RawSQL(
name="hello_world",
reverse_sql="DROP FUNCTION IF EXISTS hello_world",
sql="CREATE OR REPLACE FUNCTION hello_world()\nRETURNS varchar\nAS $$\nBEGIN\n RETURN 'Hello World!';\nEND;\n$$ LANGUAGE plpgsql;COMMENT ON FUNCTION hello_world IS 'Print hello world!';",
),
),
]
Empty file.
46 changes: 46 additions & 0 deletions func_as_constraint/models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import textwrap

from django.db import models
from django.db.backends.postgresql.psycopg_any import sql as psycopg_any_sql
from django.db.models.expressions import Func

from abusing_constraints.constraints import RawSQL


def func_as_constraint(func_class):
sql = textwrap.dedent(func_class.create_function).strip()
if func_class.__doc__:
comment = psycopg_any_sql.quote(textwrap.dedent(func_class.__doc__).strip())
if sql[-1] != ";":
sql += ";"
sql += f"COMMENT ON FUNCTION {func_class.function} IS {comment};"
return RawSQL(
name=func_class.function,
sql=sql,
reverse_sql=f"DROP FUNCTION IF EXISTS {func_class.function}",
)


class HelloWorld(Func):
"""
Print hello world!
"""

function = "hello_world"
output_field = models.CharField()
create_function = """\
CREATE OR REPLACE FUNCTION hello_world()
RETURNS varchar
AS $$
BEGIN
RETURN 'Hello World!';
END;
$$ LANGUAGE plpgsql;
"""


class PlaceholderModel(models.Model):
class Meta:
constraints = [
func_as_constraint(HelloWorld),
]
13 changes: 13 additions & 0 deletions func_as_constraint/tests.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import pytest

from .models import HelloWorld, PlaceholderModel

pytestmark = pytest.mark.django_db


def test_hello_world():
PlaceholderModel.objects.create()

qs = PlaceholderModel.objects.annotate(hello_world=HelloWorld())

assert qs[0].hello_world == "Hello World!"
1 change: 1 addition & 0 deletions stupid_django_tricks/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
# Application definition

INSTALLED_APPS = [
"func_as_constraint",
"clone_db_testcase",
"negative_indexing_querysets",
"xor_function",
Expand Down

0 comments on commit a9f24ed

Please sign in to comment.