aws / aws/aws-xray-sdk-python

Bug: XRay SDK doesn't support pg8000 native driver

Open
#451 2 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
338
Forks
147
PR merge metrics
No merged PRs in 30d

Description

## Summary
pg8000 comes with two APIs, the pg8000 native API and the DB-API 2.0 [Standard API](https://peps.python.org/pep-0249/#introduction) (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](https://github.com/aws/aws-xray-sdk-python/blob/master/aws_xray_sdk/ext/pg8000/patch.py#L13) 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:
```python
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.

```json
{
"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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.