You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@age.apache.org by mt...@apache.org on 2026/04/19 12:10:44 UTC
(age) branch master updated: Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement (#2366)
This is an automated email from the ASF dual-hosted git repository.
MuhammadTahaNaveed pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/age.git
The following commit(s) were added to refs/heads/master by this push:
new ce500459 Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement (#2366)
ce500459 is described below
commit ce5004593ea5acb2cec3a2dc94a06241d7f23c17
Author: Ueslei Lima <ue...@gmail.com>
AuthorDate: Sun Apr 19 09:10:39 2026 -0300
Python driver: Add `skip_load` parameter to skip `LOAD 'age'` statement (#2366)
* Add skip_load parameter to connect and setUpAge functions to control plugin loading.
* 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
* Address review feedback: ValueError for contradictory flags, e2e test
- 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
* 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
* fix(python-driver): use quoted $user in search_path and run TestSetUpAge 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
---
drivers/python/README.md | 20 ++++++++++--
drivers/python/age/__init__.py | 4 +--
drivers/python/age/age.py | 24 +++++++++-----
drivers/python/test_age_py.py | 73 ++++++++++++++++++++++++++++++++++++++++++
4 files changed, 109 insertions(+), 12 deletions(-)
diff --git a/drivers/python/README.md b/drivers/python/README.md
index e64f9de6..f4fa4391 100644
--- a/drivers/python/README.md
+++ b/drivers/python/README.md
@@ -89,9 +89,25 @@ 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')
+ ```
+* **Connection pools:** If you manage connections externally (e.g. via ```psycopg_pool.ConnectionPool```),
+ you can call ```setUpAge()``` with ```skip_load=True``` on each pooled connection:
+ ```python
+ from age.age import setUpAge
+ setUpAge(conn, 'graph_name', skip_load=True)
```
### License
diff --git a/drivers/python/age/__init__.py b/drivers/python/age/__init__.py
index 685f0fe7..caee6a43 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 fad1f27b..f85b1ab1 100644
--- a/drivers/python/age/age.py
+++ b/drivers/python/age/age.py
@@ -137,14 +137,22 @@ class AgeLoader(psycopg.adapt.Loader):
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):
+ 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 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;")
+ cursor.execute('SET search_path = ag_catalog, "$user", public;')
ag_info = TypeInfo.fetch(conn, 'agtype')
@@ -333,9 +341,9 @@ class Age:
# 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
diff --git a/drivers/python/test_age_py.py b/drivers/python/test_age_py.py
index f904fb9e..960d28d4 100644
--- a/drivers/python/test_age_py.py
+++ b/drivers/python/test_age_py.py
@@ -16,6 +16,7 @@ import json
from age.models import Vertex
import unittest
+import unittest.mock
import decimal
import age
import argparse
@@ -28,6 +29,76 @@ TEST_PASSWORD = "agens"
TEST_GRAPH_NAME = "test_graph"
+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()
+ 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)
+ 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, 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)
+ 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, 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)
+ 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, 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;'
+ )
+
+ 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):
ag = None
args: argparse.Namespace = argparse.Namespace(
@@ -485,6 +556,8 @@ if __name__ == "__main__":
args = parser.parse_args()
suite = unittest.TestSuite()
+ loader = unittest.TestLoader()
+ suite.addTests(loader.loadTestsFromTestCase(TestSetUpAge))
suite.addTest(TestAgeBasic("testExec"))
suite.addTest(TestAgeBasic("testQuery"))
suite.addTest(TestAgeBasic("testChangeData"))