feat: add db.save shortcut function

Author: amisadmin

README.md

@@ -94,7 +94,7 @@ from sqlalchemy import insert, select, update, delete
 async def fast_execute():
     # update
     stmt = update(User).where(User.id == 1).values({'username': 'new_user'})
-    result = await db.execute(stmt, commit=True)
+    result = await db.execute(stmt)
 
     # select
     stmt = select(User).where(User.id == 1)
@@ -121,7 +121,11 @@ async def fast_execute():
     # delete
     user = User(id=1, name='test')
     await db.delete(user)
-
+    
+    # save(insert or update)
+    user = User(name='new_user')
+    await db.save(user)
+    
     # run_sync
     await db.run_sync(Base.metadata.create_all, is_session=False)
 
@@ -154,7 +158,7 @@ from sqlalchemy import insert, select, update, delete
 def fast_execute():
     # update
     stmt = update(User).where(User.id == 1).values({'username': 'new_user'})
-    result = db.execute(stmt, commit=True)
+    result = db.execute(stmt)
 
     # select
     stmt = select(User).where(User.id == 1)
@@ -181,7 +185,11 @@ def fast_execute():
     # delete
     user = User(id=1, name='test')
     db.delete(user)
-
+    
+    # save(insert or update)
+    user = User(name='new_user')
+    db.save(user)
+    
     # run_sync
     db.run_sync(Base.metadata.create_all, is_session=False)
 
@@ -242,7 +250,11 @@ async def fast_execute(db: Union[AsyncDatabase, Database]):
     # delete
     user = User(id=1, name='test')
     await db.async_delete(user)
-
+    
+    # save(insert or update)
+    user = User(name='new_user')
+    await db.async_save(user)
+    
     # run_sync
     await db.async_run_sync(Base.metadata.create_all, is_session=False)
 

README.zh.md

@@ -122,6 +122,10 @@ async def fast_execute():
     user = User(id=1, name='test')
     await db.delete(user)
 
+    # save(insert or update)
+    user = User(name='new_user')
+    await db.save(user)
+    
     # run_sync
     await db.run_sync(Base.metadata.create_all, is_session=False)
 
@@ -182,6 +186,10 @@ def fast_execute():
     user = User(id=1, name='test')
     db.delete(user)
 
+    # save(insert or update)
+    user = User(name='new_user')
+    db.save(user)
+    
     # run_sync
     db.run_sync(Base.metadata.create_all, is_session=False)
 
@@ -239,7 +247,11 @@ async def fast_execute(db: Union[AsyncDatabase, Database]):
     # delete
     user = User(id=1, name='test')
     await db.async_delete(user)
-
+    
+    # save(insert or update)
+    user = User(name='new_user')
+    await db.async_save(user)
+    
     # run_sync
     await db.async_run_sync(Base.metadata.create_all, is_session=False)
 

sqlalchemy_database/_abc_async_database.py

@@ -23,7 +23,7 @@ async def to_thread(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs)
 class AbcAsyncDatabase(metaclass=abc.ABCMeta):
 
     def __init__(self) -> None:
-        for func_name in ['execute', 'scalar', 'scalars_all', 'get', 'delete', 'run_sync']:
+        for func_name in ['execute', 'scalar', 'scalars_all', 'get', 'delete', 'save', 'run_sync']:
             func = getattr(self, func_name)
             if not asyncio.iscoroutinefunction(func):
                 func = functools.partial(to_thread, func)

sqlalchemy_database/_abc_async_database.pyi

@@ -62,6 +62,8 @@ class AbcAsyncDatabase(metaclass=abc.ABCMeta):
 
     async def async_delete(self, instance: Any) -> None: ...
 
+    async def async_save(self, *instances: Any) -> None: ...
+
     async def async_run_sync(
             self,
             fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],
