Migrate from psycopg2 to psycopg3 by estoel81 · Pull Request #1549 · workfloworchestrator/orchestrator-core

estoel81 · 2026-04-14T09:04:40Z

Summary

Replace psycopg2-binary with psycopg[binary]>=3.3.2 and update all postgresql:// URIs to postgresql+psycopg:// across CI, docs, settings, and test fixtures
Fix idle-in-transaction connections caused by psycopg3's autobegin: add pool checkin rollback, pool_recycle, and wrap all bare DB reads on Celery workers in transactional()
Make transactional() reentrant so nested calls (e.g. a step wrapping a Celery task entry point) don't roll back the outer transaction
Serialize step state to a plain dict in _db_log_step before the outer flush, preventing re-evaluation of expensive @computed_field properties and ensuring the next step never receives live SubscriptionModel instances
Add OrchestratorCore.register_table() to copy extra column_property attributes from a custom table subclass onto the base mapper
Emit DeprecationWarning when DATABASE_URI uses the bare postgresql:// scheme

Test Plan

uv run pytest test/unit_tests/db/test_database.py — transactional reentrancy
uv run pytest test/unit_tests/services/test_processes.py — state serialization + idle-in-transaction
uv run pytest test/unit_tests/test_subscription_table_registry.py — register_table
Full unit test suite: uv run pytest test/unit_tests/

Replace psycopg2-binary with psycopg[binary]>=3.3.3. psycopg3 uses autobegin, meaning any SELECT starts an implicit transaction that stays open until an explicit commit/rollback. This PR fixes the resulting idle-in-transaction connections by managing transaction boundaries explicitly throughout the workflow execution path. Key changes: database.py: - Add pool_recycle=3600 and a checkin event listener that rolls back any leftover transaction when a connection returns to the pool. - Make transactional() reentrant: nested calls yield without commit/rollback; the outermost call owns the transaction. - Disable expire_on_commit during commit so ORM objects retain their in-memory values and avoid lazy-load queries that would start new unmanaged transactions. services/processes.py: - JSON-roundtrip current_step.state in _db_log_step before flush. This serializes live SubscriptionModel instances exactly once, preventing double model_dump on the outer flush and ensuring the next step receives a clean plain-dict state. workflow.py / threadpool.py / tasks.py: - Wrap engine_status check, post-step logging, retrieve_input_state, and Celery task entry points in transactional() so every DB query on the Celery worker's empty-scope session runs inside a managed transaction. app.py: - Emit DeprecationWarning when DATABASE_URI uses the bare 'postgresql://' scheme (defaults to psycopg2 driver). - Add OrchestratorCore.register_table() to copy extra column properties from a custom table subclass onto the base mapper. Update all DATABASE_URI references in CI, docs, README, settings, and test fixtures to use 'postgresql+psycopg://'.

Signed-off-by: Erik Stoel <estoel81@hotmail.com>

sentry · 2026-04-14T09:11:51Z

Codecov Report

❌ Patch coverage is 94.28571% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.19%. Comparing base (b3b2643) to head (f94abe8).

Files with missing lines	Patch %	Lines
orchestrator/app.py	81.25%	2 Missing and 1 partial ⚠️
orchestrator/cli/generator/generator/migration.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1549      +/-   ##
==========================================
+ Coverage   89.84%   90.19%   +0.34%     
==========================================
  Files         270      270              
  Lines       13707    13743      +36     
  Branches     1342     1346       +4     
==========================================
+ Hits        12315    12395      +80     
+ Misses       1116     1071      -45     
- Partials      276      277       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

codspeed-hq · 2026-04-14T09:11:53Z

Merging this PR will not alter performance

✅ 13 untouched benchmarks

_{Comparing psycopg3-clean (f94abe8) with main (b3b2643)}

- Fix Mapper[Any] assignment in register_table() via cast() - Fix tasks.py: kombu attr-defined, return type, remove stale type: ignore comments - Add test_pool_events.py: pool checkin rollback registration and handler - Add test_database_uri_validation.py: DeprecationWarning for postgresql:// URI - Add test_dbutils.py: handle_missing_tables() with psycopg3 error types

…ion and task execution

Tests cover the psycopg3-critical transactional() wrapping in start_process and resume_process, plus get_celery_task, register_custom_serializer, and initialise_celery task route/registration behaviour.

… SQL

pboers1988

LGTM minor nitpicks

pboers1988

@estoel81 Also need some 5.0.0 upgrade docs.

The pool checkin rollback handler is no longer needed with the removal of pool_recycle, and the DATABASE_URI validation tests are obsolete. This cleanup eliminates the _register_pool_events() function and its associated test coverage.

