From 12ac331aa4a9c5dd09c91b1e6d6bbd40d6ef8a6d Mon Sep 17 00:00:00 2001 From: Ueslei Santos Lima Date: Sun, 5 Apr 2026 12:24:06 +0200 Subject: [PATCH 1/4] Add skip_load parameter to connect and setUpAge functions to control plugin loading. Fix for #2325 --- drivers/python/age/__init__.py | 4 ++-- drivers/python/age/age.py | 15 ++++++++------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/drivers/python/age/__init__.py b/drivers/python/age/__init__.py index 685f0fe74..caee6a43c 100644 --- a/drivers/python/age/__init__.py +++ b/drivers/python/age/__init__.py @@ -25,13 +25,13 @@ def version(): def connect(dsn=None, graph=None, connection_factory=None, cursor_factory=ClientCursor, load_from_plugins=False, - **kwargs): + skip_load=False, **kwargs): dsn = conninfo.make_conninfo('' if dsn is None else dsn, **kwargs) ag = Age() ag.connect(dsn=dsn, graph=graph, connection_factory=connection_factory, cursor_factory=cursor_factory, - load_from_plugins=load_from_plugins, **kwargs) + load_from_plugins=load_from_plugins, skip_load=skip_load, **kwargs) return ag # Dummy ResultHandler diff --git a/drivers/python/age/age.py b/drivers/python/age/age.py index fad1f27b1..011dff285 100644 --- a/drivers/python/age/age.py +++ b/drivers/python/age/age.py @@ -137,12 +137,13 @@ def load(self, data: bytes | bytearray | memoryview) -> Any | None: return parseAgeValue(data_bytes.decode('utf-8')) -def setUpAge(conn:psycopg.connection, graphName:str, load_from_plugins:bool=False): +def setUpAge(conn:psycopg.connection, graphName:str, load_from_plugins:bool=False, skip_load:bool=False): with conn.cursor() as cursor: - if load_from_plugins: - cursor.execute("LOAD '$libdir/plugins/age';") - else: - cursor.execute("LOAD 'age';") + if not skip_load: + if load_from_plugins: + cursor.execute("LOAD '$libdir/plugins/age';") + else: + cursor.execute("LOAD 'age';") cursor.execute("SET search_path = ag_catalog, '$user', public;") @@ -333,9 +334,9 @@ def __init__(self): # Connect to PostgreSQL Server and establish session and type extension environment. def connect(self, graph:str=None, dsn:str=None, connection_factory=None, cursor_factory=ClientCursor, - load_from_plugins:bool=False, **kwargs): + load_from_plugins:bool=False, skip_load:bool=False, **kwargs): conn = psycopg.connect(dsn, cursor_factory=cursor_factory, **kwargs) - setUpAge(conn, graph, load_from_plugins) + setUpAge(conn, graph, load_from_plugins, skip_load=skip_load) self.connection = conn self.graphName = graph return self From 8be3cbb19ee0dc042b56c19220af1718f3006752 Mon Sep 17 00:00:00 2001 From: Ueslei Santos Lima Date: Sun, 5 Apr 2026 15:16:12 +0200 Subject: [PATCH 2/4] Add tests and README docs for skip_load parameter - 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 --- drivers/python/README.md | 14 +++++- drivers/python/test_age_py.py | 85 +++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/drivers/python/README.md b/drivers/python/README.md index e64f9de67..deadfd116 100644 --- a/drivers/python/README.md +++ b/drivers/python/README.md @@ -89,9 +89,19 @@ SET search_path = ag_catalog, "$user", public; * Make sure to give your non-superuser db account proper permissions to the graph schemas and corresponding objects * Make sure to initiate the Apache Age python driver with the ```load_from_plugins``` parameter. This parameter tries to load the Apache Age extension from the PostgreSQL plugins directory located at ```$libdir/plugins/age```. Example: - ```python. + ```python ag = age.connect(host='localhost', port=5432, user='dbuser', password='strong_password', - dbname=postgres, load_from_plugins=True, graph='graph_name) + dbname='postgres', load_from_plugins=True, graph='graph_name') + ``` + +### Managed PostgreSQL Usage (Azure, AWS RDS, etc.) +* On managed PostgreSQL services where the AGE extension is loaded server-side via ```shared_preload_libraries```, + the ```LOAD 'age'``` command may fail because the binary is not at the expected file path. Use the ```skip_load``` + parameter to skip the ```LOAD``` statement while still performing all other setup: + ```python + ag = age.connect(host='myserver.postgres.database.azure.com', port=5432, + user='dbuser', password='strong_password', + dbname='postgres', skip_load=True, graph='graph_name') ``` ### License diff --git a/drivers/python/test_age_py.py b/drivers/python/test_age_py.py index f904fb9e3..81dd53ec1 100644 --- a/drivers/python/test_age_py.py +++ b/drivers/python/test_age_py.py @@ -16,6 +16,7 @@ from age.models import Vertex import unittest +import unittest.mock import decimal import age import argparse @@ -28,6 +29,90 @@ TEST_GRAPH_NAME = "test_graph" +class TestSetUpAge(unittest.TestCase): + """Unit tests for setUpAge() skip_load parameter — no DB required.""" + + def test_skip_load_true_does_not_execute_load(self): + """When skip_load=True, LOAD 'age' must not be executed.""" + 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() + + mock_type_info = unittest.mock.MagicMock() + mock_type_info.oid = 1 + mock_type_info.array_oid = 2 + + 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) + + executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] + for stmt in executed_stmts: + self.assertNotIn("LOAD", stmt, f"LOAD should not be called when skip_load=True, got: {stmt}") + + def test_skip_load_false_executes_load(self): + """When skip_load=False (default), LOAD 'age' must be executed.""" + 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() + + mock_type_info = unittest.mock.MagicMock() + mock_type_info.oid = 1 + mock_type_info.array_oid = 2 + + 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=False) + + executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] + load_calls = [s for s in executed_stmts if "LOAD" in s] + self.assertTrue(len(load_calls) > 0, "LOAD should be called when skip_load=False") + + def test_skip_load_with_load_from_plugins(self): + """When skip_load=False and load_from_plugins=True, LOAD from plugins path.""" + 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() + + mock_type_info = unittest.mock.MagicMock() + mock_type_info.oid = 1 + mock_type_info.array_oid = 2 + + 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", load_from_plugins=True, skip_load=False) + + executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] + plugins_calls = [s for s in executed_stmts if "plugins" in s] + self.assertTrue(len(plugins_calls) > 0, "LOAD from plugins path should be called") + + def test_skip_load_true_still_sets_search_path(self): + """When skip_load=True, search_path must still be set.""" + 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() + + mock_type_info = unittest.mock.MagicMock() + mock_type_info.oid = 1 + mock_type_info.array_oid = 2 + + 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) + + executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] + search_path_calls = [s for s in executed_stmts if "search_path" in s] + self.assertTrue(len(search_path_calls) > 0, "search_path should be set even when skip_load=True") + + class TestAgeBasic(unittest.TestCase): ag = None args: argparse.Namespace = argparse.Namespace( From 56256f922fc774e050edf61c6edb4172e4e2e71d Mon Sep 17 00:00:00 2001 From: Ueslei Santos Lima Date: Mon, 6 Apr 2026 00:57:52 +0200 Subject: [PATCH 3/4] Address review feedback: ValueError for contradictory flags, e2e test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- drivers/python/README.md | 7 +++ drivers/python/age/age.py | 7 +++ drivers/python/test_age_py.py | 82 +++++++++++++++-------------------- 3 files changed, 48 insertions(+), 48 deletions(-) diff --git a/drivers/python/README.md b/drivers/python/README.md index deadfd116..bf7bd1c17 100644 --- a/drivers/python/README.md +++ b/drivers/python/README.md @@ -103,6 +103,13 @@ SET search_path = ag_catalog, "$user", public; user='dbuser', password='strong_password', dbname='postgres', skip_load=True, graph='graph_name') ``` +* **Connection pools:** If you manage connections externally (e.g. via ```psycopg_pool.ConnectionPool```), + use ```configure_connection()``` instead, which skips ```LOAD 'age'``` by default and is designed for + per-connection setup callbacks: + ```python + from age import configure_connection + configure_connection(conn, graph_name='graph_name') + ``` ### License Apache-2.0 License diff --git a/drivers/python/age/age.py b/drivers/python/age/age.py index 011dff285..5feba1bd7 100644 --- a/drivers/python/age/age.py +++ b/drivers/python/age/age.py @@ -138,6 +138,13 @@ def load(self, data: bytes | bytearray | memoryview) -> Any | None: def setUpAge(conn:psycopg.connection, graphName:str, load_from_plugins:bool=False, skip_load:bool=False): + if skip_load and load_from_plugins: + raise ValueError( + "skip_load=True and load_from_plugins=True are contradictory. " + "Set skip_load=False to load the extension from the plugins path, " + "or remove load_from_plugins to skip loading entirely." + ) + with conn.cursor() as cursor: if not skip_load: if load_from_plugins: diff --git a/drivers/python/test_age_py.py b/drivers/python/test_age_py.py index 81dd53ec1..fe4b91d1b 100644 --- a/drivers/python/test_age_py.py +++ b/drivers/python/test_age_py.py @@ -32,85 +32,71 @@ class TestSetUpAge(unittest.TestCase): """Unit tests for setUpAge() skip_load parameter — no DB required.""" - def test_skip_load_true_does_not_execute_load(self): - """When skip_load=True, LOAD 'age' must not be executed.""" + 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() - mock_type_info = unittest.mock.MagicMock() mock_type_info.oid = 1 mock_type_info.array_oid = 2 + return mock_conn, mock_cursor, mock_type_info + 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) - - executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] - for stmt in executed_stmts: - self.assertNotIn("LOAD", stmt, f"LOAD should not be called when skip_load=True, got: {stmt}") + mock_cursor.execute.assert_called_once_with( + "SET search_path = ag_catalog, '$user', public;" + ) def test_skip_load_false_executes_load(self): """When skip_load=False (default), LOAD 'age' must be executed.""" - 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() - - mock_type_info = unittest.mock.MagicMock() - mock_type_info.oid = 1 - mock_type_info.array_oid = 2 - + 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=False) - - executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] - load_calls = [s for s in executed_stmts if "LOAD" in s] - self.assertTrue(len(load_calls) > 0, "LOAD should be called when skip_load=False") + mock_cursor.execute.assert_any_call("LOAD 'age';") def test_skip_load_with_load_from_plugins(self): """When skip_load=False and load_from_plugins=True, LOAD from plugins path.""" - 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() - - mock_type_info = unittest.mock.MagicMock() - mock_type_info.oid = 1 - mock_type_info.array_oid = 2 - + 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", load_from_plugins=True, skip_load=False) - - executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] - plugins_calls = [s for s in executed_stmts if "plugins" in s] - self.assertTrue(len(plugins_calls) > 0, "LOAD from plugins path should be called") + mock_cursor.execute.assert_any_call("LOAD '$libdir/plugins/age';") def test_skip_load_true_still_sets_search_path(self): """When skip_load=True, search_path must still be set.""" - 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() - - mock_type_info = unittest.mock.MagicMock() - mock_type_info.oid = 1 - mock_type_info.array_oid = 2 - + 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_any_call( + "SET search_path = ag_catalog, '$user', public;" + ) - executed_stmts = [str(call) for call in mock_cursor.execute.call_args_list] - search_path_calls = [s for s in executed_stmts if "search_path" in s] - self.assertTrue(len(search_path_calls) > 0, "search_path should be set even when skip_load=True") + def test_contradictory_skip_load_and_load_from_plugins_raises(self): + """skip_load=True + load_from_plugins=True must raise ValueError.""" + mock_conn, _, _ = self._make_mock_conn() + with self.assertRaises(ValueError): + age.age.setUpAge(mock_conn, "test_graph", load_from_plugins=True, skip_load=True) + + def test_connect_forwards_skip_load_to_setup(self): + """age.connect(skip_load=True) must forward skip_load through the full call chain.""" + with unittest.mock.patch("age.age.psycopg.connect") as mock_psycopg, \ + unittest.mock.patch("age.age.setUpAge") as mock_setup: + mock_psycopg.return_value = unittest.mock.MagicMock() + age.connect(dsn="host=localhost", graph="test_graph", skip_load=True) + mock_setup.assert_called_once() + _, kwargs = mock_setup.call_args + self.assertTrue( + kwargs.get("skip_load", False), + "skip_load must be forwarded from age.connect() to setUpAge()" + ) class TestAgeBasic(unittest.TestCase): From 3ce6bfed674a2571f12ddd0a2a5fac27be409db1 Mon Sep 17 00:00:00 2001 From: Ueslei Santos Lima Date: Mon, 6 Apr 2026 22:36:56 +0200 Subject: [PATCH 4/4] Fix README: use setUpAge(skip_load=True) for pool example 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 --- drivers/python/README.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/drivers/python/README.md b/drivers/python/README.md index bf7bd1c17..f4fa43919 100644 --- a/drivers/python/README.md +++ b/drivers/python/README.md @@ -104,11 +104,10 @@ SET search_path = ag_catalog, "$user", public; dbname='postgres', skip_load=True, graph='graph_name') ``` * **Connection pools:** If you manage connections externally (e.g. via ```psycopg_pool.ConnectionPool```), - use ```configure_connection()``` instead, which skips ```LOAD 'age'``` by default and is designed for - per-connection setup callbacks: + you can call ```setUpAge()``` with ```skip_load=True``` on each pooled connection: ```python - from age import configure_connection - configure_connection(conn, graph_name='graph_name') + from age.age import setUpAge + setUpAge(conn, 'graph_name', skip_load=True) ``` ### License