New Relic python client, SQLAlchemy and psycopg2 library

README.md

If you attempt to register an extension with psycopg2 and happen to use the newrelic python client library you may get an error that looks like this

  File "/opt/balanced/embedded/lib/python2.7/site-packages/psycopg2/_json.py", line 142, in register_default_json
    loads=loads, oid=JSON_OID, array_oid=JSONARRAY_OID)
  File "/opt/balanced/embedded/lib/python2.7/site-packages/psycopg2/_json.py", line 125, in register_json
    register_type(JSON, not globally and conn_or_curs or None)
TypeError: argument 2 must be a connection, cursor or None

The issue you're running into is that newrelic decides to play god and replace the connection object with a proxy that it implements called ConnectionWrapper.

This is a problem because psycopg2 has some code that inspects the type of object in order to determine what how to register the extension.

Luckily SQLAlchemy checks for a handy method called _sqla_unwrap which can be used to get the underlying connection and pass it to psycopg. I guess the engineers at newrelic didn't get around to testing this scenario and their code includes a bunch of anonymous classes so it's hard to inherit the object and patch it yourself. Here's how you can fix it yourself:

Write some code to patch the connection class

class NewRelicConnectionWrapperProxy(object):
    def __init__(self, connect):
        self.__connect = connect

    def __call__(self, *args, **kwargs):
        cxn = self.__connect(*args, **kwargs)
        object.__setattr__(cxn, '_sqla_unwrap', cxn._nr_connection)
        return cxn

Now, after you initialize the newrelic agent but before you initialize your database connection you can monkey patch everything:

import newrelic.agent
newrelic.agent.initialize(config_file, environment)
import psycopg2
psycopg2.connect = NewRelicConnectionWrapperProxy(psycopg2.connect)

That's all there is to it. I hope this is handy for someone.

To New Relic devs: If you fix this, please do not use anonymous classes, it's really hard to extend your code when you cannot inherit and override a class!

Thanks for the solution, but it only works with newrelic<=2.0.0.1. In newrelic>=2.10.0.8 there are a lot of new proxy/wrapper magic with __setattr__ and __getattr__ definitions. In new versions the psycopg.connection is stored in cxn.__wrapped__ property instead of _nr_connection (as it was before). But there is still an issue with setattr: the statement

object.__setattr__(cxn, '_sqla_unwrap', cxn.__wrapped__)

raises the following exception:

...
  File "/home/spy/Workspace/venv_dev/lib/python2.6/site-packages/sqlalchemy/pool.py", line 396, in __init__
    self.connection = self.__connect()
  File "/home/spy/Workspace/venv_dev/lib/python2.6/site-packages/sqlalchemy/pool.py", line 509, in __connect
    connection = self.__pool._creator()
  File "/home/spy/Workspace/venv_dev/lib/python2.6/site-packages/sqlalchemy/engine/strategies.py", line 89, in connect
    return dialect.connect(*cargs, **cparams)
  File "/home/spy/Workspace/venv_dev/lib/python2.6/site-packages/sqlalchemy/engine/default.py", line 366, in connect
    return self.dbapi.connect(*cargs, **cparams)
  File "/home/spy/Workspace/app/wsgi.py", line 22, in __call__
    object.__setattr__(cxn, '_sqla_unwrap', cxn.__wrapped__)
TypeError: can't apply this __setattr__ to ObjectProxy object

Any ideas or suggestions on how to fix this?

New Relic has for some time had instrumentation around psycopg2.register_type() specifically to unwrap the object passed in if required and avoid this issue. We include checks to make sure it is working in our test suites. If this is not doing its job, then you likely aren't initialising the New Relic agent prior to psycopg2 modules having being imported.

In future if you have issues it would be better to contact New Relic support directly rather than simply come up with a workaround. We will always address such issues if we are told about them.

FWIW, the issue is caused by psycopg2 bypassing its own public psycopg2.extensions module to access register_type(), going direct to the C extension module. A workaround should be able to be setup from the agent configuration file to remap what our instrumentation is applied to, until we can add a more permanent fix into the agent package itself. When I have been able to test the workaround and a more permanent fix I will add the details of the workaround here.

Thanks @GrahamDumpleton! Adding the following to my newrelic.ini seemed to work around this issue for me:

[import-hook:psycopg2]
enabled = false

However, I am curious if a workaround has been implemented on the newrelic agent side, so that we can work around the problem without having to disable our NR database instrumentation.

The only way I've found to avoid this issue is to use register_json() instead of register_type(). When all arguments are provided it doesn't really have any JSON-specific logic, contrary to what the name may suggest.

For example, the following snippet will raise TypeError: argument 2 must be a connection, cursor or None
when used with Newrelic's agent:

from psycopg2.extensions import register_type, new_type, new_array_type

numeric_type = new_type((1700,), "DEC2FLOAT",  lambda v, _: float(v) if v is not None else None)
register_type(numeric_type, cursor)

numeric_array_type = new_array_type((1231,), "DEC2FLOATARRAY", numeric_type)
register_type(numeric_array_type, cursor)

While this snippet won't raise the error and will produce an equivalent result:

from psycopg2.extras import register_json

register_json(
    cursor,
    globally=False,
    loads=float,
    oid=1700,
    array_oid=1231,
    name="DEC2FLOAT",
)

I'm not really sure why one works and the other doesn't. The internal implementation of the register_json function uses practically the same code as the first snippet, so I'm guessing it's due to some import order shenanigans.