Copilot commented on code in PR #2374:
URL: https://github.com/apache/age/pull/2374#discussion_r3039036779
##########
drivers/python/age/models.py:
##########
@@ -118,6 +118,14 @@ def toJson(self) -> str:
return buf.getvalue()
+ def to_dict(self) -> list:
+ # Non-AGObj elements (e.g. raw dicts/strings from malformed paths)
+ # are included as-is via str() to guarantee JSON-serializable output.
+ return [
+ e.to_dict() if isinstance(e, AGObj) else str(e)
+ for e in self.entities
+ ]
Review Comment:
`Path.to_dict()` currently coerces any non-`AGObj` element with `str(e)`. If
a path contains a plain `dict`/`list` (or other JSON-native types), this
converts structured data into a string representation (e.g., dicts become
single-quoted strings), which is lossy and not equivalent to JSON
serialization. Consider leaving JSON-native types unchanged and only
stringifying truly non-serializable objects (or recursively normalizing
lists/dicts). Also handle `self.entities is None` (current default ctor) to
avoid a `TypeError` when iterating.
##########
drivers/python/age/age.py:
##########
@@ -162,6 +162,71 @@ def setUpAge(conn:psycopg.connection, graphName:str,
load_from_plugins:bool=Fals
if graphName != None:
checkGraphCreated(conn, graphName)
+
+def configure_connection(
+ conn: psycopg.connection,
+ graph_name: str | None = None,
+ load: bool = False,
+ load_from_plugins: bool = False,
+) -> None:
+ """Register AGE agtype adapters on an existing connection.
+
+ This enables use of AGE with externally-managed connections, such as
+ those from psycopg_pool.ConnectionPool. By default the function does
+ **not** execute ``LOAD 'age'``, making it safe for managed PostgreSQL
+ services (Azure, AWS RDS) where the extension is pre-loaded via
+ ``shared_preload_libraries``.
+
+ Performs:
+ - ``SET search_path`` to include ``ag_catalog``
+ - Fetches agtype OIDs and registers ``AgeLoader``
+ - Optionally loads the AGE extension (``load=True``)
+ - Optionally checks/creates the graph
+
+ Args:
+ conn: An existing psycopg connection.
+ graph_name: Optional graph name to check/create.
+ load: If True, execute ``LOAD 'age'`` (or the plugins path).
+ Default False — suitable for environments where AGE is
+ already loaded.
+ load_from_plugins: If True (and ``load=True``), use
+ ``LOAD '$libdir/plugins/age'`` instead of ``LOAD 'age'``.
+
+ Raises:
+ ValueError: If ``load_from_plugins=True`` but ``load=False``.
+ AgeNotSet: If the agtype type is not found in the database.
+ """
+ if load_from_plugins and not load:
+ raise ValueError(
+ "load_from_plugins=True requires load=True. "
+ "Set load=True to enable extension loading."
+ )
+
+ with conn.cursor() as cursor:
+ if 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;")
+
+ ag_info = TypeInfo.fetch(conn, 'agtype')
+
+ if not ag_info:
+ raise AgeNotSet(
+ "AGE agtype type not found. Ensure the AGE extension is "
+ "installed and loaded in the current database. "
+ "Run CREATE EXTENSION age; first."
+ )
+
+ conn.adapters.register_loader(ag_info.oid, AgeLoader)
+ conn.adapters.register_loader(ag_info.array_oid, AgeLoader)
+
+ if graph_name is not None:
+ checkGraphCreated(conn, graph_name)
+
Review Comment:
`configure_connection()` largely duplicates `setUpAge()`
(LOAD/search_path/agtype TypeInfo fetch/loader registration/graph check). This
duplication risks the two code paths drifting over time (e.g., changes to error
handling or adapter registration). Consider refactoring so `setUpAge()` calls
`configure_connection(load=True, ...)` (or vice versa) and share a single
implementation.
##########
drivers/python/age/age.py:
##########
@@ -162,6 +162,71 @@ def setUpAge(conn:psycopg.connection, graphName:str,
load_from_plugins:bool=Fals
if graphName != None:
checkGraphCreated(conn, graphName)
+
+def configure_connection(
+ conn: psycopg.connection,
+ graph_name: str | None = None,
+ load: bool = False,
+ load_from_plugins: bool = False,
+) -> None:
Review Comment:
PR description and issue summary mention a `skip_load` parameter, but the
implemented public API uses `load: bool = False` instead. Please align the PR
description/docs with the actual signature (or rename the parameter) to avoid
confusion for users upgrading based on the PR notes.
##########
drivers/python/test_age_py.py:
##########
@@ -28,6 +29,160 @@
TEST_GRAPH_NAME = "test_graph"
+class TestModelToDict(unittest.TestCase):
+ """Unit tests for Vertex/Edge/Path to_dict() — no DB required."""
+
+ def test_vertex_to_dict(self):
+ v = Vertex(id=123, label="Person", properties={"name": "Alice", "age":
30})
+ d = v.to_dict()
+ self.assertEqual(d["id"], 123)
+ self.assertEqual(d["label"], "Person")
+ self.assertEqual(d["properties"], {"name": "Alice", "age": 30})
+ # Verify it's a plain dict (JSON-serializable)
+ json_str = json.dumps(d)
+ self.assertIn("Alice", json_str)
+
+ def test_vertex_to_dict_empty_properties(self):
+ v = Vertex(id=1, label="Empty", properties=None)
+ d = v.to_dict()
+ self.assertEqual(d["properties"], {})
+
+ def test_edge_to_dict(self):
+ e = Edge(id=456, label="KNOWS", properties={"since": 2020})
+ e.start_id = 123
+ e.end_id = 789
+ d = e.to_dict()
+ self.assertEqual(d["id"], 456)
+ self.assertEqual(d["label"], "KNOWS")
+ self.assertEqual(d["start_id"], 123)
+ self.assertEqual(d["end_id"], 789)
+ self.assertEqual(d["properties"], {"since": 2020})
+ json_str = json.dumps(d)
+ self.assertIn("KNOWS", json_str)
+
+ def test_path_to_dict(self):
+ v1 = Vertex(id=1, label="A", properties={"name": "start"})
+ e = Edge(id=10, label="r", properties={"w": 1})
+ e.start_id = 1
+ e.end_id = 2
+ v2 = Vertex(id=2, label="B", properties={"name": "end"})
+ p = Path([v1, e, v2])
+ d = p.to_dict()
+ self.assertEqual(len(d), 3)
+ self.assertEqual(d[0]["label"], "A")
+ self.assertEqual(d[1]["label"], "r")
+ self.assertEqual(d[1]["start_id"], 1)
+ self.assertEqual(d[2]["label"], "B")
+ # Verify the whole path is JSON-serializable
+ json_str = json.dumps(d)
+ self.assertIn("start", json_str)
+
+ def test_vertex_to_dict_is_plain_dict(self):
+ """to_dict() returns standard dict, not a model object."""
+ v = Vertex(id=1, label="X", properties={"k": "v"})
+ d = v.to_dict()
+ self.assertIsInstance(d, dict)
+ self.assertIsInstance(d["properties"], dict)
+
+
+class TestPublicImports(unittest.TestCase):
+ """Verify that public API symbols are importable without type: ignore."""
+
+ def test_import_configure_connection(self):
+ from age import configure_connection
+ self.assertTrue(callable(configure_connection))
+
+ def test_import_age_loader(self):
+ from age import AgeLoader
+ self.assertIsNotNone(AgeLoader)
+
+ def test_import_client_cursor(self):
+ from age import ClientCursor
+ self.assertIsNotNone(ClientCursor)
+
+
+class TestConfigureConnection(unittest.TestCase):
+ """Unit tests for configure_connection() — no DB required."""
+
Review Comment:
The new unit test classes (`TestModelToDict`, `TestPublicImports`,
`TestConfigureConnection`) won’t run in CI when executing `python
test_age_py.py ...` because the `__main__` block constructs a `TestSuite` that
only includes `TestAgeBasic` tests. Consider adding these new tests to the
explicit suite, or switching this file to `unittest.main()` (and moving the
DB/integration tests behind an opt-in flag) so the new tests are actually
executed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
