Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement by uesleilima · Pull Request #2366 · apache/age

uesleilima · 2026-04-05T10:27:50Z

Problem

setUpAge() and Age.connect() unconditionally execute LOAD 'age' (or LOAD '$libdir/plugins/age'). On managed PostgreSQL services like Azure Database for PostgreSQL, the AGE extension is loaded server-side via shared_preload_libraries and the binary is not accessible at the file path expected by LOAD, causing a psycopg.errors.UndefinedFile error. There is no way to use the driver in these environments without bypassing the setup entirely.

Solution

Add a skip_load parameter (default False) to setUpAge(), Age.connect(), and the module-level age.connect(). When True, the LOAD statement is skipped while all other setup is preserved:

SET search_path
agtype adapter registration
graph existence check / creation

This is fully backward-compatible — existing code is unaffected.

Usage

import age

ag = age.connect(
    dsn="host=myserver.postgres.database.azure.com ...",
    graph="my_graph",
    skip_load=True
)

…plugin loading. Fix for apache#2325

- Add 4 unit tests for setUpAge() skip_load behavior: skip_load=True skips LOAD, skip_load=False executes LOAD, load_from_plugins integration, and search_path always set. - Document skip_load in README under new "Managed PostgreSQL Usage" section for Azure/AWS RDS/etc. environments. - Fix syntax in existing load_from_plugins code example. Made-with: Cursor

tcolo

@uesleilima thanks a a lot for the PR - please addresse the issue sin the Claude Code Check.

Code Review

Summary

This PR addresses a real and valid problem — managed PostgreSQL environments (Azure, AWS RDS) where LOAD 'age' fails because the extension is pre-loaded via shared_preload_libraries. The implementation is backward-compatible and test coverage is solid.

However, there's a significant design conflict with the already-merged configure_connection() API that should be resolved before merging.

🔴 Critical

skip_load may leak into psycopg.connect() via **kwargs (age/__init__.py)

The current module-level connect() signature is:

def connect(dsn=None, graph=None, connection_factory=None, cursor_factory=ClientCursor, load_from_plugins=False, **kwargs):
    ag.connect(dsn=dsn, ..., load_from_plugins=load_from_plugins, **kwargs)

If skip_load is not explicitly named in this signature, it falls through **kwargs → Age.connect(**kwargs) → psycopg.connect(**kwargs), which will raise:

ProgrammingError: invalid dsn parameter 'skip_load'

skip_load must be extracted explicitly, the same way load_from_plugins is.

🟡 Suggestions

1. API duplication with configure_connection() (age/age.py — setUpAge())

The already-merged configure_connection(load=False) (see commits 7d64707e, 2a8a7c2d) solves the exact same problem — skipping LOAD 'age' on managed services — with a safer opt-in default. This PR introduces a parallel API for the same use case.

2. Inverted semantics between the two APIs

Function	Pattern	Safe default
configure_connection()	load=False (opt-in)	✅ skips LOAD by default
setUpAge() (this PR)	skip_load=False (opt-out)	⚠️ executes LOAD by default

A user reading both APIs won't know which to use for managed services, and the meaning of "safe default" is different for each. If both must coexist, the semantics should be aligned and the distinction documented clearly.

3. Undocumented interaction between skip_load=True and load_from_plugins=True

When both flags are set, LOAD is silently skipped — which is probably correct, but surprising. The docstring should document the precedence, or a ValueError should be raised for the contradictory combination.

4. Call chain not covered by tests

The unit tests mock setUpAge() in isolation, but the path age.connect(skip_load=True) → Age.connect() → setUpAge() isn't tested end-to-end. A test verifying the parameter is forwarded correctly through the full call chain would catch regressions like the **kwargs leak described above.

5. README should mention configure_connection() as an alternative

The new README section shows age.connect(skip_load=True) as the solution for managed environments. Since configure_connection() also exists and is better suited for pool-based setups, the docs should mention both and clarify the trade-offs.

✅ What Looks Good

Fully backward-compatible (skip_load=False default preserves existing behavior)
Search path and adapter registration are correctly preserved when skip_load=True
Unit tests use mocking cleanly — no live DB required
5 test methods give solid coverage of setUpAge() in isolation
The problem (UndefinedFile on Azure PostgreSQL) is well-documented in the PR description

