feat: Add the refresh option to the db.save shortcut function.

Author: amisadmin

sqlalchemy_database/database.py

@@ -271,6 +271,8 @@ async def save(self, *instances: Any, refresh: bool = False) -> None:
         async with self.session_maker() as session:
             async with session.begin():
                 session.add_all(instances)
+            if refresh:
+                [await session.refresh(instance) for instance in instances]
 
     async def run_sync(
             self,
@@ -435,6 +437,8 @@ def save(self, *instances: Any, refresh: bool = False) -> None:
         with self.session_maker() as session:
             with session.begin():
                 session.add_all(instances)
+            if refresh:
+                [session.refresh(instance) for instance in instances]
 
     def run_sync(
             self,

tests/test_AbcAsyncDatabase.py

@@ -4,7 +4,6 @@
 import pytest
 from sqlalchemy import insert, select, update, delete
 from sqlalchemy.orm import Session
-
 from sqlalchemy_database import AsyncDatabase, Database
 from tests.conftest import async_db, sync_db, Base, User
 
@@ -104,6 +103,10 @@ async def test_async_save(db, fake_users):
     await db.async_save(user2)
     u = await db.async_scalar(select(User).where(User.username == 'new_user2'))
     assert u.username == 'new_user2'
+    # test refresh
+    user3 = User(username='new_user3')
+    await db.async_save(user3, refresh=True)
+    assert user3.id
 
 
 async def test_async_run_sync(db, fake_users):

tests/test_AsyncDatabase.py

@@ -111,7 +111,10 @@ async def test_save(fake_users):
     await db.save(user2)
     u = await db.scalar(select(User).where(User.username == 'new_user2'))
     assert u.username == 'new_user2'
-
+    # test refresh
+    user3 = User(username='new_user3')
+    await db.save(user3, refresh=True)
+    assert user3.id
 
 async def test_run_sync(fake_users):
     def delete_user(session: Session, instance: User):

tests/test_Database.py

@@ -5,7 +5,6 @@
 import pytest
 from sqlalchemy import insert, select, update, delete
 from sqlalchemy.orm import Session
-
 from tests.conftest import sync_db as db, Base, User
 
 
@@ -111,6 +110,10 @@ def test_save(fake_users):
     db.save(user2)
     u = db.scalar(select(User).where(User.username == 'new_user2'))
     assert u.username == 'new_user2'
+    # test refresh
+    user3 = User(username='new_user3')
+    db.save(user3, refresh=True)
+    assert user3.id
 
 
 def test_run_sync(fake_users):