Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions drivers/python/age/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import psycopg.conninfo as conninfo
from . import age
from .age import *
from .age import AgeLoader, ClientCursor, configure_connection
from .models import *
from .builder import ResultHandler, DummyResultHandler, parseAgeValue, newResultHandler
from . import VERSION
Expand Down
67 changes: 67 additions & 0 deletions drivers/python/age/age.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
# under the License.

import re
from typing import Any, Optional

import psycopg
from psycopg.types import TypeInfo
from psycopg import sql
Expand Down Expand Up @@ -170,6 +172,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: Optional[str] = None,
load: bool = False,
load_from_plugins: bool = False,
Comment on lines +176 to +180
Copy link

Copilot AI Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

configure_connection() uses PEP 604 union syntax (str | None) in its public signature. This is a syntax error on Python 3.9, but the package metadata declares requires-python = ">=3.9" (drivers/python/pyproject.toml:25). Either adjust the type hints to be 3.9-compatible (e.g., Optional[str] / Union[...]) or bump the declared minimum Python version to 3.10+ so users on 3.9 don’t break at import time.

Copilot uses AI. Check for mistakes.
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1d92c31: graph_name is now Optional[str] with from typing import Optional so the module parses on Python 3.9 while keeping requires-python >= 3.9. Also aligned SET search_path with "$user" and updated tests.

) -> None:
Comment thread
jrgemignani marked this conversation as resolved.
"""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)

Comment thread
jrgemignani marked this conversation as resolved.

# Create the graph, if it does not exist
def checkGraphCreated(conn:psycopg.connection, graphName:str):
validate_graph_name(graphName)
Expand Down
40 changes: 40 additions & 0 deletions drivers/python/age/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,20 @@ def toJson(self) -> str:

return buf.getvalue()

def to_dict(self) -> list:
# AGObj elements are recursively converted; JSON-native types
# (dict, list, str, int, float, bool, None) pass through unchanged.
# Non-serializable objects fall back to str() as a safety net.
result = []
for e in (self.entities or []):
if isinstance(e, AGObj):
result.append(e.to_dict())
elif isinstance(e, (dict, list, str, int, float, bool, type(None))):
result.append(e)
else:
result.append(str(e))
return result




Expand Down Expand Up @@ -146,6 +160,18 @@ def __str__(self) -> str:
def __repr__(self) -> str:
return self.toString()

def to_dict(self) -> dict:
"""Return a plain dict suitable for JSON serialization.

Properties are shallow-copied; nested mutable values will share
references with the original Vertex.
"""
return {
"id": self.id,
"label": self.label,
"properties": dict(self.properties) if self.properties else {},
}

def toString(self) -> str:
return nodeToString(self)

Expand Down Expand Up @@ -186,6 +212,20 @@ def __str__(self) -> str:
def __repr__(self) -> str:
return self.toString()

def to_dict(self) -> dict:
"""Return a plain dict suitable for JSON serialization.

Properties are shallow-copied; nested mutable values will share
references with the original Edge.
"""
return {
"id": self.id,
"label": self.label,
"start_id": self.start_id,
"end_id": self.end_id,
"properties": dict(self.properties) if self.properties else {},
}

def extraStrFormat(node, buf):
if node.start_id != None:
buf.write(", start_id:")
Expand Down
Loading