Repository: elastic/apm-agent-python

Issue:
Exception when using `quote_ident` in psycopg2
**Describe the bug**:
If you make use of the function `psycopg2.extensions.quote_ident` [docs](http://initd.org/psycopg/docs/extensions.html#psycopg2.extensions.quote_ident), 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](https://github.com/psycopg/psycopg2/blob/2_7_6_1/psycopg/psycopgmodule.c#L181). 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**

```python
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:
```python
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 https://github.com/DataDog/dd-trace-py/issues/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.

```python
def quote_ident(string, cursor):
    return psycopg2.extensions.quote_ident(string, cursor.__wrapped__)
```




Discussion:

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py
--- a/elasticapm/instrumentation/packages/psycopg2.py
+++ b/elasticapm/instrumentation/packages/psycopg2.py
@@ -87,13 +87,21 @@ def call(self, module, method, wrapped, instance, args, kwargs):
             return PGConnectionProxy(wrapped(*args, **kwargs))
 
 
-class Psycopg2RegisterTypeInstrumentation(DbApi2Instrumentation):
-    name = "psycopg2-register-type"
+class Psycopg2ExtensionsInstrumentation(DbApi2Instrumentation):
+    """
+    Some extensions do a type check on the Connection/Cursor in C-code, which our
+    proxy fails. For these extensions, we need to ensure that the unwrapped
+    Connection/Cursor is passed.
+    """
+
+    name = "psycopg2"
 
     instrument_list = [
         ("psycopg2.extensions", "register_type"),
         # specifically instrument `register_json` as it bypasses `register_type`
         ("psycopg2._json", "register_json"),
+        ("psycopg2.extensions", "quote_ident"),
+        ("psycopg2.extensions", "encrypt_password"),
     ]
 
     def call(self, module, method, wrapped, instance, args, kwargs):
@@ -108,4 +116,11 @@ def call(self, module, method, wrapped, instance, args, kwargs):
             if args and hasattr(args[0], "__wrapped__"):
                 args = (args[0].__wrapped__,) + args[1:]
 
+        elif method == "encrypt_password":
+            # connection/cursor is either 3rd argument, or "scope" keyword argument
+            if len(args) >= 3 and hasattr(args[2], "__wrapped__"):
+                args = args[:2] + (args[2].__wrapped__,) + args[3:]
+            elif "scope" in kwargs and hasattr(kwargs["scope"], "__wrapped__"):
+                kwargs["scope"] = kwargs["scope"].__wrapped__
+
         return wrapped(*args, **kwargs)
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/elasticapm/instrumentation/register.py
@@ -35,7 +35,7 @@
     "elasticapm.instrumentation.packages.botocore.BotocoreInstrumentation",
     "elasticapm.instrumentation.packages.jinja2.Jinja2Instrumentation",
     "elasticapm.instrumentation.packages.psycopg2.Psycopg2Instrumentation",
-    "elasticapm.instrumentation.packages.psycopg2.Psycopg2RegisterTypeInstrumentation",
+    "elasticapm.instrumentation.packages.psycopg2.Psycopg2ExtensionsInstrumentation",
     "elasticapm.instrumentation.packages.mysql.MySQLInstrumentation",
     "elasticapm.instrumentation.packages.pylibmc.PyLibMcInstrumentation",
     "elasticapm.instrumentation.packages.pymongo.PyMongoInstrumentation",


Test patch:
diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py
--- a/tests/instrumentation/psycopg2_tests.py
+++ b/tests/instrumentation/psycopg2_tests.py
@@ -48,6 +48,11 @@
     # as of Jan 2018, psycopg2cffi doesn't have this module
     has_sql_module = False
 
+try:
+    import psycopg2.extensions
+except ImportError:
+    psycopg2.extensions = None
+
 
 pytestmark = pytest.mark.psycopg2
 
@@ -259,7 +264,7 @@ def test_fully_qualified_table_name():
 
 @pytest.mark.integrationtest
 @pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
-def test_psycopg2_register_type(postgres_connection, elasticapm_client):
+def test_psycopg2_register_type(instrument, postgres_connection, elasticapm_client):
     import psycopg2.extras
 
     elasticapm_client.begin_transaction("web.django")
@@ -271,7 +276,7 @@ def test_psycopg2_register_type(postgres_connection, elasticapm_client):
 
 @pytest.mark.integrationtest
 @pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
-def test_psycopg2_register_json(postgres_connection, elasticapm_client):
+def test_psycopg2_register_json(instrument, postgres_connection, elasticapm_client):
     # register_json bypasses register_type, so we have to test unwrapping
     # separately
     import psycopg2.extras
@@ -286,6 +291,40 @@ def test_psycopg2_register_json(postgres_connection, elasticapm_client):
     elasticapm_client.end_transaction(None, "test-transaction")
 
 
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+@pytest.mark.skipif(
+    not hasattr(psycopg2.extensions, "quote_ident"), reason="psycopg2 driver doesn't have quote_ident extension"
+)
+def test_psycopg2_quote_ident(instrument, postgres_connection, elasticapm_client):
+    elasticapm_client.begin_transaction("web.django")
+    ident = psycopg2.extensions.quote_ident("x'x", postgres_connection)
+    elasticapm_client.end_transaction(None, "test-transaction")
+
+    assert ident == '"x\'x"'
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+@pytest.mark.skipif(
+    not hasattr(psycopg2.extensions, "encrypt_password"),
+    reason="psycopg2 driver doesn't have encrypt_password extension",
+)
+@pytest.mark.skipif(
+    hasattr(psycopg2, "__libpq_version__") and psycopg2.__libpq_version__ < 100000,
+    reason="test code requires libpq >= 10",
+)
+def test_psycopg2_encrypt_password(instrument, postgres_connection, elasticapm_client):
+    elasticapm_client.begin_transaction("web.django")
+    pw1 = psycopg2.extensions.encrypt_password("user", "password", postgres_connection)
+    pw2 = psycopg2.extensions.encrypt_password("user", "password", postgres_connection, None)
+    pw3 = psycopg2.extensions.encrypt_password("user", "password", postgres_connection, algorithm=None)
+    pw4 = psycopg2.extensions.encrypt_password("user", "password", scope=postgres_connection, algorithm=None)
+    elasticapm_client.end_transaction(None, "test-transaction")
+
+    assert pw1.startswith("md5") and (pw1 == pw2 == pw3 == pw4)
+
+
 @pytest.mark.integrationtest
 @pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
 def test_psycopg2_tracing_outside_of_elasticapm_transaction(instrument, postgres_connection, elasticapm_client):