From a6c560b4c0f5d1d195f66db6fe72bd96a5bd7c88 Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@edgedb.com>
Date: Tue, 3 Aug 2021 13:55:29 -0700
Subject: [PATCH] Disable JIT while doing type introspection

The misapplication of JIT to asyncpg introspection queries has been a
constant source of user complaints.

Closes: #530
Closes: #1078
Previously: #875, #794, #782, #741, #727 (and probably more).
---
 asyncpg/connection.py       | 44 ++++++++++++++++++++++++++++++++++---
 tests/test_introspection.py |  8 ++++++-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/asyncpg/connection.py b/asyncpg/connection.py
index 45cf99b1..06e4ce23 100644
--- a/asyncpg/connection.py
+++ b/asyncpg/connection.py
@@ -461,7 +461,26 @@ async def _get_statement(
         return statement
 
     async def _introspect_types(self, typeoids, timeout):
-        return await self.__execute(
+        if self._server_caps.jit:
+            try:
+                cfgrow, _ = await self.__execute(
+                    """
+                    SELECT
+                        current_setting('jit') AS cur,
+                        set_config('jit', 'off', false) AS new
+                    """,
+                    (),
+                    0,
+                    timeout,
+                    ignore_custom_codec=True,
+                )
+                jit_state = cfgrow[0]['cur']
+            except exceptions.UndefinedObjectError:
+                jit_state = 'off'
+        else:
+            jit_state = 'off'
+
+        result = await self.__execute(
             self._intro_query,
             (list(typeoids),),
             0,
@@ -469,6 +488,20 @@ async def _introspect_types(self, typeoids, timeout):
             ignore_custom_codec=True,
         )
 
+        if jit_state != 'off':
+            await self.__execute(
+                """
+                SELECT
+                    set_config('jit', $1, false)
+                """,
+                (jit_state,),
+                0,
+                timeout,
+                ignore_custom_codec=True,
+            )
+
+        return result
+
     async def _introspect_type(self, typename, schema):
         if (
             schema == 'pg_catalog'
@@ -2370,7 +2403,7 @@ class _ConnectionProxy:
 ServerCapabilities = collections.namedtuple(
     'ServerCapabilities',
     ['advisory_locks', 'notifications', 'plpgsql', 'sql_reset',
-     'sql_close_all'])
+     'sql_close_all', 'jit'])
 ServerCapabilities.__doc__ = 'PostgreSQL server capabilities.'
 
 
@@ -2382,6 +2415,7 @@ def _detect_server_capabilities(server_version, connection_settings):
         plpgsql = False
         sql_reset = True
         sql_close_all = False
+        jit = False
     elif hasattr(connection_settings, 'crdb_version'):
         # CockroachDB detected.
         advisory_locks = False
@@ -2389,6 +2423,7 @@ def _detect_server_capabilities(server_version, connection_settings):
         plpgsql = False
         sql_reset = False
         sql_close_all = False
+        jit = False
     elif hasattr(connection_settings, 'crate_version'):
         # CrateDB detected.
         advisory_locks = False
@@ -2396,6 +2431,7 @@ def _detect_server_capabilities(server_version, connection_settings):
         plpgsql = False
         sql_reset = False
         sql_close_all = False
+        jit = False
     else:
         # Standard PostgreSQL server assumed.
         advisory_locks = True
@@ -2403,13 +2439,15 @@ def _detect_server_capabilities(server_version, connection_settings):
         plpgsql = True
         sql_reset = True
         sql_close_all = True
+        jit = server_version >= (11, 0)
 
     return ServerCapabilities(
         advisory_locks=advisory_locks,
         notifications=notifications,
         plpgsql=plpgsql,
         sql_reset=sql_reset,
-        sql_close_all=sql_close_all
+        sql_close_all=sql_close_all,
+        jit=jit,
     )
 
 
diff --git a/tests/test_introspection.py b/tests/test_introspection.py
index 78561dd0..bf95537a 100644
--- a/tests/test_introspection.py
+++ b/tests/test_introspection.py
@@ -43,6 +43,12 @@ def tearDownClass(cls):
 
         super().tearDownClass()
 
+    @classmethod
+    def get_server_settings(cls):
+        settings = super().get_server_settings()
+        settings.pop('jit', None)
+        return settings
+
     def setUp(self):
         super().setUp()
         self.loop.run_until_complete(self._add_custom_codec(self.con))
@@ -124,7 +130,7 @@ async def test_introspection_no_stmt_cache_03(self):
         await self.con.fetchval(
             "SELECT $1::int[], '{foo}'".format(foo='a' * 10000), [1, 2])
 
-        self.assertEqual(apg_con._uid, old_uid + 1)
+        self.assertGreater(apg_con._uid, old_uid)
 
     async def test_introspection_sticks_for_ps(self):
         # Test that the introspected codec pipeline for a prepared