make get_or_create work for complex where clauses (piccolo-orm#195)

* make get_or_create work for complex where clauses

* update get_or_create docs

* fix LGTM warning, and add `_was_created` attribute

Author: dantownsend

docs/src/piccolo/query_types/objects.rst

@@ -113,6 +113,32 @@ or create a new one with the ``defaults`` arguments:
         Band.name == 'Pythonistas', defaults={'popularity': 100}
     ).run_sync()
 
+You can find out if an existing row was found, or if a new row was created:
+
+.. code-block:: python
+
+    band = Band.objects.get_or_create(
+        Band.name == 'Pythonistas'
+    ).run_sync()
+    band._was_created  # True if it was created, otherwise False if it was already in the db
+
+Complex where clauses are supported, but only within reason. For example:
+
+.. code-block:: python
+
+    # This works OK:
+    band = Band.objects().get_or_create(
+        (Band.name == 'Pythonistas') & (Band.popularity == 1000),
+    ).run_sync()
+
+    # This is problematic, as it's unclear what the name should be if we
+    # need to create the row:
+    band = Band.objects().get_or_create(
+        (Band.name == 'Pythonistas') | (Band.name == 'Rustaceans'),
+        defaults={'popularity': 100}
+    ).run_sync()
+
+
 Query clauses
 -------------
 

piccolo/columns/combination.py

@@ -2,7 +2,7 @@
 
 import typing as t
 
-from piccolo.columns.operators.comparison import ComparisonOperator
+from piccolo.columns.operators.comparison import ComparisonOperator, Equal
 from piccolo.custom_types import Combinable, Iterable
 from piccolo.querystring import QueryString
 from piccolo.utils.sql_values import convert_to_sql_value
@@ -44,6 +44,35 @@ def __str__(self):
 class And(Combination):
     operator = "AND"
 
+    def get_column_values(self) -> t.Dict[Column, t.Any]:
+        """
+        This is used by `get_or_create` to know which values to assign if
+        the row doesn't exist in the database.
+
+        For example, if we have:
+
+        (Band.name == 'Pythonistas') & (Band.popularity == 1000)
+
+        We will return {Band.name: 'Pythonistas', Band.popularity: 1000}.
+
+        If the operator is anything besides equals, we don't return it, for
+        example:
+
+        (Band.name == 'Pythonistas') & (Band.popularity > 1000)
+
+        Returns {Band.name: 'Pythonistas'}
+
+        """
+        output = {}
+        for combinable in (self.first, self.second):
+            if isinstance(combinable, Where):
+                if combinable.operator == Equal:
+                    output[combinable.column] = combinable.value
+            elif isinstance(combinable, And):
+                output.update(combinable.get_column_values())
+
+        return output
+
 
 class Or(Combination):
     operator = "OR"

piccolo/query/methods/objects.py

@@ -3,6 +3,7 @@
 import typing as t
 from dataclasses import dataclass
 
+from piccolo.columns.combination import And, Where
 from piccolo.custom_types import Combinable
 from piccolo.engine.base import Batch
 from piccolo.query.base import Query
@@ -32,14 +33,22 @@ class GetOrCreate:
     async def run(self):
         instance = await self.query.get(self.where).run()
         if instance:
+            instance._was_created = False
             return instance
 
         instance = self.query.table()
-        setattr(
-            instance,
-            self.where.column._meta.name,  # type: ignore
-            self.where.value,  # type: ignore
-        )
+
+        # If it's a complex `where`, there can be several column values to
+        # extract e.g. (Band.name == 'Pythonistas') & (Band.popularity == 1000)
+        if isinstance(self.where, Where):
+            setattr(
+                instance,
+                self.where.column._meta.name,  # type: ignore
+                self.where.value,  # type: ignore
+            )
+        elif isinstance(self.where, And):
+            for column, value in self.where.get_column_values().items():
+                setattr(instance, column._meta.name, value)
 
         for column, value in self.defaults.items():
             if isinstance(column, str):
@@ -48,6 +57,8 @@ async def run(self):
 
         await instance.save().run()
 
+        instance._was_created = True
+
         return instance
 
     def __await__(self):

tests/table/test_objects.py

@@ -68,6 +68,10 @@ def test_get(self):
         self.assertTrue(band.name == "Pythonistas")
 
     def test_get_or_create(self):
+        """
+        Make sure `get_or_create` works for simple where clauses.
+        """
+        # When the row doesn't exist in the db:
         Band.objects().get_or_create(
             Band.name == "Pink Floyd", defaults={"popularity": 100}
         ).run_sync()
@@ -76,10 +80,11 @@ def test_get_or_create(self):
             Band.objects().where(Band.name == "Pink Floyd").first().run_sync()
         )
 
-        self.assertTrue(isinstance(instance, Band))
+        self.assertIsInstance(instance, Band)
         self.assertTrue(instance.name == "Pink Floyd")
         self.assertTrue(instance.popularity == 100)
 
+        # When the row already exists in the db:
         Band.objects().get_or_create(
             Band.name == "Pink Floyd", defaults={Band.popularity: 100}
         ).run_sync()
@@ -88,6 +93,70 @@ def test_get_or_create(self):
             Band.objects().where(Band.name == "Pink Floyd").first().run_sync()
         )
 
-        self.assertTrue(isinstance(instance, Band))
+        self.assertIsInstance(instance, Band)
         self.assertTrue(instance.name == "Pink Floyd")
         self.assertTrue(instance.popularity == 100)
+
+    def test_get_or_create_complex(self):
+        """
+        Make sure `get_or_create` works with complex where clauses.
+        """
+        self.insert_rows()
+
+        # When the row already exists in the db:
+        instance = (
+            Band.objects()
+            .get_or_create(
+                (Band.name == "Pythonistas") & (Band.popularity == 1000)
+            )
+            .run_sync()
+        )
+        self.assertIsInstance(instance, Band)
+        self.assertEqual(instance._was_created, False)
+
+        # When the row doesn't exist in the db:
+        instance = (
+            Band.objects()
+            .get_or_create(
+                (Band.name == "Pythonistas2") & (Band.popularity == 2000)
+            )
+            .run_sync()
+        )
+        self.assertIsInstance(instance, Band)
+        self.assertEqual(instance._was_created, True)
+
+    def test_get_or_create_very_complex(self):
+        """
+        Make sure `get_or_create` works with very complex where clauses.
+        """
+        self.insert_rows()
+
+        # When the row already exists in the db:
+        instance = (
+            Band.objects()
+            .get_or_create(
+                (Band.name == "Pythonistas")
+                & (Band.popularity > 0)
+                & (Band.popularity < 5000)
+            )
+            .run_sync()
+        )
+        self.assertIsInstance(instance, Band)
+        self.assertEqual(instance._was_created, False)
+
+        # When the row doesn't exist in the db:
+        instance = (
+            Band.objects()
+            .get_or_create(
+                (Band.name == "Pythonistas2")
+                & (Band.popularity > 10)
+                & (Band.popularity < 5000)
+            )
+            .run_sync()
+        )
+        self.assertIsInstance(instance, Band)
+        self.assertEqual(instance._was_created, True)
+
+        # The values in the > and < should be ignored, and the default should
+        # be used for the column.
+        self.assertEqual(instance.popularity, 0)