Add support for explicit table-level locks

Author: Photonios

docs/source/api_reference.rst

@@ -37,6 +37,9 @@ API Reference
    .. autoclass:: ConditionalUniqueIndex
    .. autoclass:: CaseInsensitiveUniqueIndex
 
+.. automodule:: psqlextra.locking
+   :members:
+
 .. automodule:: psqlextra.partitioning
    :members:
 

docs/source/index.rst

@@ -35,6 +35,10 @@ Explore the documentation to learn about all features:
 
    Support for ``TRUNCATE TABLE`` statements (including cascading).
 
+* :ref:`Locking models & tables <locking_page>`
+
+   Support for explicit table-level locks.
+
 
 .. toctree::
    :maxdepth: 2
@@ -49,6 +53,7 @@ Explore the documentation to learn about all features:
    table_partitioning
    expressions
    annotations
+   locking
    settings
    api_reference
    major_releases

docs/source/locking.rst

@@ -0,0 +1,56 @@
+.. include:: ./snippets/postgres_doc_links.rst
+
+.. _locking_page:
+
+Locking
+=======
+
+`Explicit table-level locks`_ are supported through the :meth:`psqlextra.locking.postgres_lock_model` and :meth:`psqlextra.locking.postgres_lock_table` methods. All table-level lock methods are supported.
+
+Locks are always bound to the current transaction and are released when the transaction is comitted or rolled back. There is no support (in PostgreSQL) for explicitly releasing a lock.
+
+.. warning::
+
+    Locks are only released when the *outer* transaction commits or when a nested transaction is rolled back. You can ensure that the transaction you created is the outermost one by passing the ``durable=True`` argument to ``transaction.atomic``.
+
+.. note::
+
+    Use `django-pglocks <https://pypi.org/project/django-pglocks/>`_ if you need a advisory lock.
+
+Locking a model
+---------------
+
+Use :class:`psqlextra.locking.PostgresTableLockMode` to indicate the type of lock to acquire.
+
+.. code-block:: python
+
+    from django.db import transaction
+
+    from psqlextra.locking import PostgresTableLockMode, postgres_lock_table
+
+    with transaction.atomic(durable=True):
+        postgres_lock_model(MyModel, PostgresTableLockMode.EXCLUSIVE)
+
+    # locks are released here, when the transaction comitted
+
+
+Locking a table
+---------------
+
+Use :meth:`psqlextra.locking.postgres_lock_table` to lock arbitrary tables in arbitrary schemas.
+
+.. code-block:: python
+
+    from django.db import transaction
+
+    from psqlextra.locking import PostgresTableLockMode, postgres_lock_table
+
+    with transaction.atomic(durable=True):
+        postgres_lock_table("mytable", PostgresTableLockMode.EXCLUSIVE)
+        postgres_lock_table(
+            "tableinotherschema",
+            PostgresTableLockMode.EXCLUSIVE,
+            schema_name="myschema"
+        )
+
+    # locks are released here, when the transaction comitted

docs/source/snippets/postgres_doc_links.rst

@@ -2,3 +2,4 @@
 .. _TRUNCATE TABLE: https://www.postgresql.org/docs/9.1/sql-truncate.html
 .. _hstore: https://www.postgresql.org/docs/11/hstore.html
 .. _PostgreSQL Declarative Table Partitioning: https://www.postgresql.org/docs/current/ddl-partitioning.html#DDL-PARTITIONING-DECLARATIVE
+.. _Explicit table-level locks: https://www.postgresql.org/docs/current/explicit-locking.html#LOCKING-TABLES

psqlextra/backend/introspection.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from psqlextra.types import PostgresPartitioningMethod
 
@@ -198,3 +198,19 @@ def get_constraints(self, cursor, table_name: str):
                 constraint["definition"] = definition
 
         return constraints
+
+    def get_table_locks(self, cursor) -> List[Tuple[str, str, str]]:
+        cursor.execute(
+            """
+         SELECT
+          n.nspname,
+          t.relname,
+          l.mode
+        FROM pg_locks l
+        INNER JOIN pg_class t ON t.oid = l.relation
+        INNER JOIN pg_namespace n ON n.oid = t.relnamespace
+        WHERE t.relnamespace >= 2200
+        ORDER BY n.nspname, t.relname, l.mode"""
+        )
+
+        return cursor.fetchall()

psqlextra/locking.py

