unwrap pg connection/cursor for new psycopg2 extensions (#621)

psycopg2 2.8 added two new extensions, `quote_ident` and
`encrypt_password`. Both extensions do a type check on the
connection/cursor object, which fails if they are wrapped with our
object proxies. Similar to `register_type`, we need to unwrap the
connection/cursor before handing it to these extensions.

fixes #620

Author: beniwohli

CHANGELOG.md

@@ -6,6 +6,7 @@
 ### Bugfixes
 
  * ensure that metrics with value 0 are not collected if they have the `reset_on_collect` flag set (#615)
+ * unwrap postgres cursor for newly introduced psycopg2 extensions (#621)
 
 ## v5.2.2
 [Check the diff](https://github.com/elastic/apm-agent-python/compare/v5.2.1...v5.2.2)

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)

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",

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):