Verdict: Request Changes

The **kwargs leak is a likely runtime error that needs to be addressed. Beyond that, the team should align on whether skip_load on setUpAge() is needed given configure_connection(load=False) already exists, or document a clear distinction between the two. Merging both without resolving this will leave users with two confusingly similar solutions with opposite flag polarities.

- Raise ValueError when skip_load=True and load_from_plugins=True are both set (contradictory combination) - Add end-to-end test verifying skip_load is forwarded through the full age.connect() → Age.connect() → setUpAge() call chain - Replace fragile string assertions with assert_called_with/assert_any_call - README: mention configure_connection() as the pool-based alternative for managed PostgreSQL environments Made-with: Cursor

Copilot

Pull request overview

Adds an opt-out for executing LOAD 'age' during Python driver connection/setup, enabling use on managed PostgreSQL services (e.g., Azure) where AGE is preloaded via shared_preload_libraries and LOAD fails.

Changes:

Add skip_load parameter to setUpAge(), Age.connect(), and module-level age.connect() to bypass LOAD while keeping the rest of setup.
Add unit tests validating skip_load behavior and argument forwarding.
Update Python driver README with managed-Postgres usage guidance.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
drivers/python/age/age.py	Introduces `skip_load` behavior in `setUpAge()` and forwards it through `Age.connect()`.
drivers/python/age/init.py	Exposes `skip_load` via module-level `connect()`.
drivers/python/test_age_py.py	Adds unit tests for `skip_load` behavior without requiring a DB connection.
drivers/python/README.md	Documents `skip_load` for managed PostgreSQL environments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

drivers/python/README.md

drivers/python/age/age.py

configure_connection() is not part of this PR; use the available setUpAge() API with skip_load=True for the connection pool example. Made-with: Cursor

uesleilima · 2026-04-06T20:42:49Z

All review feedback has been addressed:

tcolo's review:

✅ skip_load is already explicitly named in __init__.py:connect() and Age.connect() — no kwargs leak.
✅ Added ValueError for contradictory skip_load=True + load_from_plugins=True.
✅ Added end-to-end test verifying skip_load is forwarded through the full age.connect() → Age.connect() → setUpAge() call chain.
✅ README now mentions setUpAge(skip_load=True) for pool-based setups.
✅ Replaced fragile string assertions with assert_called_with / assert_any_call.

Copilot review:

✅ Fixed README: removed configure_connection() reference (not part of this branch); now uses setUpAge(conn, 'graph_name', skip_load=True) for the pool example.
⚠️ connection_factory being unused is a pre-existing issue, out of scope for this PR.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-10T18:19:37Z

drivers/python/age/age.py

+            else:
+                cursor.execute("LOAD 'age';")

        cursor.execute("SET search_path = ag_catalog, '$user', public;")


SET search_path uses '$user' (single quotes), which PostgreSQL treats as a literal schema name rather than the special $user placeholder. This likely prevents the current user's schema from being included in search_path and diverges from the documented setup elsewhere in the repo (which uses "$user"). Consider changing this statement to use "$user" so it behaves as intended, and update tests accordingly.

Suggested change

cursor.execute("SET search_path = ag_catalog, '$user', public;")

cursor.execute('SET search_path = ag_catalog, "$user", public;')

Addressed in 56d3df3: SET search_path now uses "$user" so PostgreSQL expands the session user schema. Thanks for catching this.

Copilot · 2026-04-10T18:19:38Z

drivers/python/test_age_py.py

+    def test_skip_load_true_does_not_execute_load(self):
+        """When skip_load=True, LOAD 'age' must not be executed."""
+        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
+        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
+             unittest.mock.patch("age.age.checkGraphCreated"):
+            age.age.setUpAge(mock_conn, "test_graph", skip_load=True)
+        mock_cursor.execute.assert_called_once_with(
+            "SET search_path = ag_catalog, '$user', public;"
+        )


