"""PostgreSQL locks shared by assignment model writes and services."""

from django.apps import apps
from django.contrib.auth import get_user_model
from django.db import connections


def lock_user(user_id, using="default"):
    get_user_model().objects.using(using).select_for_update().get(pk=user_id)


def share_companies(company_ids, using="default"):
    # FOR SHARE is compatible with other assignment writers, but conflicts with
    # the FOR UPDATE company lock used by organization lifecycle operations.
    ids = sorted({value for value in company_ids if value is not None}, key=str)
    if not ids:
        return
    connection = connections[using]
    table = connection.ops.quote_name(apps.get_model("organization", "Company")._meta.db_table)
    with connection.cursor() as cursor:
        cursor.execute(f"SELECT id FROM {table} WHERE id = ANY(%s) ORDER BY id FOR SHARE", [ids])
        cursor.fetchall()