Add allocation methods to the datastore

3 files changed, 88 insertions, 3 deletions

diff --git a/guix-build-coordinator/datastore.scm b/guix-build-coordinator/datastore.scm
index 83fe0cc..04ef0d9 100644
--- a/guix-build-coordinator/datastore.scm
+++ b/guix-build-coordinator/datastore.scm
@@ -12,7 +12,9 @@
                datastore-new-agent-password
                datastore-agent-password-exists?
                datastore-list-unprocessed-builds
-               datastore-replace-build-allocation-plan)
+               datastore-replace-build-allocation-plan
+               datastore-allocate-builds-to-agent
+               datastore-list-allocation-plan-builds)
   #:export (database-uri->datastore))
 
 (define (database-uri->datastore database)
diff --git a/guix-build-coordinator/datastore/abstract.scm b/guix-build-coordinator/datastore/abstract.scm
index 53a2c7f..9d4e7b1 100644
--- a/guix-build-coordinator/datastore/abstract.scm
+++ b/guix-build-coordinator/datastore/abstract.scm
@@ -8,7 +8,9 @@
             datastore-find-agent
             datastore-new-agent
             datastore-new-agent-password
-            datastore-agent-password-exists?))
+            datastore-agent-password-exists?
+            datastore-allocate-builds-to-agent
+            datastore-list-allocation-plan-builds))
 
 (define-class <abstract-datastore> ())
 
@@ -22,3 +24,5 @@
 (define-generic datastore-agent-password-exists?)
 (define-generic datastore-agent-list-unprocessed-builds)
 (define-generic datastore-agent-replace-build-allocation-plan)
+(define-generic datastore-allocate-builds-to-agent)
+(define-generic datastore-list-allocation-plan-builds)
diff --git a/guix-build-coordinator/datastore/sqlite.scm b/guix-build-coordinator/datastore/sqlite.scm
index 749ddcc..3ee3f6c 100644
--- a/guix-build-coordinator/datastore/sqlite.scm
+++ b/guix-build-coordinator/datastore/sqlite.scm
@@ -18,7 +18,9 @@
             datastore-new-agent-password
             datastore-agent-password-exists?
             datastore-list-unprocessed-builds
-            datastore-replace-build-allocation-plan))
+            datastore-replace-build-allocation-plan
+            datastore-allocate-builds-to-agent
+            datastore-list-allocation-plan-builds))
 
 (define-class <sqlite-datastore> (<abstract-datastore>)
   database-file
@@ -231,6 +233,83 @@ INSERT INTO build_allocation_plan (build_id, agent_id, ordering) VALUES "
      (sqlite-exec db "COMMIT TRANSACTION;")))
   #t)
 
+(define-method (datastore-allocate-builds-to-agent
+                (datastore <sqlite-datastore>)
+                agent-id
+                build-ids)
+  (define (insert-to-allocated-builds db agent-id build-ids)
+    (sqlite-exec
+     db
+     (string-append
+      "
+INSERT INTO allocated_builds (build_id, agent_id) VALUES "
+      (string-join
+       (map (lambda (build-id)
+              (simple-format
+               #f
+               "('~A', '~A')"
+               build-id
+               agent-id))
+            build-ids)
+       ", ")
+      ";")))
+
+  (define (remove-builds-from-plan db build-ids)
+    (sqlite-exec
+     db
+     (string-append
+      "
+DELETE FROM build_allocation_plan
+WHERE build_id IN ("
+      (string-join
+       (map (lambda (build-id)
+              (string-append "'" build-id "'"))
+            build-ids)
+       ", ")
+      ")")))
+
+  (call-with-worker-thread
+   (slot-ref datastore 'worker-thread-channel)
+   (lambda (db)
+     (sqlite-exec db "BEGIN TRANSACTION;")
+     (insert-to-allocated-builds db agent-id build-ids)
+     (remove-builds-from-plan db build-ids)
+     (sqlite-exec db "COMMIT TRANSACTION;")))
+  #t)
+
+(define-method (datastore-list-allocation-plan-builds
+                (datastore <sqlite-datastore>)
+                agent-id
+                limit)
+  (call-with-worker-thread
+   (slot-ref datastore 'worker-thread-channel)
+   (lambda (db)
+     (let ((statement
+            (sqlite-prepare
+             db
+             "
+SELECT builds.uuid, builds.derivation_name
+FROM builds
+INNER JOIN build_allocation_plan
+  ON builds.uuid = build_allocation_plan.build_id
+WHERE build_allocation_plan.agent_id = :agent_id
+LIMIT :limit")))
+
+       (sqlite-bind-arguments
+        statement
+        #:agent_id agent-id
+        #:limit limit)
+
+       (let ((builds (sqlite-map
+                      (match-lambda
+                        (#(uuid derivation_name)
+                         `((uuid            . ,uuid)
+                           (derivation-name . ,derivation_name))))
+                      statement)))
+         (sqlite-reset statement)
+
+         builds)))))
+
 (define (db-open database)
   (define flags
     (list SQLITE_OPEN_READWRITE