db: store task_id from the scheduler

sql/swh-vault-schema.sql

@@ -22,6 +22,7 @@ create table vault_bundle (
   type cook_type not null,  -- requested cooking type
   object_id obj_hash not null,  -- requested object ID
 
+  task_id integer,  -- scheduler task id
   task_status cook_status not null default 'new',  -- status of the task
   sticky boolean not null default false, -- bundle cannot expire
 

swh/vault/backend.py

@@ -127,7 +127,7 @@ class VaultBackend:
         """Fetch information from a bundle"""
         obj_id = hashutil.hash_to_bytes(obj_id)
         cursor.execute('''
-            SELECT id, type, object_id, task_status, sticky,
+            SELECT id, type, object_id, task_id, task_status, sticky,
                    ts_created, ts_done, ts_last_access, progress_msg
             FROM vault_bundle
             WHERE type = %s AND object_id = %s''', (obj_type, obj_id))
@@ -139,7 +139,8 @@ class VaultBackend:
     def _send_task(self, args):
         """Send a cooking task to the celery scheduler"""
         task = create_oneshot_task_dict('swh-vault-cooking', *args)
-        self.scheduler.create_tasks([task])
+        added_tasks = self.scheduler.create_tasks([task])
+        return added_tasks[0]['id']
 
     @autocommit
     def create_task(self, obj_type, obj_id, sticky=False, cursor=None):
@@ -157,7 +158,12 @@ class VaultBackend:
             VALUES (%s, %s, %s)''', (obj_type, obj_id, sticky))
         self.commit()
 
-        self._send_task(args)
+        task_id = self._send_task(args)
+
+        cursor.execute('''
+            UPDATE vault_bundle
+            SET task_id = %s
+            WHERE type = %s AND object_id = %s''', (task_id, obj_type, obj_id))
 
     @autocommit
     def add_notif_email(self, obj_type, obj_id, email, cursor=None):

swh/vault/tests/test_backend.py

@@ -20,6 +20,7 @@ class BaseTestBackend(VaultTestFixture, StorageTestFixture, DbTestFixture):
     @contextlib.contextmanager
     def mock_cooking(self):
         with patch.object(self.vault_backend, '_send_task') as mt:
+            mt.return_value = 42
             with patch('swh.vault.backend.get_cooker') as mg:
                 mcc = unittest.mock.MagicMock()
                 mc = unittest.mock.MagicMock()
@@ -76,6 +77,7 @@ class TestBackend(BaseTestBackend, unittest.TestCase):
         self.assertEqual(info['object_id'], TEST_OBJ_ID)
         self.assertEqual(info['type'], TEST_TYPE)
         self.assertEqual(info['task_status'], 'new')
+        self.assertEqual(info['task_id'], 42)
 
         self.assertTimestampAlmostNow(info['ts_created'])