# Conflicts: # pyproject.toml # uv.lock

…de guide

tjeerddie · 2026-04-15T07:46:58Z

+
+class MySubscriptionTable(SubscriptionTable):
+    customer_name: Mapped[str] = column_property(
+        select(CustomerTable.fullname)


maybe add a table example above this, the only reference of CustomerTable is somewhere in the docs about adding it, could also link to that place

tjeerddie · 2026-04-15T08:51:23Z

+    # Serialize state to a plain dict now, before the outer transactional() flushes it.
+    # This ensures the next step never receives live SubscriptionModel instances and
+    # avoids re-evaluating expensive @computed_field properties on flush.
+    current_step.state = json_loads(json_dumps(current_step.state))


maybe override the pydanticmodel.model_dump, so its in done in the model

tjeerddie · 2026-04-15T08:57:32Z

-            ensure_correct_process_status(process_id, ProcessStatus.CREATED)
+            # Celery workers use the empty-scope session (no database_scope); wrap in
+            # transactional() to prevent psycopg3 autobegin leaving the connection idle-in-transaction.
+            with transactional(db, local_logger):


could move the transactional inside the _get_process, load_process and ensure_correct_process_status. so any re-uses in other places don't need to be wrapped inside a transactional.

tjeerddie · 2026-04-15T08:58:54Z


        self.include_router(api_router, prefix="/api")

+        # Validate DATABASE_URI dialect before initializing the database.


I don't believe this comment is necessary

tjeerddie · 2026-04-15T09:06:58Z

-                    "Not executing Step as the workflow engine is Paused. Process will remain in state 'running'"
-                )
-                return process
+            with transactional(db, logger):


move the transactional inside the get_engine_settings_table?

tjeerddie · 2026-04-15T09:08:43Z

+        # Wrap logging in transactional() so SELECTs triggered by mutationlogger and
+        # dblogstep run inside a managed transaction (psycopg3 autobegin prevention).
+        with transactional(db, logger):
+            result_to_log.on_success(mutationlogger).on_failed(errorlogger).on_waiting(errorlogger)


I don't think mutationlogger triggers any selects and move the transactional inside the dblogstep

tjeerddie · 2026-04-15T09:12:47Z

+    log = MagicMock()
+
+    with transactional(db, log):
+        pass


wouldn't you need a db insert or db error before you can check that a commit or rollback isn't called.

tjeerddie · 2026-04-15T09:16:39Z

+
+def test_handle_missing_tables_suppresses_undefined_table_error() -> None:
+    """ProgrammingError wrapping UndefinedTable is caught and not re-raised."""
+    orig = psycopg_errors.UndefinedTable.__new__(psycopg_errors.UndefinedTable)


this looks weird, cant you directly make the class?

Suggested change

orig = psycopg_errors.UndefinedTable.__new__(psycopg_errors.UndefinedTable)

orig = psycopg_errors.UndefinedTable()

tjeerddie · 2026-04-15T09:20:37Z

+    original_description = base_mapper.column_attrs["description"]
+
+    class CustomSubscriptionTable(SubscriptionTable):
+        pass


you aren't overwriting the description here?

tjeerddie · 2026-04-15T09:23:06Z

+
+
+@pytest.mark.usefixtures("_cleanup_extra_field")
+def test_register_table_columns_visible_in_inspect():


this seems redundant since test_register_table_copies_column_properties already checks if the extra_fields is inside mapper.column_attrs

tjeerddie · 2026-04-15T09:32:13Z

+    # Build a minimal in-memory workflow function. The workflow is never executed
+    # because we mock _run_process_async; we only need a non-removed_workflow value
+    # for the early branch in thread_start_process and for ProcessStat construction.
+    wf = make_workflow(_workflow_test_fn, "wf description", None, Target.SYSTEM, StepList())


can't you use the simple_workflow?

tjeerddie · 2026-04-15T09:33:05Z

+
+    # Mock _run_process_async to skip actually running the workflow thread; we are
+    # only testing the pre-runwf section of thread_start_process.
+    with mock.patch("orchestrator.services.executors.threadpool._run_process_async"):


could add the patch as decorator on the test function

tjeerddie · 2026-04-15T09:37:47Z

+
+
+def test_register_custom_serializer_registers_orchestrator_json():
+    with patch("orchestrator.services.tasks.registry") as mock_registry:


use patch as decorator, also update this in the other tests

tjeerddie · 2026-04-15T09:39:29Z

+    }
+
+
+def test_initialise_celery_registers_four_named_tasks():


redundant, you already check it in the previous test

tjeerddie · 2026-04-15T09:45:00Z

