prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  DB interactions not traced when using context manager with psycopg2 connections or cursors
  When using a context manager with psycopg2 connections or cursors, db interactions are not captured in spans.
  
  The code below generates a span for `psycopg2.connect`, but not the query:
  ```
  with psycopg2.connect(DSN) as conn:
      with conn.cursor() as curs:
          curs.execute("SELECT * FROM data.portfolio;")
          portfolios = curs.fetchall()
  ```
  
  whereas the following captures both spans as expected:
  ```
  conn = psycopg2.connect(DSN)
  curs = conn.cursor()
  curs.execute("SELECT * FROM data.portfolio;")
  portfolios = curs.fetchall()
  curs.close()
  conn.close()
  ```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `elasticapm/instrumentation/packages/psycopg2.py`
  Content:
  ```
  1 #  BSD 3-Clause License
  2 #
  3 #  Copyright (c) 2019, Elasticsearch BV
  4 #  All rights reserved.
  5 #
  6 #  Redistribution and use in source and binary forms, with or without
  7 #  modification, are permitted provided that the following conditions are met:
  8 #
  9 #  * Redistributions of source code must retain the above copyright notice, this
  10 #    list of conditions and the following disclaimer.
  11 #
  12 #  * Redistributions in binary form must reproduce the above copyright notice,
  13 #    this list of conditions and the following disclaimer in the documentation
  14 #    and/or other materials provided with the distribution.
  15 #
  16 #  * Neither the name of the copyright holder nor the names of its
  17 #    contributors may be used to endorse or promote products derived from
  18 #    this software without specific prior written permission.
  19 #
  20 #  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  21 #  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  22 #  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  23 #  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  24 #  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  25 #  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  26 #  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  27 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  28 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  29 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  30 
  31 from elasticapm.instrumentation.packages.dbapi2 import (
  32     ConnectionProxy,
  33     CursorProxy,
  34     DbApi2Instrumentation,
  35     extract_signature,
  36 )
  37 from elasticapm.traces import capture_span
  38 from elasticapm.utils import default_ports
  39 
  40 
  41 class PGCursorProxy(CursorProxy):
  42     provider_name = "postgresql"
  43 
  44     def _bake_sql(self, sql):
  45         # if this is a Composable object, use its `as_string` method
  46         # see http://initd.org/psycopg/docs/sql.html
  47         if hasattr(sql, "as_string"):
  48             return sql.as_string(self.__wrapped__)
  49         return sql
  50 
  51     def extract_signature(self, sql):
  52         return extract_signature(sql)
  53 
  54 
  55 class PGConnectionProxy(ConnectionProxy):
  56     cursor_proxy = PGCursorProxy
  57 
  58 
  59 class Psycopg2Instrumentation(DbApi2Instrumentation):
  60     name = "psycopg2"
  61 
  62     instrument_list = [("psycopg2", "connect")]
  63 
  64     def call(self, module, method, wrapped, instance, args, kwargs):
  65         signature = "psycopg2.connect"
  66 
  67         host = kwargs.get("host")
  68         if host:
  69             signature += " " + str(host)
  70 
  71             port = kwargs.get("port")
  72             if port:
  73                 port = str(port)
  74                 if int(port) != default_ports.get("postgresql"):
  75                     signature += ":" + port
  76         else:
  77             # Parse connection string and extract host/port
  78             pass
  79 
  80         with capture_span(signature, span_type="db", span_subtype="postgresql", span_action="connect"):
  81             return PGConnectionProxy(wrapped(*args, **kwargs))
  82 
  83 
  84 class Psycopg2RegisterTypeInstrumentation(DbApi2Instrumentation):
  85     name = "psycopg2-register-type"
  86 
  87     instrument_list = [
  88         ("psycopg2.extensions", "register_type"),
  89         # specifically instrument `register_json` as it bypasses `register_type`
  90         ("psycopg2._json", "register_json"),
  91     ]
  92 
  93     def call(self, module, method, wrapped, instance, args, kwargs):
  94         if "conn_or_curs" in kwargs and hasattr(kwargs["conn_or_curs"], "__wrapped__"):
  95             kwargs["conn_or_curs"] = kwargs["conn_or_curs"].__wrapped__
  96         # register_type takes the connection as second argument
  97         elif len(args) == 2 and hasattr(args[1], "__wrapped__"):
  98             args = (args[0], args[1].__wrapped__)
  99         # register_json takes the connection as first argument, and can have
  100         # several more arguments
  101         elif method == "register_json":
  102             if args and hasattr(args[0], "__wrapped__"):
  103                 args = (args[0].__wrapped__,) + args[1:]
  104 
  105         return wrapped(*args, **kwargs)
  106
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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
  @@ -51,10 +51,16 @@
       def extract_signature(self, sql):
           return extract_signature(sql)
   
  +    def __enter__(self):
  +        return PGCursorProxy(self.__wrapped__.__enter__())
  +
   
   class PGConnectionProxy(ConnectionProxy):
       cursor_proxy = PGCursorProxy
   
  +    def __enter__(self):
  +        return PGConnectionProxy(self.__wrapped__.__enter__())
  +
   
   class Psycopg2Instrumentation(DbApi2Instrumentation):
       name = "psycopg2"
  
verification_info: {"golden_diff": "diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py\n--- a/elasticapm/instrumentation/packages/psycopg2.py\n+++ b/elasticapm/instrumentation/packages/psycopg2.py\n@@ -51,10 +51,16 @@\n     def extract_signature(self, sql):\n         return extract_signature(sql)\n \n+    def __enter__(self):\n+        return PGCursorProxy(self.__wrapped__.__enter__())\n+\n \n class PGConnectionProxy(ConnectionProxy):\n     cursor_proxy = PGCursorProxy\n \n+    def __enter__(self):\n+        return PGConnectionProxy(self.__wrapped__.__enter__())\n+\n \n class Psycopg2Instrumentation(DbApi2Instrumentation):\n     name = \"psycopg2\"\n", "issue": "DB interactions not traced when using context manager with psycopg2 connections or cursors\nWhen using a context manager with psycopg2 connections or cursors, db interactions are not captured in spans.\r\n\r\nThe code below generates a span for `psycopg2.connect`, but not the query:\r\n```\r\nwith psycopg2.connect(DSN) as conn:\r\n    with conn.cursor() as curs:\r\n        curs.execute(\"SELECT * FROM data.portfolio;\")\r\n        portfolios = curs.fetchall()\r\n```\r\n\r\nwhereas the following captures both spans as expected:\r\n```\r\nconn = psycopg2.connect(DSN)\r\ncurs = conn.cursor()\r\ncurs.execute(\"SELECT * FROM data.portfolio;\")\r\nportfolios = curs.fetchall()\r\ncurs.close()\r\nconn.close()\r\n```\n", "before_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n\nfrom elasticapm.instrumentation.packages.dbapi2 import (\n    ConnectionProxy,\n    CursorProxy,\n    DbApi2Instrumentation,\n    extract_signature,\n)\nfrom elasticapm.traces import capture_span\nfrom elasticapm.utils import default_ports\n\n\nclass PGCursorProxy(CursorProxy):\n    provider_name = \"postgresql\"\n\n    def _bake_sql(self, sql):\n        # if this is a Composable object, use its `as_string` method\n        # see http://initd.org/psycopg/docs/sql.html\n        if hasattr(sql, \"as_string\"):\n            return sql.as_string(self.__wrapped__)\n        return sql\n\n    def extract_signature(self, sql):\n        return extract_signature(sql)\n\n\nclass PGConnectionProxy(ConnectionProxy):\n    cursor_proxy = PGCursorProxy\n\n\nclass Psycopg2Instrumentation(DbApi2Instrumentation):\n    name = \"psycopg2\"\n\n    instrument_list = [(\"psycopg2\", \"connect\")]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        signature = \"psycopg2.connect\"\n\n        host = kwargs.get(\"host\")\n        if host:\n            signature += \" \" + str(host)\n\n            port = kwargs.get(\"port\")\n            if port:\n                port = str(port)\n                if int(port) != default_ports.get(\"postgresql\"):\n                    signature += \":\" + port\n        else:\n            # Parse connection string and extract host/port\n            pass\n\n        with capture_span(signature, span_type=\"db\", span_subtype=\"postgresql\", span_action=\"connect\"):\n            return PGConnectionProxy(wrapped(*args, **kwargs))\n\n\nclass Psycopg2RegisterTypeInstrumentation(DbApi2Instrumentation):\n    name = \"psycopg2-register-type\"\n\n    instrument_list = [\n        (\"psycopg2.extensions\", \"register_type\"),\n        # specifically instrument `register_json` as it bypasses `register_type`\n        (\"psycopg2._json\", \"register_json\"),\n    ]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        if \"conn_or_curs\" in kwargs and hasattr(kwargs[\"conn_or_curs\"], \"__wrapped__\"):\n            kwargs[\"conn_or_curs\"] = kwargs[\"conn_or_curs\"].__wrapped__\n        # register_type takes the connection as second argument\n        elif len(args) == 2 and hasattr(args[1], \"__wrapped__\"):\n            args = (args[0], args[1].__wrapped__)\n        # register_json takes the connection as first argument, and can have\n        # several more arguments\n        elif method == \"register_json\":\n            if args and hasattr(args[0], \"__wrapped__\"):\n                args = (args[0].__wrapped__,) + args[1:]\n\n        return wrapped(*args, **kwargs)\n", "path": "elasticapm/instrumentation/packages/psycopg2.py"}], "after_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n\nfrom elasticapm.instrumentation.packages.dbapi2 import (\n    ConnectionProxy,\n    CursorProxy,\n    DbApi2Instrumentation,\n    extract_signature,\n)\nfrom elasticapm.traces import capture_span\nfrom elasticapm.utils import default_ports\n\n\nclass PGCursorProxy(CursorProxy):\n    provider_name = \"postgresql\"\n\n    def _bake_sql(self, sql):\n        # if this is a Composable object, use its `as_string` method\n        # see http://initd.org/psycopg/docs/sql.html\n        if hasattr(sql, \"as_string\"):\n            return sql.as_string(self.__wrapped__)\n        return sql\n\n    def extract_signature(self, sql):\n        return extract_signature(sql)\n\n    def __enter__(self):\n        return PGCursorProxy(self.__wrapped__.__enter__())\n\n\nclass PGConnectionProxy(ConnectionProxy):\n    cursor_proxy = PGCursorProxy\n\n    def __enter__(self):\n        return PGConnectionProxy(self.__wrapped__.__enter__())\n\n\nclass Psycopg2Instrumentation(DbApi2Instrumentation):\n    name = \"psycopg2\"\n\n    instrument_list = [(\"psycopg2\", \"connect\")]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        signature = \"psycopg2.connect\"\n\n        host = kwargs.get(\"host\")\n        if host:\n            signature += \" \" + str(host)\n\n            port = kwargs.get(\"port\")\n            if port:\n                port = str(port)\n                if int(port) != default_ports.get(\"postgresql\"):\n                    signature += \":\" + port\n        else:\n            # Parse connection string and extract host/port\n            pass\n\n        with capture_span(signature, span_type=\"db\", span_subtype=\"postgresql\", span_action=\"connect\"):\n            return PGConnectionProxy(wrapped(*args, **kwargs))\n\n\nclass Psycopg2RegisterTypeInstrumentation(DbApi2Instrumentation):\n    name = \"psycopg2-register-type\"\n\n    instrument_list = [\n        (\"psycopg2.extensions\", \"register_type\"),\n        # specifically instrument `register_json` as it bypasses `register_type`\n        (\"psycopg2._json\", \"register_json\"),\n    ]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        if \"conn_or_curs\" in kwargs and hasattr(kwargs[\"conn_or_curs\"], \"__wrapped__\"):\n            kwargs[\"conn_or_curs\"] = kwargs[\"conn_or_curs\"].__wrapped__\n        # register_type takes the connection as second argument\n        elif len(args) == 2 and hasattr(args[1], \"__wrapped__\"):\n            args = (args[0], args[1].__wrapped__)\n        # register_json takes the connection as first argument, and can have\n        # several more arguments\n        elif method == \"register_json\":\n            if args and hasattr(args[0], \"__wrapped__\"):\n                args = (args[0].__wrapped__,) + args[1:]\n\n        return wrapped(*args, **kwargs)\n", "path": "elasticapm/instrumentation/packages/psycopg2.py"}]}