fix: Add explicit commit/rollback in PostgreSQL context manager exit
- On successful execution (no exception): explicitly commit before closing - On exception: explicitly rollback before closing - Add try-except to handle commit/rollback failures gracefully This ensures that all inserted data is committed to the database when the context manager exits. Previously, commits were only done per-batch in insert_batch(), but the final context exit wasn't ensuring a final commit. Fixes: Data not appearing in PostgreSQL after migration completes
This commit is contained in:
@@ -49,6 +49,18 @@ class PostgreSQLConnector:
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit."""
|
||||
if exc_type is None:
|
||||
# No exception, commit before closing
|
||||
try:
|
||||
self.connection.commit()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to commit on exit: {e}")
|
||||
else:
|
||||
# Exception occurred, rollback
|
||||
try:
|
||||
self.connection.rollback()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to rollback on exit: {e}")
|
||||
self.disconnect()
|
||||
|
||||
def execute(self, query: str, params: Optional[tuple] = None) -> None:
|
||||
|
||||
Reference in New Issue
Block a user