Python driver: Age.connect() / setUpAge() unconditionally runs LOAD 'age', failing on managed PostgreSQL (Azure) where extension is pre-loaded via shared_preload_libraries

[uesleilima]

Description

The setUpAge() function unconditionally executes LOAD 'age' (or LOAD '$libdir/plugins/age'). The Age.connect() method always calls setUpAge() with no option to skip the LOAD statement.

On managed PostgreSQL services like Azure Database for PostgreSQL, the AGE extension is loaded server-side via shared_preload_libraries. The extension binary may not be accessible at the file path expected by LOAD, causing a psycopg.errors.UndefinedFile error. Users have no way to use Age.connect() in these environments without bypassing the driver's setup entirely.

This is related to #839, which discussed non-superuser access. However, our scenario is different: on Azure managed PostgreSQL the LOAD command fails regardless of privileges because the binary is not at the expected file path. The extension is already loaded and functional — the driver just needs a way to skip the LOAD statement.

Steps to Reproduce

1. Set up Azure Database for PostgreSQL (Flexible Server) with AGE extension enabled via shared_preload_libraries
2. Create the extension: CREATE EXTENSION IF NOT EXISTS age;
3. Attempt to connect using the Python driver:

import age

ag = age.connect(
    dsn="host=myserver.postgres.database.azure.com port=5432 dbname=mydb user=myuser password=mypass",
    graph="my_graph"
)

Expected Behavior

The connection should succeed since AGE is already loaded and available in the database.

Actual Behavior

With load_from_plugins=False (the default):

psycopg.errors.UndefinedFile: could not access file "age": No such file or directory

With load_from_plugins=True:

psycopg.errors.UndefinedFile: could not access file "$libdir/plugins/age": No such file or directory

Neither option works because the binary path is managed by the cloud provider and not directly accessible.

Root Cause

In age.py setUpAge():

def setUpAge(conn, graphName, load_from_plugins=False):
    with conn.cursor() as cursor:
        if load_from_plugins:
            cursor.execute("LOAD '$libdir/plugins/age';")
        else:
            cursor.execute("LOAD 'age';")
        # ... rest of setup (search_path, adapters, graph check)

And in Age.connect():

def connect(self, graph=None, dsn=None, ..., load_from_plugins=False, **kwargs):
    conn = psycopg.connect(dsn, cursor_factory=cursor_factory, **kwargs)
    setUpAge(conn, graph, load_from_plugins)  # always called, no way to skip LOAD

The LOAD statement is only needed when the AGE shared library has not been pre-loaded into the PostgreSQL backend. On managed services, the library is already loaded via shared_preload_libraries, making the LOAD statement unnecessary and harmful.

Suggested Fix

Add a skip_load parameter to both setUpAge() and Age.connect() that skips the LOAD statement while still performing all other setup (search_path, agtype adapter registration, graph creation):

def setUpAge(conn, graphName, load_from_plugins=False, skip_load=False):
    with conn.cursor() as cursor:
        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;")
        # ... rest of adapter registration and graph setup unchanged

def connect(self, graph=None, dsn=None, ..., load_from_plugins=False, skip_load=False, **kwargs):
    conn = psycopg.connect(dsn, cursor_factory=cursor_factory, **kwargs)
    setUpAge(conn, graph, load_from_plugins, skip_load=skip_load)
    # ...

Current Workaround

Users must bypass Age.connect() entirely and manually replicate the setup logic, which is fragile and tightly coupled to internal driver implementation details:

import age
import psycopg
from psycopg.types import TypeInfo

ag = age.Age()
ag.connection = psycopg.connect(dsn, cursor_factory=age.age.ClientCursor)

with ag.connection.cursor() as cur:
    cur.execute('SET search_path = ag_catalog, "$user", public;')

ag_info = TypeInfo.fetch(ag.connection, "agtype")
ag.connection.adapters.register_loader(ag_info.oid, age.age.AgeLoader)
ag.connection.adapters.register_loader(ag_info.array_oid, age.age.AgeLoader)

ag.graphName = "my_graph"
age.checkGraphCreated(ag.connection, "my_graph")
ag.connection.commit()

Environment

• Apache AGE: 1.6.0 (also affects master)
• Python driver: master branch (psycopg3 version)
• psycopg: 3.2.x
• Python: 3.13
• Managed PostgreSQL: Azure Database for PostgreSQL – Flexible Server

[jrgemignani]

@uesleilima If you would like to create a PR to fix this, I can review it :)

[uesleilima]

PR #2366 created @jrgemignani . Thanks