trying to set up assignment of create operation

capsulflask/db_model.py

@@ -1,4 +1,6 @@
 
+# I was never able to get this type hinting to work correctly 
+# from psycopg2.extensions import connection as Psycopg2Connection, cursor as Psycopg2Cursor
 from nanoid import generate
 from flask import current_app
 from typing import List
@@ -10,6 +12,7 @@ class OnlineHost:
     self.url = url
 
 class DBModel:
+  #def __init__(self, connection: Psycopg2Connection, cursor: Psycopg2Cursor):
   def __init__(self, connection, cursor):
     self.connection = connection
     self.cursor = cursor
@@ -311,6 +314,30 @@ class DBModel:
     else:
       return None
 
+  def host_operation_exists(self, operation_id: int, host_id: str) -> bool:
+    self.cursor.execute("SELECT operation FROM host_operation WHERE host = %s AND operation = %s",(host_id, operation_id))
+    return len(self.cursor.fetchall()) != 0
+
+  def claim_operation(self, operation_id: int, host_id: str) -> bool:
+    self.cursor.execute("""
+      SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;
+      BEGIN TRANSACTION;
+      UPDATE host_operation SET assignment_status = 'assigned' 
+        WHERE host = %s AND operation = %s AND operation != (
+          SELECT COALESCE(
+            (SELECT operation FROM host_operation WHERE operation = %s AND assignment_status = 'assigned'),
+            -1
+          ) as already_assigned_operation_id
+        );
+      COMMIT TRANSACTION;
+    """,  (host_id, operation_id, operation_id))
+
+    to_return = self.cursor.rowcount != 0
+
+    self.connection.commit()
+
+    return to_return
+