From 0d61c94250dba67e689b34d1c530b62ce234e7fb Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Fri, 31 Mar 2023 23:09:46 -0400
Subject: [PATCH] add db_lock function and locks table

---
 docs/schema.sql |  6 ++++++
 kojihub/db.py   | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/docs/schema.sql b/docs/schema.sql
index 505f0b8d..da3f3efb 100644
--- a/docs/schema.sql
+++ b/docs/schema.sql
@@ -983,4 +983,10 @@ CREATE TABLE rpm_checksum (
 ) WITHOUT OIDS;
 CREATE INDEX rpm_checksum_rpm_id ON rpm_checksum(rpm_id);
 
+
+CREATE TABLE locks (
+        name TEXT NOT NULL PRIMARY KEY
+) WITHOUT OIDS;
+
+
 COMMIT WORK;
diff --git a/kojihub/db.py b/kojihub/db.py
index 478c30fa..d331bee8 100644
--- a/kojihub/db.py
+++ b/kojihub/db.py
@@ -321,6 +321,41 @@ def currval(sequence):
     return _singleValue("SELECT currval(%(sequence)s)", data, strict=True)
 
 
+def db_lock(name, wait=True):
+    """Obtain lock for name
+
+    :param string name: the lock name
+    :param bool wait: whether to wait for the lock (default: True)
+    :return: True if locked, False otherwise
+
+    This function is implemented using db row locks and the locks table
+    """
+    # first see if we need to add the row
+    query = "SELECT name FROM locks WHERE name=%(name)s"
+    data = {"name": name}
+    rows =_fetchMulti(query, data)
+    if not rows:
+        insert = "INSERT INTO locks (name) VALUES (%(name)s) ON CONFLICT DO NOTHING"
+        _dml(insert, data)
+
+    # and then actually lock the row
+    if wait:
+        query = "SELECT name FROM locks WHERE name=%(name)s FOR UPDATE"
+    else:
+        # using SKIP LOCKED rather than NOWAIT to avoid error messages
+        query = "SELECT name FROM locks WHERE name=%(name)s FOR UPDATE SKIP LOCKED"
+    rows = _fetchMulti(query, data)
+
+    if rows:
+        # we have the lock
+        return True
+    elif wait:
+        # should not happen
+        raise koji.LockError(f"Failed to read lock {name}")
+    else:
+        return False
+
+
 class Savepoint(object):
 
     def __init__(self, name):