From 71a2515b273c2f2d02258772c51f7e955f54b420 Mon Sep 17 00:00:00 2001
From: Milo Casagrande <milo.casagrande@linaro.org>
Date: Wed, 26 Nov 2014 09:41:23 +0100
Subject: Fix boot import and saving.

Change-Id: I9f6e05a37bb4c306ba852c0cbfbbad90cc15d963
---
 app/handlers/boot.py    | 33 +++++++++++++++++----------------
 app/utils/bootimport.py | 28 ++++++++++++++++++++--------
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/app/handlers/boot.py b/app/handlers/boot.py
index 01a75aa..cd76e92 100644
--- a/app/handlers/boot.py
+++ b/app/handlers/boot.py
@@ -88,23 +88,24 @@ class BootHandler(hbase.BaseHandler):
             self.db[models.LAB_COLLECTION], [lab_name], field=models.NAME_KEY
         )
 
-        lab_token_doc = utils.db.find_one(
-            self.db[models.TOKEN_COLLECTION], [lab_doc[models.TOKEN_KEY]]
-        )
+        if lab_doc:
+            lab_token_doc = utils.db.find_one(
+                self.db[models.TOKEN_COLLECTION], [lab_doc[models.TOKEN_KEY]]
+            )
 
-        if all([lab_doc, lab_token_doc]):
-            lab_token = mtoken.Token.from_json(lab_token_doc)
-            if all([req_token == lab_token.token, not lab_token.expired]):
-                valid_lab = True
-            elif all([lab_token.is_admin, not lab_token.expired]):
-                valid_lab = True
-                utils.LOG.warn(
-                    "Received boot POST request from an admin token")
-            else:
-                utils.LOG.warn(
-                    "Received token (%s) is not associated with lab '%s'",
-                    req_token, lab_name
-                )
+            if lab_token_doc:
+                lab_token = mtoken.Token.from_json(lab_token_doc)
+                if all([req_token == lab_token.token, not lab_token.expired]):
+                    valid_lab = True
+                elif all([lab_token.is_admin, not lab_token.expired]):
+                    valid_lab = True
+                    utils.LOG.warn(
+                        "Received boot POST request from an admin token")
+                else:
+                    utils.LOG.warn(
+                        "Received token (%s) is not associated with lab '%s'",
+                        req_token, lab_name
+                    )
 
         return valid_lab
 
diff --git a/app/utils/bootimport.py b/app/utils/bootimport.py
index a8dc3fe..62ba468 100644
--- a/app/utils/bootimport.py
+++ b/app/utils/bootimport.py
@@ -80,17 +80,29 @@ def save_or_update(boot_doc, database):
     :param database: The database connection.
     :return The save action return code and the doc ID.
     """
-    prev_doc = utils.db.find_one(
-        database[models.BOOT_COLLECTION],
-        [boot_doc.name],
-        field=models.NAME_KEY,
-        fields=[models.ID_KEY, models.CREATED_KEY]
-    )
+    spec = {
+        models.LAB_NAME_KEY: boot_doc.lab_name,
+        models.NAME_KEY: boot_doc.name,
+    }
+
+    fields = [
+        models.CREATED_KEY,
+        models.ID_KEY,
+    ]
+
+    found_doc = utils.db.find(
+        database[models.BOOT_COLLECTION], 1, 0, spec=spec, fields=fields)
+
+    prev_doc = None
+    doc_len = found_doc.count()
+    if doc_len == 1:
+        prev_doc = found_doc[0]
 
     if prev_doc:
-        doc_id = prev_doc.id
+        doc_get = prev_doc.get
+        doc_id = doc_get(models.ID_KEY)
         boot_doc.id = doc_id
-        boot_doc.created_on = prev_doc.created_on
+        boot_doc.created_on = doc_get(models.CREATED_KEY)
 
         utils.LOG.info("Updating boot document with id '%s'", doc_id)
         ret_val, _ = utils.db.save(database, boot_doc)
-- 
cgit v1.2.3