@@ -0,0 +1,97 @@
+from enum import Enum
+from typing import Optional, Type
+
+from django.db import DEFAULT_DB_ALIAS, connections, models
+
+
+class PostgresTableLockMode(Enum):
+    """List of table locking modes.
+
+    See: https://www.postgresql.org/docs/current/explicit-locking.html
+    """
+
+    ACCESS_SHARE = "ACCESS SHARE"
+    ROW_SHARE = "ROW SHARE"
+    ROW_EXCLUSIVE = "ROW EXCLUSIVE"
+    SHARE_UPDATE_EXCLUSIVE = "SHARE UPDATE EXCLUSIVE"
+    SHARE = "SHARE"
+    SHARE_ROW_EXCLUSIVE = "SHARE ROW EXCLUSIVE"
+    EXCLUSIVE = "EXCLUSIVE"
+    ACCESS_EXCLUSIVE = "ACCESS EXCLUSIVE"
+
+
+def postgres_lock_table(
+    table_name: str,
+    lock_mode: PostgresTableLockMode,
+    *,
+    schema_name: Optional[str] = None,
+    using: str = DEFAULT_DB_ALIAS,
+) -> None:
+    """Locks the specified table with the specified mode.
+
+    The lock is held until the end of the current transaction.
+
+    Arguments:
+        table_name:
+            Unquoted table name to acquire the lock on.
+
+        lock_mode:
+            Type of lock to acquire.
+
+        schema_name:
+            Optionally, the unquoted name of the schema
+            the table to lock is in. If not specified,
+            the table name is resolved by PostgreSQL
+            using it's ``search_path``.
+
+        using:
+            Name of the database alias to use.
+    """
+
+    connection = connections[using]
+
+    with connection.cursor() as cursor:
+        quoted_fqn = connection.ops.quote_name(table_name)
+        if schema_name:
+            quoted_fqn = (
+                connection.ops.quote_name(schema_name) + "." + quoted_fqn
+            )
+
+        cursor.execute(f"LOCK TABLE {quoted_fqn} IN {lock_mode.value} MODE")
+
+
+def postgres_lock_model(
+    model: Type[models.Model],
+    lock_mode: PostgresTableLockMode,
+    *,
+    using: str = DEFAULT_DB_ALIAS,
+    schema_name: Optional[str] = None,
+) -> None:
+    """Locks the specified model with the specified mode.
+
+    The lock is held until the end of the current transaction.
+
+    Arguments:
+        model:
+            The model of which to lock the table.
+
+        lock_mode:
+            Type of lock to acquire.
+
+        schema_name:
+            Optionally, the unquoted name of the schema
+            the table to lock is in. If not specified,
+            the table name is resolved by PostgreSQL
+            using it's ``search_path``.
+
+            Django models always reside in the default
+            ("public") schema. You should not specify
+            this unless you're doing something special.
+
+        using:
+            Name of the database alias to use.
+    """
+
+    postgres_lock_table(
+        model._meta.db_table, lock_mode, schema_name=schema_name, using=using
+    )

tests/test_locking.py

@@ -0,0 +1,107 @@
+import uuid
+
+import pytest
+
+from django.db import connection, models, transaction
+
+from psqlextra.locking import (
+    PostgresTableLockMode,
+    postgres_lock_model,
+    postgres_lock_table,
+)
+
+from .fake_model import get_fake_model
+
+
+@pytest.fixture
+def mocked_model():
+    return get_fake_model(
+        {
+            "name": models.TextField(),
+        }
+    )
+
+
+def get_table_locks():
+    with connection.cursor() as cursor:
+        return connection.introspection.get_table_locks(cursor)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_postgres_lock_table(mocked_model):
+    lock_signature = (
+        "public",
+        mocked_model._meta.db_table,
+        "AccessExclusiveLock",
+    )
+    with transaction.atomic():
+        postgres_lock_table(
+            mocked_model._meta.db_table, PostgresTableLockMode.ACCESS_EXCLUSIVE
+        )
+        assert lock_signature in get_table_locks()
+
+    assert lock_signature not in get_table_locks()
+
+
+@pytest.mark.django_db(transaction=True)
+def test_postgres_lock_table_in_schema():
+    schema_name = str(uuid.uuid4())[:8]
+    table_name = str(uuid.uuid4())[:8]
+    quoted_schema_name = connection.ops.quote_name(schema_name)
+    quoted_table_name = connection.ops.quote_name(table_name)
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"CREATE SCHEMA {quoted_schema_name}")
+        cursor.execute(
+            f"CREATE TABLE {quoted_schema_name}.{quoted_table_name} AS SELECT 'hello world'"
+        )
+
+    lock_signature = (schema_name, table_name, "ExclusiveLock")
+    with transaction.atomic():
+        postgres_lock_table(
+            table_name, PostgresTableLockMode.EXCLUSIVE, schema_name=schema_name
+        )
+        assert lock_signature in get_table_locks()
+
+    assert lock_signature not in get_table_locks()
+
+
+@pytest.mark.django_db(transaction=True)
+def test_postgres_lock_mode(mocked_model):
+    lock_signature = (
+        "public",
+        mocked_model._meta.db_table,
+        "AccessExclusiveLock",
+    )
+
+    with transaction.atomic():
+        postgres_lock_model(
+            mocked_model, PostgresTableLockMode.ACCESS_EXCLUSIVE
+        )
+        assert lock_signature in get_table_locks()
+
+    assert lock_signature not in get_table_locks()
+
+
+@pytest.mark.django_db(transaction=True)
+def test_postgres_lock_model_in_schema(mocked_model):
+    schema_name = str(uuid.uuid4())[:8]
+    quoted_schema_name = connection.ops.quote_name(schema_name)
+    quoted_table_name = connection.ops.quote_name(mocked_model._meta.db_table)
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"CREATE SCHEMA {quoted_schema_name}")
+        cursor.execute(
+            f"CREATE TABLE {quoted_schema_name}.{quoted_table_name} (LIKE public.{quoted_table_name} INCLUDING ALL)"
+        )
+
+    lock_signature = (schema_name, mocked_model._meta.db_table, "ExclusiveLock")
+    with transaction.atomic():
+        postgres_lock_model(
+            mocked_model,
+            PostgresTableLockMode.EXCLUSIVE,
+            schema_name=schema_name,
+        )
+        assert lock_signature in get_table_locks()
+
+    assert lock_signature not in get_table_locks()