added a batch insertion function

Author: yoann-diep

ecodev_core/db_insertion.py

@@ -10,19 +10,21 @@
 from typing import Union
 
 import pandas as pd
+import progressbar
 from fastapi import BackgroundTasks
 from fastapi import UploadFile
 from pandas import ExcelFile
 from sqlmodel import Session
 from sqlmodel import SQLModel
+from sqlmodel.main import SQLModelMetaclass
 from sqlmodel.sql.expression import SelectOfScalar
 
+from ecodev_core.db_upsertion import BATCH_SIZE
 from ecodev_core.logger import log_critical
 from ecodev_core.logger import logger_get
 from ecodev_core.pydantic_utils import CustomFrozen
 from ecodev_core.safe_utils import SimpleReturn
 
-
 log = logger_get(__name__)
 
 
@@ -72,7 +74,7 @@ async def insert_file(file: UploadFile, insertor: Insertor, session: Session) ->
     insert_data(df_raw, insertor, session)
 
 
-def insert_data(df:  Union[pd.DataFrame, ExcelFile], insertor: Insertor, session: Session) -> None:
+def insert_data(df: Union[pd.DataFrame, ExcelFile], insertor: Insertor, session: Session) -> None:
     """
     Inserts a csv/df into a database
     """
@@ -81,6 +83,24 @@ def insert_data(df:  Union[pd.DataFrame, ExcelFile], insertor: Insertor, session
     session.commit()
 
 
+def insert_batch_data(data: list[dict | SQLModelMetaclass],
+                      session: Session,
+                      raw_db_schema: SQLModelMetaclass | None = None) -> None:
+    """
+    Insert the passed list of dicts (corresponding to db_schema) into db_schema db.
+    Warning: this only inserts data, without checking for pre-existence.
+    Ensure deleting the data before using it to avoid duplicates.
+    """
+    db_schema = raw_db_schema or data[0].__class__
+    batches = [data[i:i + BATCH_SIZE] for i in range(0, len(data), BATCH_SIZE)]
+
+    for batch in progressbar.progressbar(batches, redirect_stdout=False):
+        for row in batch:
+            new_object = db_schema(**row) if isinstance(row, dict) else row
+            session.add(new_object)
+        session.commit()
+
+
 def create_or_update(session: Session, row: Dict, insertor: Insertor) -> SQLModel:
     """
     Create a new row in db if the selector insertor does not find existing row in db. Update the row

tests/functional/test_db_insertion.py

@@ -9,8 +9,10 @@
 from fastapi import status
 from fastapi import UploadFile
 from fastapi.testclient import TestClient
+from sqlmodel import Field
 from sqlmodel import select
 from sqlmodel import Session
+from sqlmodel import SQLModel
 
 from ecodev_core import AppUser
 from ecodev_core import attempt_to_log
@@ -27,9 +29,9 @@
 from ecodev_core import SimpleReturn
 from ecodev_core import upsert_app_users
 from ecodev_core.app_user import USER_INSERTOR
+from ecodev_core.db_insertion import insert_batch_data
 from ecodev_core.db_insertion import insert_file
 
-
 app = FastAPI()
 test_client = TestClient(app)
 DATA_DIR = Path('/app/tests/unitary/data')
@@ -141,3 +143,43 @@ def test_user_insertion(self):
                 test_client.post('/user-insert', files={'file': ('filename', f)},
                                  headers={'Authorization': f'Bearer {admin_token}'})
                 self.assertTrue(len(session.exec(select(AppUser)).all()) == 4)
+
+
+class InFoo(SQLModel, table=True):  # type: ignore
+    """
+    Test class to test DB insertion
+    """
+    __tablename__ = 'in_foo'
+    id: int | None = Field(default=None, primary_key=True)
+    bar1: int = Field()
+
+
+class BatchInsertorTest(SafeTestCase):
+    """
+    Class testing db batch insertion
+    """
+
+    def setUp(self):
+        """
+        Class set up
+        """
+        super().setUp()
+        create_db_and_tables(InFoo)
+        delete_table(InFoo)
+
+    def tearDown(self) -> None:
+        return super().tearDown()
+
+    def test_insert_batch_data(self):
+        """
+        Testing that the insert_batch_data function behaves as expected
+        """
+
+        data = [InFoo(bar1=1), InFoo(bar1=2), InFoo(bar1=3)]
+
+        with Session(engine) as session:
+            insert_batch_data(data, session)
+
+            db_data = session.exec(select(InFoo)).all()
+
+        self.assertEqual(len(db_data), 3)