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):
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||||
"""Context manager exit."""
|
"""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()
|
self.disconnect()
|
||||||
|
|
||||||
def execute(self, query: str, params: Optional[tuple] = None) -> None:
|
def execute(self, query: str, params: Optional[tuple] = None) -> None:
|
||||||
|
|||||||
Reference in New Issue
Block a user