uesleilima commented on code in PR #2374:
URL: https://github.com/apache/age/pull/2374#discussion_r3041587173
##########
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:
Fixed — `Path.to_dict()` now preserves JSON-native types (dict, list, str,
int, float, bool, None) unchanged; `str()` is only used as a safety net for
non-serializable objects. Also handles `self.entities is None` safely.
##########
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:
Valid observation. Refactoring `setUpAge()` to delegate to
`configure_connection()` (or vice versa) would reduce duplication, but it
changes the behavior semantics of `setUpAge` which is the existing public API.
Better addressed in a follow-up PR to avoid scope creep.
##########
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:
This is intentional — `configure_connection(load=False)` uses positive
opt-in semantics (`load=True` to enable loading), unlike
`setUpAge(skip_load=True)` which uses opt-out. The `load` parameter was chosen
because the default use case for `configure_connection` is connection pooling
where AGE is already loaded. The docstring documents this clearly.
##########
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:
Fixed — added `TestModelToDict`, `TestPublicImports`, and
`TestConfigureConnection` to the `__main__` suite via
`unittest.TestLoader().loadTestsFromTestCase()`. They now run both via `python
-m unittest` discovery and direct script execution.
--
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]
