From 1b3fe22c8a9fc24f702b3796ccf92d8372445930 Mon Sep 17 00:00:00 2001
From: kalombo <nogamemorebrain@gmail.com>
Date: Sat, 15 Jun 2024 23:23:11 +0500
Subject: [PATCH] feat: add aio_get_or_create

---
 peewee_async.py                   | 21 +++++++++++++++++++++
 peewee_async_compat.py            |  8 ++++++++
 tests/aio_model/test_shortcuts.py | 15 +++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/peewee_async.py b/peewee_async.py
index 70239ef..e504dd0 100644
--- a/peewee_async.py
+++ b/peewee_async.py
@@ -882,3 +882,24 @@ async def aio_create(cls, **query):
         inst = cls(**query)
         await inst.aio_save(force_insert=True)
         return inst
+
+    @classmethod
+    async def aio_get_or_create(cls, **kwargs):
+        defaults = kwargs.pop('defaults', {})
+        query = cls.select()
+        for field, value in kwargs.items():
+            query = query.where(getattr(cls, field) == value)
+
+        try:
+            return await query.aio_get(), False
+        except cls.DoesNotExist:
+            try:
+                if defaults:
+                    kwargs.update(defaults)
+                async with cls._meta.database.aio_atomic():
+                    return await cls.aio_create(**kwargs), True
+            except peewee.IntegrityError as exc:
+                try:
+                    return await query.aio_get(), False
+                except cls.DoesNotExist:
+                    raise exc
diff --git a/peewee_async_compat.py b/peewee_async_compat.py
index ecd4e3e..d7776ac 100644
--- a/peewee_async_compat.py
+++ b/peewee_async_compat.py
@@ -265,6 +265,10 @@ async def get_or_create(self, model_, defaults=None, **kwargs):
         Return 2-tuple containing the model instance and a boolean
         indicating whether the instance was created.
         """
+        warnings.warn(
+            "`get_or_create` method is deprecated, use `AioModel.aio_get_or_create` instead.",
+            DeprecationWarning
+        )
         try:
             return (await self.get(model_, **kwargs)), False
         except model_.DoesNotExist:
@@ -342,6 +346,10 @@ async def create_or_get(self, model_, **kwargs):
         Try to create new object with specified data. If object already
         exists, then try to get it by unique fields.
         """
+        warnings.warn(
+            "`create_or_get` method is deprecated, use `AioModel.aio_get_or_create` instead.",
+            DeprecationWarning
+        )
         try:
             return (await self.create(model_, **kwargs)), True
         except IntegrityErrors:
diff --git a/tests/aio_model/test_shortcuts.py b/tests/aio_model/test_shortcuts.py
index 9816e2f..79a387e 100644
--- a/tests/aio_model/test_shortcuts.py
+++ b/tests/aio_model/test_shortcuts.py
@@ -109,3 +109,18 @@ async def test_aio_save__force_insert(db):
 
     with pytest.raises(peewee.IntegrityError):
         await t.aio_save(force_insert=True)
+
+
+@dbs_all
+async def test_aio_get_or_create__get(db):
+    t1 = await TestModel.aio_create(text="text", data="data")
+    t2, created = await TestModel.aio_get_or_create(text="text")
+    assert t1.id == t2.id
+    assert created is False
+
+
+@dbs_all
+async def test_aio_get_or_create__created(db):
+    t2, created = await TestModel.aio_get_or_create(text="text")
+    assert t2.text == "text"
+    assert created is True