Exception when using `quote_ident` in psycopg2

[HUSSTECH]

Describe the bug:
If you make use of the function psycopg2.extensions.quote_ident docs, a TypeError exception is thrown. This is because the cursor object, when under instrumentation from ES-APM, is an instance of PGCursorProxy, not the actual cursor, and quote_ident does not allow this because the type is checked in the C code link. With the error message saying TypeError: argument 2 must be a connection or a cursor. Inspecting the cur object at a debug breakpoint we can see it is the proxy object:

>>> cur
<PGCursorProxy at 0x7fd7f70f9a88 for NamedTupleCursor at 0x7fd7f70f0148>
>>> type(cur)
<class 'elasticapm.instrumentation.packages.psycopg2.PGCursorProxy'>
>>> type(cur.__wrapped__)
<class 'psycopg2.extras.NamedTupleCursor'>

To Reproduce

from psycopg2.extensions import quote_ident
....
....
with psycopg2.connect(DSN) as conn:
    with conn.cursor() as curs:
        ident = quote_ident("column_name", cur)
        curs.execute(f"SELECT {column_name} FROM data.table;")
        data = curs.fetchall()

passing the underlying wrapped cursor works:

from psycopg2.extensions import quote_ident
....
....
with psycopg2.connect(DSN) as conn:
    with conn.cursor() as curs:
        ident = quote_ident("column_name", cur.__wrapped__)
        curs.execute(f"SELECT {column_name} FROM data.table;")
        data = curs.fetchall()

Environment (please complete the following information)

• OS: Linux
• Python version: 3.6.4
• Agent version: 5.2.2

Additional context
Looks like the same problem was encountered here DataDog/dd-trace-py#474, and was fixed by also patching quote_ident to pass the __wrapped__ object. Testing this out with a basic in-module proxy function worked, but obviously a patch right at the top level from the apm module would sort it out.

def quote_ident(string, cursor):
    return psycopg2.extensions.quote_ident(string, cursor.__wrapped__)

[beniwohli]

@HUSSTECH thanks for the very detailed bug report. I opened #621 to fix this (and also added handling for the encrypt_password extension while I'm at it). I will release this in the coming days.

If you have the time, it would be super cool if you could give the fix a quick test run in your development environment. You can install it like so:

pip install "https://github.com/beniwohli/apm-agent-python/archive/fix-620.zip#egg=elastic-apm==100"

[HUSSTECH]

Hey @beniwohli, thanks for the quick turnaround. I've been attempting to test out the fix, but I'm still getting the same error message. Not sure if you were ready for me to test now, but I tried both pip install "https://github.com/beniwohli/apm-agent-python/archive/fix-620.zip#egg=elastic-apm==100" and pip install "https://github.com/beniwohli/apm-agent-python/archive/fix-620.zip incase I wasn't getting the head of the branch.

[beniwohli]

Thanks for testing it! I'm not quite sure why this isn't working for you, but one possible option is that it might be that you import quote_ident before we can patch it, so you get the original, unpatched version.

Can you try to either

• change the import from from psycopg2.extensions import quote_ident to import psycopg2.extensions and then use pyscopg2.extensions.quote_ident(my_string, cursor)
• import quote_ident inside the function/method where you use it

Unfortunately, both options aren't really satisfactory, as it requires code changes on your side :(

[HUSSTECH]

Ok, that's super interesting. Both suggested workarounds work. They aren't major changes to my code, and option 1 is probably better of the two. So it just means doing a more broad import for psycopg2.extensions first. But wondering why there is a difference, is the patching done only x levels deep into the module?

I wanted to test as pure of a case as possible. So I set up a minimal single file Flask application.

1. This one does not work (Basically this one is replicating the issue I opened):

from elasticapm.contrib.flask import ElasticAPM
from flask import Flask, jsonify
import psycopg2
from psycopg2.extensions import quote_ident


app = Flask(__name__)


app.config['ELASTIC_APM'] = {
    'SERVICE_NAME': 'development',
    'SERVER_URL': 'xxxx',
    'ENVIRONMENT': 'development',
    'SECRET_TOKEN': 'xxxxx',
    'AUTO_LOG_STACKS': False,
    'COLLECT_LOCAL_VARIABLES': 'off',
    'SOURCE_LINES_ERROR_APP_FRAMES': 0,
    'SOURCE_LINES_ERROR_LIBRARY_FRAMES': 0,
    'SOURCE_LINES_SPAN_APP_FRAMES': 0,
    'SOURCE_LINES_SPAN_LIBRARY_FRAMES': 0
}

