Add Truncate Functionality (#14)

Author: JonasScholl

.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 .DS_Store
 
 test.db
+test.db-journal
 
 # C extensions
 *.so

database_setup_tools/setup.py

@@ -1,8 +1,9 @@
 import threading
-from typing import Optional
+from typing import List, Optional
 
 import sqlalchemy_utils
-from sqlalchemy import MetaData
+from sqlalchemy import MetaData, Table
+from sqlmodel import SQLModel
 
 from database_setup_tools.session_manager import SessionManager
 
@@ -47,6 +48,15 @@ def model_metadata(self) -> MetaData:
         """
         return self._model_metadata
 
+    @property
+    def session_manager(self) -> SessionManager:
+        """Getter for the session manager
+
+        Returns:
+            SessionManager: The session manager
+        """
+        return SessionManager(database_uri=self.database_uri)
+
     @property
     def database_uri(self) -> str:
         """Getter for the database URI
@@ -71,11 +81,26 @@ def create_database(self) -> bool:
         """Create the database and the tables if not done yet"""
         if not sqlalchemy_utils.database_exists(self.database_uri):
             sqlalchemy_utils.create_database(self.database_uri)
-            session_manager = SessionManager(self.database_uri)
-            self.model_metadata.create_all(session_manager.engine)
+            self.model_metadata.create_all(self.session_manager.engine)
             return True
         return False
 
+    def truncate(self, tables: Optional[List[SQLModel]] = None):
+        """Truncate all tables in the database"""
+        tables_to_truncate: List[Table] = self.model_metadata.sorted_tables
+        if tables is not None:
+            table_names = [table.__tablename__ for table in tables]
+            tables_to_truncate = filter(lambda table: table.name in table_names, tables_to_truncate)
+
+        session = next(self.session_manager.get_session())
+
+        try:
+            tables_with_schema = [f"{table.schema or 'public'}.\"{table.name}\"" for table in tables_to_truncate]
+            session.execute(f"TRUNCATE TABLE {', '.join(tables_with_schema)} CASCADE;")
+            session.commit()
+        finally:
+            session.close()
+
     @classmethod
     def _get_cached_instance(cls, args: tuple, kwargs: dict) -> Optional[object]:
         """Provides a cached instance of the SessionManager class if existing"""

poetry.lock

@@ -1,7 +1,5 @@
-SQLITE_DATABASE_URI = "sqlite:///test.db"
 POSTGRESQL_DATABASE_URI = "postgresql+psycopg2://postgres:postgres@localhost:5432/test"
 
 DATABASE_URIS = [
-    SQLITE_DATABASE_URI,
     POSTGRESQL_DATABASE_URI,
 ]

tests/integration/database_config.py

@@ -3,45 +3,104 @@
 import pytest
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.orm.scoping import ScopedSession
+from sqlmodel import Field, SQLModel
 
 from database_setup_tools.session_manager import SessionManager
 from database_setup_tools.setup import DatabaseSetup
 from tests.integration.database_config import DATABASE_URIS
-from tests.sample_model import User, model_metadata
+from tests.sample_model import Customer, model_metadata
 
 
 @pytest.mark.parametrize("database_uri", DATABASE_URIS)
 class TestDatabaseIntegration:
     @pytest.fixture
     def database_setup(self, database_uri: str) -> DatabaseSetup:
         setup = DatabaseSetup(model_metadata=model_metadata, database_uri=database_uri)
+        setup.drop_database()
+        setup.create_database()
         yield setup
         setup.drop_database()
 
     @pytest.fixture
-    def database_session(self, database_uri: str) -> Iterator[ScopedSession]:
+    def database_session(self, database_setup: DatabaseSetup) -> Iterator[ScopedSession]:
         """Get a database session"""
-        session_manager = SessionManager(database_uri)
-        return next(session_manager.get_session())
+        return next(database_setup.session_manager.get_session())
 
     def test_create_database_and_tables(self, database_setup: DatabaseSetup, database_session: ScopedSession):
         """Test that the tables are created correctly"""
-        # noinspection SqlInjection,SqlDialectInspection
-        database_session.execute(f"SELECT * FROM {User.__tablename__}")
+        database_session.execute(f"SELECT * FROM {Customer.__tablename__}")
 
     def test_create_database_multiple_times(self, database_setup: DatabaseSetup, database_session: ScopedSession):
         """Test that creating the database multiple times does not cause problems"""
         database_setup.create_database()
-        # noinspection SqlInjection,SqlDialectInspection
-        database_session.execute(f"SELECT * FROM {User.__tablename__}")
+        database_session.execute(f"SELECT * FROM {Customer.__tablename__}")
 
     def test_drop_database(self, database_setup: DatabaseSetup, database_session: ScopedSession):
         """Test that the database is dropped correctly"""
-        database_setup.create_database()
         assert database_setup.drop_database() is True
 
         with pytest.raises(OperationalError):
-            # noinspection SqlDialectInspection
-            database_session.execute(f"SELECT * FROM {User.__tablename__}")
+            database_session.execute(f"SELECT * FROM {Customer.__tablename__}")
 
         assert database_setup.drop_database() is False
+
+    def test_truncate_all_tables(self, database_setup: DatabaseSetup, database_session: ScopedSession):
+        """Test that all tables are truncated correctly"""
+
+        setup_statements = [
+            f"""CREATE TABLE delivery (
+                    id INTEGER,
+                    name TEXT NOT NULL,
+                    customer_id INTEGER,
+                    PRIMARY KEY(id),
+                    CONSTRAINT fk_user
+                        FOREIGN KEY(customer_id)
+                            REFERENCES "{Customer.__tablename__}"(id)
+                )
+            """,
+            f'SELECT * FROM "{Customer.__tablename__}"',
+            'SELECT * FROM "delivery"',
+            f"INSERT INTO \"{Customer.__tablename__}\" VALUES (1, 'John Doe')",
+            "INSERT INTO \"delivery\" VALUES (1, 'Delivery 1', 1)",
+        ]
+        for statement in setup_statements:
+            database_session.execute(statement)
+
+        assert database_session.execute(f"SELECT * FROM {Customer.__tablename__}").rowcount == 1
+        assert database_session.execute(f'SELECT * FROM "delivery"').rowcount == 1
+        database_session.commit()
+
+        database_setup.truncate()
+
+        assert database_session.execute(f"SELECT * FROM {Customer.__tablename__}").rowcount == 0
+        assert database_session.execute(f'SELECT * FROM "delivery"').rowcount == 0
+
+    def test_truncate_custom_tables(self, database_uri: str):
+        """Test that only specified tables are truncated correctly"""
+
+        class TableToTruncate(SQLModel, table=True):
+            id: int = Field(index=True, primary_key=True)
+            name: str
+
+        setup = DatabaseSetup(model_metadata=model_metadata, database_uri=database_uri)
+        setup.drop_database()
+        setup.create_database()
+        database_session = next(setup.session_manager.get_session())
+
+        setup_statements = [
+            f'SELECT * FROM "{Customer.__tablename__}"',
+            f'SELECT * FROM "{TableToTruncate.__tablename__}"',
+            f"INSERT INTO \"{Customer.__tablename__}\" VALUES (1, 'John Doe')",
+            f"INSERT INTO \"{TableToTruncate.__tablename__}\" VALUES (1, 'Test')",
+        ]
+        for statement in setup_statements:
+            database_session.execute(statement)
+
+        assert database_session.execute(f"SELECT * FROM {Customer.__tablename__}").rowcount == 1
+        assert database_session.execute(f'SELECT * FROM "{TableToTruncate.__tablename__}"').rowcount == 1
+        database_session.commit()
+
+        setup.truncate(tables=[TableToTruncate])
+
+        assert database_session.execute(f"SELECT * FROM {Customer.__tablename__}").rowcount == 1
+        assert database_session.execute(f'SELECT * FROM "{TableToTruncate.__tablename__}"').rowcount == 0

tests/integration/test_database_integration.py

@@ -1,8 +1,8 @@
-from sqlmodel import SQLModel, Field
+from sqlmodel import Field, SQLModel
 
 
-class User(SQLModel, table=True):
-    """User model"""
+class Customer(SQLModel, table=True):
+    """Customer model"""
 
     id: int = Field(index=True, primary_key=True)
     name: str

tests/sample_model.py

@@ -43,12 +43,15 @@ def test_create_database_setup_fail_database_uri_invalid_type(invalid_database_u
         with pytest.raises(TypeError):
             DatabaseSetup(model_metadata=model_metadata, database_uri=invalid_database_uri)
 
-    def test_database_uri(self, database_setup: DatabaseSetup, database_uri: str):
+    def test_database_uri_property(self, database_setup: DatabaseSetup, database_uri: str):
         assert database_setup.database_uri == database_uri
 
-    def test_model_metadata(self, database_setup: DatabaseSetup):
+    def test_model_metadata_property(self, database_setup: DatabaseSetup):
         assert database_setup.model_metadata == model_metadata
 
+    def test_session_manager_property(self, database_setup: DatabaseSetup, database_uri: str):
+        assert database_setup.session_manager == SessionManager(database_uri=database_uri)
+
     def test_create_database(self, database_setup: DatabaseSetup, database_uri: str, when: Callable, expect: Callable):
         unstub()  # remove stub for create_database method