+def test_start_process_returns_none_on_exception(celery_start_fn, failing_fn):
+    process_id = uuid4()
+
+    patches = {


probably better to use variables instead of a dict or change the keys to just the function name.

tjeerddie · 2026-04-15T09:47:05Z

+
+
+# ---------------------------------------------------------------------------
+# CeleryJobWorkerStatus (existing fixture kept below)


the (existing fixture kept below) is useless in the comment, could have more context about the tests below instead.

Signed-off-by: Peter Boers <peter.boers@surf.nl>

github-actions · 2026-04-15T18:03:05Z

uv Lockfile Report

Added

Package	Version
psycopg	3.3.3
psycopg-binary	3.3.3

Removed

Package	Version
psycopg2-binary	2.9.11

estoel81 added 6 commits April 14, 2026 09:00

Remove redundant expire_on_commit manipulation around commit

7d2e21a

Remove stale expire_on_commit references after feature was removed

faccaac

Align search/core/types.py with main (remove stale type: ignore comment)

14b3772

Trim verbose psycopg3 comments to one-liners

c75be29

Trim verbose docstrings to one-liners

914127c

estoel81 enabled auto-merge (squash) April 14, 2026 09:06

Merge branch 'main' into psycopg3-clean

ec2e9ae

Signed-off-by: Erik Stoel <estoel81@hotmail.com>

estoel81 added 10 commits April 14, 2026 11:45

Fix type annotations in tasks.py and database URI validation tests

2ba7225

Add noqa comment and refactor docstrings and set comprehension

9e97078

Remove unused imports and rename unused parameter

0b25f4a

Add comprehensive unit tests for tasks.py covering Celery initializat…

edb7baa

…ion and task execution

Remove stale type: ignore comment from _noop_transactional

90a1655

Add unit tests for initialise_celery, start_process, and resume_process

767d509

Tests cover the psycopg3-critical transactional() wrapping in start_process and resume_process, plus get_celery_task, register_custom_serializer, and initialise_celery task route/registration behaviour.

Fix ltree filter bind parameters to use lquery type instead of TEXT

46dc1f4

Add unit tests for _LQuery type and verify CAST AS lquery in compiled…

ed8a0da

… SQL

Remove unused imports and apply black formatting to ltree filters

8e9b2c0

pboers1988 approved these changes Apr 14, 2026

View reviewed changes

Comment thread test/unit_tests/db/test_pool_events.py Outdated

Comment thread test/unit_tests/test_database_uri_validation.py Outdated

pboers1988 previously approved these changes Apr 14, 2026

View reviewed changes

estoel81 dismissed pboers1988’s stale review via 3bcf142 April 15, 2026 06:59

estoel81 added 4 commits April 15, 2026 09:02

Merge branch 'main' into psycopg3-clean

e748668

# Conflicts: # pyproject.toml # uv.lock

Document psycopg3 migration and new register_table() API in 5.0 upgra…

18991a2

…de guide

Remove docstring from ltree_filters unit test module

ff918a2

tjeerddie reviewed Apr 15, 2026

View reviewed changes

pboers1988 previously approved these changes Apr 15, 2026

View reviewed changes

Merge branch 'main' into psycopg3-clean

f94abe8

Signed-off-by: Peter Boers <peter.boers@surf.nl>

pboers1988 dismissed their stale review via f94abe8 April 15, 2026 18:02


		self.include_router(api_router, prefix="/api")

		# Validate DATABASE_URI dialect before initializing the database.

	orig = psycopg_errors.UndefinedTable.__new__(psycopg_errors.UndefinedTable)
	orig = psycopg_errors.UndefinedTable()



		@pytest.mark.usefixtures("_cleanup_extra_field")
		def test_register_table_columns_visible_in_inspect():



		def test_register_custom_serializer_registers_orchestrator_json():
		with patch("orchestrator.services.tasks.registry") as mock_registry:



		# ---------------------------------------------------------------------------
		# CeleryJobWorkerStatus (existing fixture kept below)

Conversation

estoel81 commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Plan

Uh oh!

sentry bot commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

codspeed-hq bot commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will not alter performance

Uh oh!

pboers1988 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

pboers1988 left a comment

Choose a reason for hiding this comment

Uh oh!

tjeerddie Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tjeerddie Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tjeerddie Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Apr 15, 2026

uv Lockfile Report

Added

Removed

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

estoel81 commented Apr 14, 2026 •

edited

Loading

sentry bot commented Apr 14, 2026 •

edited

Loading

codspeed-hq bot commented Apr 14, 2026 •

edited

Loading

tjeerddie Apr 15, 2026 •

edited

Loading

tjeerddie Apr 15, 2026 •

edited

Loading

tjeerddie Apr 15, 2026 •

edited

Loading