apm = ElasticAPM(app)


def connect_db():
    return psycopg2.connect("DSN")


@app.route('/test')
def test():
    conn = connect_db()
    cur = conn.cursor()
    ident = quote_ident("column", cur)
    cur.execute(f"SELECT {ident} FROM data.table;")
    data = cur.fetchall()
    cur.close()
    conn.close()
    return jsonify(data)

2. The following works (Using the more broad import of psycopg2):

from elasticapm.contrib.flask import ElasticAPM
from flask import Flask, jsonify
import psycopg2


app = Flask(__name__)


app.config['ELASTIC_APM'] = {
    'SERVICE_NAME': 'development',
    'SERVER_URL': 'xxxx',
    'ENVIRONMENT': 'development',
    'SECRET_TOKEN': 'xxxxx',
    'AUTO_LOG_STACKS': False,
    'COLLECT_LOCAL_VARIABLES': 'off',
    'SOURCE_LINES_ERROR_APP_FRAMES': 0,
    'SOURCE_LINES_ERROR_LIBRARY_FRAMES': 0,
    'SOURCE_LINES_SPAN_APP_FRAMES': 0,
    'SOURCE_LINES_SPAN_LIBRARY_FRAMES': 0
}

apm = ElasticAPM(app)


def connect_db():
    return psycopg2.connect("DSN")


@app.route('/test')
def test():
    conn = connect_db()
    cur = conn.cursor()
    ident = psycopg2.extensions.quote_ident("column", cur)
    cur.execute(f"SELECT {ident} FROM data.table;")
    data = cur.fetchall()
    cur.close()
    conn.close()
    return jsonify(data)

3. The following also works (By specifically importing quote_ident after the ElasticAPM(app)):

from elasticapm.contrib.flask import ElasticAPM
from flask import Flask, jsonify
import psycopg2


app = Flask(__name__)


app.config['ELASTIC_APM'] = {
    'SERVICE_NAME': 'development',
    'SERVER_URL': 'xxxx',
    'ENVIRONMENT': 'development',
    'SECRET_TOKEN': 'xxxxx',
    'AUTO_LOG_STACKS': False,
    'COLLECT_LOCAL_VARIABLES': 'off',
    'SOURCE_LINES_ERROR_APP_FRAMES': 0,
    'SOURCE_LINES_ERROR_LIBRARY_FRAMES': 0,
    'SOURCE_LINES_SPAN_APP_FRAMES': 0,
    'SOURCE_LINES_SPAN_LIBRARY_FRAMES': 0
}

apm = ElasticAPM(app)

from psycopg2.extensions import quote_ident

def connect_db():
    return psycopg2.connect("DSN")


@app.route('/test')
def test():
    conn = connect_db()
    cur = conn.cursor()
    ident = quote_ident("column", cur)
    cur.execute(f"SELECT {ident} FROM data.table;")
    data = cur.fetchall()
    cur.close()
    conn.close()
    return jsonify(data)

in particular for this third example, it just doesn't seem great to have this behaviour be so dependent on module import order, or the developer just so happening to use a more broad import.

Maybe this is just how it has to be for instrumentation libraries in Python.

[beniwohli]

The technical reason why from psycopg.extensions import quote_ident doesn't work and the other two approaches work has to do with the import cache (specifically, sys.modules).

When you import a module, Python first checks sys.modules to see if that module has already been imported, and if yes, uses that module. That basically makes Python modules singletons (you always get back the same "module" object when importing a certain module). Our patching mechanisms import a module (psycopg2.extensions in this case), and apply the patch to the given function (quote_ident).

This means that if you import a module, it doesn't matter if our patching happens before or after your import, since we both operate on the same object.

On the other hand, when you do from psycopg.extensions import quote_ident, you import the psycopg2.extensions module, then take the quote_ident function from that module and assign it to a variable in the local scope. If this happens before our patching, quote_ident in your local scope will not be affected by our patching. As far as I know, there is no way around this :(

[HUSSTECH]

ah ok, I thought "leaf" objects at the end of import chains would still resolve themselves (walking back up the call path) against a single version/instance of the entire library. So this fix for quote_ident, just makes it in line with the other instrumentation in the APM code.

The issue will come with quote_ident being used in a library. This is the case for us, but luckily it is one of our libraries, so not a problem to tweak the code to make it work with the instrumentation.

Anyway, with all that, guess we can say this issue has been pretty thoroughly tested then. I can confirm I see database traces in our test APM UI. Am I able to control the logging of the database statement anywhere? thanks for all the help!