These new tests assert the driver sets SET search_path = ag_catalog, '$user', public;, but PostgreSQL's special $user placeholder requires double quotes ("$user") to be expanded. As written, the tests lock in the likely-buggy single-quoted form; if the driver is corrected to "$user", update these assertions to match.

Updated test expectations to match the corrected "$user" form (same commit as the driver fix).

Copilot · 2026-04-10T18:19:38Z

drivers/python/test_age_py.py

+class TestSetUpAge(unittest.TestCase):
+    """Unit tests for setUpAge() skip_load parameter — no DB required."""
+
+    def _make_mock_conn(self):
+        mock_conn = unittest.mock.MagicMock()
+        mock_cursor = unittest.mock.MagicMock()
+        mock_conn.cursor.return_value.__enter__ = unittest.mock.Mock(return_value=mock_cursor)
+        mock_conn.cursor.return_value.__exit__ = unittest.mock.Mock(return_value=False)
+        mock_conn.adapters = unittest.mock.MagicMock()


The CI workflow runs python test_age_py.py ..., which executes only the custom TestSuite built under if __name__ == '__main__' and does not include TestSetUpAge. As a result, these unit tests for skip_load won’t run in CI. Consider adding TestSetUpAge to the scripted suite, or switching the script to unittest.main() (and passing through args) so discovery runs all tests.

TestSetUpAge is now included in the __main__ suite via TestLoader().loadTestsFromTestCase(TestSetUpAge), so CI’s python test_age_py.py … run executes the skip_load tests.

jrgemignani · 2026-04-10T18:21:38Z

@uesleilima Please see above Copilot comments

…Age in CI PostgreSQL treats single-quoted '$user' as a literal schema name; use "$user" so the session user schema is included. Include TestSetUpAge in the test_age_py __main__ suite so skip_load tests run in GitHub Actions. Made-with: Cursor

Add skip_load parameter to connect and setUpAge functions to control …

12ac331

…plugin loading. Fix for apache#2325

uesleilima mentioned this pull request Apr 5, 2026

Python driver: Age.connect() / setUpAge() unconditionally runs LOAD 'age', failing on managed PostgreSQL (Azure) where extension is pre-loaded via shared_preload_libraries #2353

Open

tcolo suggested changes Apr 5, 2026

View reviewed changes

uesleilima requested a review from tcolo April 5, 2026 23:01

MuhammadTahaNaveed requested review from Copilot and removed request for tcolo April 6, 2026 10:30

Copilot started reviewing on behalf of MuhammadTahaNaveed April 6, 2026 10:30 View session

Copilot AI reviewed Apr 6, 2026

View reviewed changes

drivers/python/README.md Outdated Show resolved Hide resolved

drivers/python/age/age.py Show resolved Hide resolved

Fix README: use setUpAge(skip_load=True) for pool example

3ce6bfe

configure_connection() is not part of this PR; use the available setUpAge() API with skip_load=True for the connection pool example. Made-with: Cursor

uesleilima requested a review from tcolo April 6, 2026 20:45

jrgemignani requested a review from Copilot April 10, 2026 18:15

Copilot started reviewing on behalf of jrgemignani April 10, 2026 18:16 View session

Copilot AI reviewed Apr 10, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement#2366

Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement#2366
uesleilima wants to merge 5 commits intoapache:masterfrom
uesleilima:master

uesleilima commented Apr 5, 2026

Uh oh!

tcolo left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

uesleilima commented Apr 6, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

uesleilima Apr 10, 2026

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

uesleilima Apr 10, 2026

Uh oh!

Copilot AI Apr 10, 2026

Uh oh!

uesleilima Apr 10, 2026

Uh oh!

jrgemignani commented Apr 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	cursor.execute("SET search_path = ag_catalog, '$user', public;")
	cursor.execute('SET search_path = ag_catalog, "$user", public;')

Conversation

uesleilima commented Apr 5, 2026

Problem

Solution

Usage

Uh oh!

tcolo left a comment

Choose a reason for hiding this comment

Code Review

Summary

🔴 Critical

🟡 Suggestions

✅ What Looks Good

Verdict: Request Changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

uesleilima commented Apr 6, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

uesleilima Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

uesleilima Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

uesleilima Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

jrgemignani commented Apr 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants