Merge pull request #123 from schireson/dc/idle-in-transaction

fix: Add ability to set session-local idle-in-transaction setting.

Author: DanCardin

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+### v2.9.1
+
+* fix: Add ability to set session-local idle-in-transaction setting
+
 ### v2.9.0
 
 * Swap to psycopg from psycopg2

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databudgie"
-version = "2.9.0"
+version = "2.9.1"
 description = "Ergonomic and flexible tool for database backup and restore"
 readme = "README.md"
 license = { text = "MIT" }
@@ -57,21 +57,7 @@ src = ["src", "tests"]
 [tool.ruff.lint]
 select = ["C", "D", "E", "F", "I", "N", "Q", "RET", "RUF", "S", "T", "UP", "YTT"]
 ignore = ["C901", "E501", "S101", "D1", "D203", "D213", "D406", "D407", "D408", "D409", "D413"]
-extend-ignore = [
-    "D1",
-
-    "D203",
-    "D204",
-    "D213",
-    "D215",
-    "D400",
-    "D404",
-    "D406",
-    "D407",
-    "D408",
-    "D409",
-    "D413",
-]
+extend-ignore = ["D1", "D203", "D204", "D213", "D215", "D400", "D404", "D406", "D407", "D408", "D409", "D413"]
 
 [tool.ruff.lint.isort]
 order-by-type = false

src/databudgie/cli/base.py

@@ -43,7 +43,14 @@ def _create_postgres_session(config: BackupConfig | RestoreConfig):
         url_obj = sqlalchemy.engine.url.make_url(url)
 
     engine = sqlalchemy.create_engine(url_obj)
-    return sqlalchemy.orm.scoping.scoped_session(sqlalchemy.orm.session.sessionmaker(bind=engine))()
+    session = sqlalchemy.orm.scoping.scoped_session(sqlalchemy.orm.session.sessionmaker(bind=engine))()
+
+    if config.idle_in_transaction_timeout is not None:
+        session.execute(
+            sqlalchemy.text(f"SET idle_in_transaction_session_timeout = '{config.idle_in_transaction_timeout}s'")
+        )
+
+    return session
 
 
 def backup_config(root_config: RootConfig):

src/databudgie/config.py

@@ -141,6 +141,7 @@ class TableParentConfig(typing.Generic[T], Config):
     s3: S3Config | None = None
     root_location: str | None = None
     adapter: str | None = None
+    idle_in_transaction_timeout: int | None = None
 
     @classmethod
     @abc.abstractmethod
@@ -168,6 +169,8 @@ def from_stack(cls, stack: ConfigStack):
 
         connections = Connection.from_collection(stack.get("connections"))
 
+        idle_in_transaction_timeout = stack.get("idle_in_transaction_timeout")
+
         return cls(
             connection=connection,
             tables=tables,
@@ -177,6 +180,7 @@ def from_stack(cls, stack: ConfigStack):
             root_location=root_location,
             adapter=adapter,
             connections=connections,
+            idle_in_transaction_timeout=idle_in_transaction_timeout,
         )
 
 

tests/test_backup.py

@@ -195,6 +195,28 @@ def test_backup_unnamed_table(pg, mf, s3_resource):
     )
 
 
+def test_backup_with_idle_in_transaction_timeout(pg, s3_resource):
+    """Validate backup works with idle_in_transaction_timeout config option."""
+    config = RootConfig.from_dict(
+        {
+            "location": "s3://sample-bucket/databudgie/test/{table}",
+            "tables": ["public.customer", "public.store"],
+            "sequences": False,
+            "strict": True,
+            "idle_in_transaction_timeout": 5000,
+            **s3_config,
+        }
+    )
+
+    backup_all(pg, config.backup)
+
+    all_object_keys = [obj.key for obj in s3_resource.Bucket("sample-bucket").objects.all()]
+    assert all_object_keys == [
+        "databudgie/test/public.customer/2021-04-26T09:00:00.csv",
+        "databudgie/test/public.store/2021-04-26T09:00:00.csv",
+    ]
+
+
 def _validate_backup_contents(buffer, expected_contents: List[Customer]):
     """Validate the contents of a backup file. Columns from the file will be raw."""
     wrapper = io.TextIOWrapper(buffer, encoding="utf-8")