get_or_create method (piccolo-orm#181)

* very simple draft of `get_or_create`

* add `GetOrCreate` and update docs

* added empty dict as default, and allow direct await for convenience

Co-authored-by: Daniel Townsend <[email protected]>

Author: aminalaee

docs/src/piccolo/query_types/objects.rst

@@ -89,6 +89,23 @@ object, you can do so using ``get_related``.
     >>> print(manager.name)
     'Guido'
 
+get_or_create
+-------------
+
+With ``get_or_create`` you can get an existing record matching the criteria,
+or create a new one with the ``defaults`` arguments:
+
+.. code-block:: python
+
+    band = Band.objects().get_or_create(
+        Band.name == 'Pythonistas', defaults={Band.popularity: 100}
+    ).run_sync()
+
+    # Or using string column names
+    band = Band.objects().get_or_create(
+        Band.name == 'Pythonistas', defaults={'popularity': 100}
+    ).run_sync()
+
 Query clauses
 -------------
 

piccolo/query/methods/objects.py

@@ -14,6 +14,7 @@
     WhereDelegate,
 )
 from piccolo.querystring import QueryString
+from piccolo.utils.sync import run_sync
 
 from .select import Select
 
@@ -22,6 +23,44 @@
     from piccolo.table import Table
 
 
+@dataclass
+class GetOrCreate:
+    query: Objects
+    where: Combinable
+    defaults: t.Dict[t.Union[Column, str], t.Any]
+
+    async def run(self):
+        instance = await self.query.where(self.where).first().run()
+        if instance:
+            return instance
+
+        instance = self.query.table()
+        setattr(
+            instance,
+            self.where.column._meta.name,  # type: ignore
+            self.where.value,  # type: ignore
+        )
+
+        for column, value in self.defaults.items():
+            if isinstance(column, str):
+                column = instance._meta.get_column_by_name(column)
+            setattr(instance, column._meta.name, value)
+
+        await instance.save().run()
+
+        return instance
+
+    def __await__(self):
+        """
+        If the user doesn't explicity call .run(), proxy to it as a
+        convenience.
+        """
+        return self.run().__await__()
+
+    def run_sync(self):
+        return run_sync(self.run())
+
+
 @dataclass
 class Objects(Query):
     """
@@ -64,6 +103,13 @@ def offset(self, number: int) -> Objects:
         self.offset_delegate.offset(number)
         return self
 
+    def get_or_create(
+        self,
+        where: Combinable,
+        defaults: t.Dict[t.Union[Column, str], t.Any] = {},
+    ):
+        return GetOrCreate(query=self, where=where, defaults=defaults)
+
     def order_by(self, *columns: Column, ascending=True) -> Objects:
         self.order_by_delegate.order_by(*columns, ascending=ascending)
         return self

tests/table/test_objects.py

@@ -59,3 +59,28 @@ def test_offset_sqlite(self):
         self.assertEqual(
             [i.name for i in response], ["Pythonistas", "Rustaceans"]
         )
+
+    def test_get_or_create(self):
+        Band.objects().get_or_create(
+            Band.name == "Pink Floyd", defaults={"popularity": 100}
+        ).run_sync()
+
+        instance = (
+            Band.objects().where(Band.name == "Pink Floyd").first().run_sync()
+        )
+
+        self.assertTrue(isinstance(instance, Band))
+        self.assertTrue(instance.name == "Pink Floyd")
+        self.assertTrue(instance.popularity == 100)
+
+        Band.objects().get_or_create(
+            Band.name == "Pink Floyd", defaults={Band.popularity: 100}
+        ).run_sync()
+
+        instance = (
+            Band.objects().where(Band.name == "Pink Floyd").first().run_sync()
+        )
+
+        self.assertTrue(isinstance(instance, Band))
+        self.assertTrue(instance.name == "Pink Floyd")
+        self.assertTrue(instance.popularity == 100)