Bug: XRay SDK doesn't support pg8000 native driver

Inizia leggendo aws_xray_sdk/ext/pg8000/patch.py, in particolare la gestione esistente di pg8000.connect, quindi confrontala con pg8000.native.Connection e con il suo punto di ingresso run descritto nell’issue. Il lavoro è completato quando le connessioni native e le chiamate a run producono sottosegmenti SQL X-Ray acquisiti, come nel percorso della Standard API.

Summary

pg8000 comes with two APIs, the pg8000 native API and the DB-API 2.0 Standard API (PEP-0249) and the XRay SDK fails to capture SQL segments when using the pg8000 native API.

Explanation, code example and output

When working with the pg8000 native API, customers initialize a connection using pg8000.native.Connection and xray is unable to intercept and wrapper this because in this line it always considers pg8000.connect and not pg8000.native.Connection. Also, when using pg8000 native API the connection instance doesn't have cursor, all the queries are executed directly using run command.

Check this code example:

from urllib.parse import quote_plus
import pg8000.native
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core import patch_all

# Initialize X-Ray
xray_recorder.configure(service='MyPostgresApp')
patch_all()

@xray_recorder.capture('connect_to_db')
def connect_to_db():
    try:

        host = "MYHOST"
        user = "MYUSER"
        database = "MYDATABASE"
        password = "MYPASSWORD"

        conn = pg8000.native.Connection(
            database=database,
            user=user,
            password=password,
            host=host,
            port=5432,
            ssl_context=True
        )
        return conn
    except Exception as e:
        print(f"Error connecting to database: {e}")
        return None

@xray_recorder.capture('query_db')
def query_db(conn):
    try:
        results = conn.run("SELECT * FROM public.todos LIMIT 5")
        return results
    except Exception as e:
        print(f"Error querying database: {e}")
        return None

def main():
    with xray_recorder.in_segment('main_function'):
        conn = connect_to_db()
        if conn:
            results = query_db(conn)
            if results:
                print("Query results:", results)
            conn.close()

if __name__ == "__main__":
    main()

This creates the json below and we see that there is no run subsegment in the query_db segment because XRAY cannot patch this driver when using native Connection.

{
   "id":"caeb26df6fe242ed",
   "name":"main_function",
   "start_time":1736462563.884077,
   "in_progress":false,
   "aws":{
      "xray":{
         "sdk":"X-Ray for Python",
         "sdk_version":"2.14.0"
      }
   },
   "subsegments":[
      {
         "id":"981eb724c5369158",
         "name":"connect_to_db",
         "start_time":1736462563.884105,
         "parent_id":"caeb26df6fe242ed",
         "in_progress":false,
         "trace_id":"1-678050e3-fbec3813f4280021b9069ff5",
         "type":"subsegment",
         "namespace":"local",
         "end_time":1736462564.930994
      },
      {
         "id":"eb5a8872a770eacd",
         "name":"query_db",
         "start_time":1736462564.931169,
         "parent_id":"caeb26df6fe242ed",
         "in_progress":false,
         "trace_id":"1-678050e3-fbec3813f4280021b9069ff5",
         "type":"subsegment",
         "namespace":"local",
         "end_time":1736462565.3793
      }
   ],
   "trace_id":"1-678050e3-fbec3813f4280021b9069ff5",
   "service":{
      "runtime":"CPython",
      "runtime_version":"3.12.8"
   },
   "end_time":1736462565.3813171
}

Solution

I created a small PoC locally and made it work by adding a new method to wrap native.Connection and creating a new class called XRayTracedNative to intercept the run method and then be able to create the subsegment.

*the names are just suggestions, you can choose whatever name you want.

Please let me know if you are interested in a PR to implement this and I can work.