Merge pull request #837 from JarriqTheTechie/feature/find_or

Feature/find_or

Author: josephmancuso

src/masoniteorm/exceptions.py

@@ -32,3 +32,7 @@ class InvalidUrlConfiguration(Exception):
 
 class MultipleRecordsFound(Exception):
     pass
+
+
+class InvalidArgument(Exception):
+    pass

src/masoniteorm/models/Model.py

@@ -183,6 +183,7 @@ class Model(TimeStampsMixin, ObservesEvents, metaclass=ModelMeta):
             "doesnt_exist",
             "doesnt_have",
             "exists",
+            "find_or",
             "find_or_404",
             "find_or_fail",
             "first_or_fail",

src/masoniteorm/query/QueryBuilder.py

@@ -1,15 +1,15 @@
 import inspect
 from copy import deepcopy
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Callable
 
 from ..collection.Collection import Collection
 from ..config import load_config
 from ..exceptions import (
     HTTP404,
     ConnectionNotRegistered,
     ModelNotFound,
-    MultipleRecordsFound,
+    MultipleRecordsFound, InvalidArgument,
 )
 from ..expressions.expressions import (
     AggregateExpression,
@@ -1797,6 +1797,30 @@ def find(self, record_id):
 
         return self.where(self._model.get_primary_key(), record_id).first()
 
+    def find_or(self, record_id: int, callback: Callable, args=None):
+        """Finds a row by the primary key ID (Requires a model) or raise a ModelNotFound exception.
+
+        Arguments:
+            record_id {int} -- The ID of the primary key to fetch.
+            callback {Callable} -- The function to call if no record is found.
+
+        Returns:
+            Model|Callable
+        """
+
+        if not callable(callback):
+            raise InvalidArgument("A callback must be callable.")
+
+        result = self.find(record_id=record_id)
+
+        if not result:
+            if not args:
+                return callback()
+            else:
+                return callback(*args)
+
+        return result
+
     def find_or_fail(self, record_id):
         """Finds a row by the primary key ID (Requires a model) or raise a ModelNotFound exception.
 

tests/sqlite/models/test_sqlite_model.py

@@ -82,6 +82,23 @@ def test_can_find_list(self):
             sql, """SELECT * FROM "users" WHERE "users"."id" IN ('1','2','3')"""
         )
 
+    def test_find_or_if_record_not_found(self):
+        # Insane record number so record cannot be found
+        record_id = 1_000_000_000_000_000
+
+        result = User.find_or(record_id, lambda: "Record not found.")
+        self.assertEqual(
+            result, "Record not found."
+        )
+
+    def test_find_or_if_record_found(self):
+        record_id = 1
+        result_id = User.find_or(record_id, lambda: "Record not found.").id
+
+        self.assertEqual(
+            result_id, record_id
+        )
+
     def test_can_set_and_retreive_attribute(self):
         user = User.hydrate({"id": 1, "name": "joe", "customer_id": 1})
         user.customer_id = "CUST1"