Enforce distinction between creating / updating expense reports.

tseaver · Aug 28, 2014 · Oct 3, 2014 · Oct 3, 2014 · Oct 9, 2014 · Oct 9, 2014
commit d9c3af11b5e350f786fa14f8a786ef27c65670ea
diff --git a/docs/examples/expenses/expenses/__init__.py b/docs/examples/expenses/expenses/__init__.py
@@ -6,34 +6,42 @@
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.query import Query
 
+
+class DuplicateReport(Exception):
+    """Attempt to create a report which already exists."""
+
+
+class InvalidReport(Exception):
+    """Attempt to update a report which does not already exist."""
+
+
 def get_dataset():
     client_email = os.environ['GCLOUD_TESTS_CLIENT_EMAIL']
     private_key_path = os.environ['GCLOUD_TESTS_KEY_FILE']
     dataset_id = os.environ['GCLOUD_TESTS_DATASET_ID']
     conn = datastore.get_connection(client_email, private_key_path)
     return conn.dataset(dataset_id)
 
-def get_employee(dataset, employee_id):
+def get_employee(dataset, employee_id, create=True):
     key = Key(dataset, path=[{'kind': 'Employee', 'name': employee_id}])
     employee = dataset.get_entity(key)
-    if employee is None:
+    if employee is None and create:
         employee = dataset.entity('Employee').key(key)
         employee.save()
     return employee
 
-def get_report(dataset, employee_id, report_id):
+def get_report(dataset, employee_id, report_id, create=True):
     key = Key(dataset,
               path=[{'kind': 'Employee', 'name': employee_id},
                     {'kind': 'Expense Report', 'name': report_id},
                    ])
     report = dataset.get_entity(key)
-    if report is None:
+    if report is None and create:
         report = dataset.entity('Report').key(key)
         report.save()
     return report
 
-def upsert_report(employee_id, report_id, rows):
-    dataset = get_dataset()
+def _upsert_report(dataset, employee_id, report_id, rows):
     with dataset.transaction():
         employee = get_employee(dataset, employee_id)
         report = get_report(dataset, employee_id, report_id)
@@ -50,3 +58,15 @@ def upsert_report(employee_id, report_id, rows):
             for k, v in row.items():
                 item[k] = v
             item.save()
+
+def create_report(employee_id, report_id, rows):
+    dataset = get_dataset()
+    if get_report(dataset, employee_id, report_id, False) is not None:
+        raise DuplicateReport()
+    _upsert_report(dataset, employee_id, report_id, rows)
+
+def update_report(employee_id, report_id, rows):
+    dataset = get_dataset()
+    if get_report(dataset, employee_id, report_id, False) is None:
+        raise InvalidReport()
+    _upsert_report(dataset, employee_id, report_id, rows)
diff --git a/docs/examples/expenses/expenses/scripts/submit_expenses.py b/docs/examples/expenses/expenses/scripts/submit_expenses.py
@@ -5,7 +5,10 @@
 import sys
 
 
-from .. import upsert_report
+from .. import DuplicateReport
+from .. import InvalidReport
+from .. import create_report
+from .. import update_report
 
 
 class InvalidCommandLine(ValueError):
@@ -77,17 +80,29 @@ class CreateReport(_Command):
     """Create an expense report from a CSV file.
     """
     def __call__(self):
-        upsert_report(self.employee_id, self.report_id, self.rows)
-        self.submitter.blather("Processed %d rows." % len(self.rows))
-        self.submitter.blather("Created, report ID: %s" % self.report_id)
+        try:
+            create_report(self.employee_id, self.report_id, self.rows)
+        except DuplicateReport:
+            self.submitter.blather("Report already exists: %s/%s"
+                                   % (self.employee_id, self.report_id))
+        else:
+            self.submitter.blather("Created, report ID: %s/%s"
+                                   % (self.employee_id, self.report_id))
+            self.submitter.blather("Processed %d rows." % len(self.rows))
 
 
 class UpdateReport(_Command):
     """Update an expense report from a CSV file.
     """
     def __call__(self):
-        self.submitter.blather("Processed %d rows." % len(self.rows))
-        self.submitter.blather("Updated, report ID: %s" % self.report_id)
+        try:
+            update_report(self.employee_id, self.report_id, self.rows)
+        except InvalidReport:
+            self.submitter.blather("No such report: %s/%s"
+                                   % (self.employee_id, self.report_id))
+        else:
+            self.submitter.blather("Updated, report ID: %s" % self.report_id)
+            self.submitter.blather("Processed %d rows." % len(self.rows))
 
 
 _COMMANDS = {