@@ -70,4 +72,4 @@ class AbcAsyncDatabase(metaclass=abc.ABCMeta):
             on_close_pre: Callable[[_T], _R] = None,
             is_session: bool = True,
             **kwargs: _P.kwargs
-    ) -> Union[_T,_R]: ...
+    ) -> Union[_T, _R]: ...

sqlalchemy_database/database.py

@@ -5,9 +5,10 @@
 from sqlalchemy.future import Engine, create_engine
 from sqlalchemy.orm import sessionmaker, Session
 from sqlalchemy.sql import Executable, Select
-from sqlalchemy_database._abc_async_database import AbcAsyncDatabase
 from typing_extensions import Concatenate, ParamSpec
 
+from sqlalchemy_database._abc_async_database import AbcAsyncDatabase
+
 _P = ParamSpec("_P")
 _T = TypeVar("_T")
 _R = TypeVar("_R")
@@ -53,8 +54,7 @@ async def execute(
         else:
             maker = self.engine.connect
         async with maker() as conn:
-
-            result = await conn.execute(statement, params, execution_options, **kw)
+            result = await conn.execute(statement, params, execution_options, **kw)  # type:ignore
             if on_close_pre:
                 result = on_close_pre(result)
             if commit and not isinstance(statement, Select):
@@ -122,6 +122,11 @@ async def delete(self, instance: Any) -> None:
             async with session.begin():
                 await session.delete(instance)
 
+    async def save(self, *instances: Any) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                session.add_all(instances)
+
     async def run_sync(
             self,
             fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],
@@ -246,6 +251,11 @@ def delete(self, instance: Any) -> None:
             with session.begin():
                 session.delete(instance)
 
+    def save(self, *instances: Any) -> None:
+        with self.session_maker() as session:
+            with session.begin():
+                session.add_all(instances)
+
     def run_sync(
             self,
             fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],

tests/test_AbcAsyncDatabase.py

@@ -91,6 +91,21 @@ async def test_async_delete(db, fake_users):
     assert user is None
 
 
+async def test_async_save(db, fake_users):
+    # test update
+    user = await db.async_get(User, 1)
+    assert user.id == 1
+    user.username = 'new_user'
+    await db.async_save(user)
+    user = await db.async_get(User, 1)
+    assert user.username == 'new_user'
+    # test insert
+    user2 = User(username='new_user2')
+    await db.async_save(user2)
+    u = await db.async_scalar(select(User).where(User.username == 'new_user2'))
+    assert u.username == 'new_user2'
+
+
 async def test_async_run_sync(db, fake_users):
     def delete_user(session: Session, instance: User):
         session.delete(instance)

tests/test_AsyncDatabase.py

@@ -98,6 +98,21 @@ async def test_delete(fake_users):
     assert user is None
 
 
+async def test_save(fake_users):
+    # test update
+    user = await db.get(User, 1)
+    assert user.id == 1
+    user.username = 'new_user'
+    await db.save(user)
+    user = await db.get(User, 1)
+    assert user.username == 'new_user'
+    # test insert
+    user2 = User(username='new_user2')
+    await db.save(user2)
+    u = await db.scalar(select(User).where(User.username == 'new_user2'))
+    assert u.username == 'new_user2'
+
+
 async def test_run_sync(fake_users):
     def delete_user(session: Session, instance: User):
         session.delete(instance)

tests/test_Database.py

@@ -98,6 +98,21 @@ def test_delete(fake_users):
     assert user is None
 
 
+def test_save(fake_users):
+    # test update
+    user = db.get(User, 1)
+    assert user.id == 1
+    user.username = 'new_user'
+    db.save(user)
+    user = db.get(User, 1)
+    assert user.username == 'new_user'
+    # test insert
+    user2 = User(username='new_user2')
+    db.save(user2)
+    u = db.scalar(select(User).where(User.username == 'new_user2'))
+    assert u.username == 'new_user2'
+
+
 def test_run_sync(fake_users):
     def delete_user(session: Session, instance: User):
         session.delete(instance)