From 691b293939adf2beb026d6e085d8e7e97881af2d Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Wed, 10 Nov 2021 23:01:55 -0800
Subject: [PATCH] Add argument raise_for_missing.

---
 voussoirkit/worms.py | 41 +++++++++++++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 8 deletions(-)

diff --git a/voussoirkit/worms.py b/voussoirkit/worms.py
index d45833c..8111d7a 100644
--- a/voussoirkit/worms.py
+++ b/voussoirkit/worms.py
@@ -193,20 +193,24 @@ class Database(metaclass=abc.ABCMeta):
             instance = object_class(self, object_row)
             yield instance
 
-    def get_objects_by_id(self, object_class, object_ids):
+    def get_objects_by_id(self, object_class, object_ids, *, raise_for_missing=False):
         '''
         Select many objects by their IDs.
+
         This is better than calling get_object_by_id in a loop because we can
         use a single SQL select to get batches of up to 999 items.
 
-        This method does not raise exceptions for IDs that are not present in
-        the database. You will need to compare the set of returned objects with
-        the set of requested IDs.
-
-        Note: The order of the output is not guaranteed to match the order of
+        Note: The order of the output will most likely not match the order of
         the input. Consider using get_objects_by_sql if that is a necessity.
+
+        raise_for_missing:
+            If any of the requested object ids are not found in the database,
+            we can raise that class's no_such_exception with the set of missing
+            IDs.
         '''
-        ids_needed = list(set(object_ids))
+        object_ids = set(object_ids)
+        ids_needed = list(object_ids)
+        ids_found = set()
 
         while ids_needed:
             # SQLite3 has a limit of 999 ? in a query, so we must batch them.
@@ -218,8 +222,14 @@ class Database(metaclass=abc.ABCMeta):
             query = f'SELECT * FROM {object_class.table} WHERE id IN {qmarks}'
             for object_row in self.select(query, id_batch):
                 instance = object_class(self, db_row=object_row)
+                ids_found.add(instance.id)
                 yield instance
 
+        if raise_for_missing:
+            missing = object_ids.difference(ids_found)
+            if missing:
+                raise object_class.no_such_exception(missing)
+
     def get_objects_by_sql(self, object_class, query, bindings=None):
         '''
         Use an arbitrary SQL query to select objects from the database.
@@ -459,20 +469,28 @@ class DatabaseWithCaching(Database, metaclass=abc.ABCMeta):
             instance = self.get_cached_instance(object_class, object_row)
             yield instance
 
-    def get_objects_by_id(self, object_class, object_ids):
+    def get_objects_by_id(self, object_class, object_ids, *, raise_for_missing=False):
         '''
         Given multiple IDs, this method will find which ones are in the cache
         and which ones need to be selected from the db.
+
         This is better than calling get_object_by_id in a loop because we can
         use a single SQL select to get batches of up to 999 items.
 
         Note: The order of the output will most likely not match the order of
         the input, because we first pull items from the cache before requesting
         the rest from the database.
+
+        raise_for_missing:
+            If any of the requested object ids are not found in the database,
+            we can raise that class's no_such_exception with the set of missing
+            IDs.
         '''
         object_cache = self.caches.get(object_class, None)
 
+        object_ids = set(object_ids)
         ids_needed = set()
+        ids_found = set()
 
         if object_cache is None:
             ids_needed.update(object_ids)
@@ -483,6 +501,7 @@ class DatabaseWithCaching(Database, metaclass=abc.ABCMeta):
                 except KeyError:
                     ids_needed.add(object_id)
                 else:
+                    ids_found.add(object_id)
                     yield instance
 
         if not ids_needed:
@@ -508,8 +527,14 @@ class DatabaseWithCaching(Database, metaclass=abc.ABCMeta):
                 instance = object_class(self, db_row=object_row)
                 if object_cache is not None:
                     object_cache[instance.id] = instance
+                ids_found.add(instance.id)
                 yield instance
 
+        if raise_for_missing:
+            missing = object_ids.difference(ids_found)
+            if missing:
+                raise object_class.no_such_exception(missing)
+
     def get_objects_by_sql(self, object_class, query, bindings=None):
         '''
         Use an arbitrary SQL query to select objects from the database.