Connection.prepare should not use statement cache

Author: vangheem

asyncpg/connection.py

@@ -264,19 +264,21 @@ async def executemany(self, command: str, args, *, timeout: float=None):
         self._check_open()
         return await self._executemany(command, args, timeout)
 
-    async def _get_statement(self, query, timeout, *, named: bool=False):
-        statement = self._stmt_cache.get(query)
-        if statement is not None:
-            return statement
-
-        # Only use the cache when:
-        #  * `statement_cache_size` is greater than 0;
-        #  * query size is less than `max_cacheable_statement_size`.
-        use_cache = self._stmt_cache.get_max_size() > 0
-        if (use_cache and
-                self._config.max_cacheable_statement_size and
-                len(query) > self._config.max_cacheable_statement_size):
-            use_cache = False
+    async def _get_statement(self, query, timeout, *, named: bool=False,
+                             use_cache: bool=True):
+        if use_cache:
+            statement = self._stmt_cache.get(query)
+            if statement is not None:
+                return statement
+
+            # Only use the cache when:
+            #  * `statement_cache_size` is greater than 0;
+            #  * query size is less than `max_cacheable_statement_size`.
+            use_cache = self._stmt_cache.get_max_size() > 0
+            if (use_cache and
+                    self._config.max_cacheable_statement_size and
+                    len(query) > self._config.max_cacheable_statement_size):
+                use_cache = False
 
         if use_cache or named:
             stmt_name = self._get_unique_id('stmt')
@@ -328,8 +330,12 @@ async def prepare(self, query, *, timeout=None):
 
         :return: A :class:`~prepared_stmt.PreparedStatement` instance.
         """
+        return await self._prepare(query, timeout=timeout, use_cache=False)
+
+    async def _prepare(self, query, *, timeout=None, use_cache: bool=False):
         self._check_open()
-        stmt = await self._get_statement(query, timeout, named=True)
+        stmt = await self._get_statement(query, timeout, named=True,
+                                         use_cache=use_cache)
         return prepared_stmt.PreparedStatement(self, query, stmt)
 
     async def fetch(self, query, *args, timeout=None) -> list:
@@ -645,7 +651,7 @@ async def copy_records_to_table(self, table_name, *, records,
         intro_query = 'SELECT {cols} FROM {tab} LIMIT 1'.format(
             tab=tabname, cols=col_list)
 
-        intro_ps = await self.prepare(intro_query)
+        intro_ps = await self._prepare(intro_query, use_cache=True)
 
         opts = '(FORMAT binary)'
 

tests/test_cache_invalidation.py

@@ -151,7 +151,8 @@ async def test_prepared_type_cache_invalidation(self):
 
         try:
             await self.con.execute('INSERT INTO tab1 VALUES (1, (2, 3))')
-            prep = await self.con.prepare('SELECT * FROM tab1')
+            prep = await self.con._prepare('SELECT * FROM tab1',
+                                           use_cache=True)
             result = await prep.fetchrow()
             self.assertEqual(result, (1, (2, 3)))
 
@@ -170,7 +171,8 @@ async def test_prepared_type_cache_invalidation(self):
                             'the prepared statement is closed'):
                         await prep.fetchrow()
 
-                    prep = await self.con.prepare('SELECT * FROM tab1')
+                    prep = await self.con._prepare('SELECT * FROM tab1',
+                                                   use_cache=True)
                     # The second PS must be correct (cache was dropped):
                     result = await prep.fetchrow()
                     self.assertEqual(result, (1, (2, 3, None)))
@@ -183,7 +185,8 @@ async def test_prepared_type_cache_invalidation(self):
                 await prep.fetchrow()
 
             # Reprepare it again after dropping cache.
-            prep = await self.con.prepare('SELECT * FROM tab1')
+            prep = await self.con._prepare('SELECT * FROM tab1',
+                                           use_cache=True)
             # This is now OK, the cache is filled after being dropped.
             result = await prep.fetchrow()
             self.assertEqual(result, (1, (2, 3)))

tests/test_prepare.py

@@ -153,7 +153,7 @@ async def test_prepare_10_stmt_lru(self):
 
         stmts = []
         for i in range(iter_max):
-            s = await self.con.prepare(query.format(i))
+            s = await self.con._prepare(query.format(i), use_cache=True)
             self.assertEqual(await s.fetchval(), i)
             stmts.append(s)
 
@@ -207,7 +207,7 @@ async def test_prepare_11_stmt_gc(self):
         # The prepared statement that we'll create will be GCed
         # right await.  However, its state should be still in
         # in the statements LRU cache.
-        await self.con.prepare('select 1')
+        await self.con._prepare('select 1', use_cache=True)
         gc.collect()
 
         self.assertEqual(len(cache), 1)
@@ -224,12 +224,12 @@ async def test_prepare_12_stmt_gc(self):
         self.assertEqual(len(cache), 0)
         self.assertEqual(len(self.con._stmts_to_close), 0)
 
-        stmt = await self.con.prepare('select 100000000')
+        stmt = await self.con._prepare('select 100000000', use_cache=True)
         self.assertEqual(len(cache), 1)
         self.assertEqual(len(self.con._stmts_to_close), 0)
 
         for i in range(cache_max):
-            await self.con.prepare('select {}'.format(i))
+            await self.con._prepare('select {}'.format(i), use_cache=True)
 
         self.assertEqual(len(cache), cache_max)
         self.assertEqual(len(self.con._stmts_to_close), 0)
@@ -293,7 +293,7 @@ async def test_prepare_15_stmt_gc_cache_disabled(self):
         # Disable cache
         cache.set_max_size(0)
 
-        stmt = await self.con.prepare('select 100000000')
+        stmt = await self.con._prepare('select 100000000', use_cache=True)
         self.assertEqual(len(cache), 0)
         self.assertEqual(len(self.con._stmts_to_close), 0)
 
@@ -305,7 +305,7 @@ async def test_prepare_15_stmt_gc_cache_disabled(self):
         self.assertEqual(len(self.con._stmts_to_close), 1)
 
         # Next "prepare" call will trigger a cleanup
-        stmt = await self.con.prepare('select 1')
+        stmt = await self.con._prepare('select 1', use_cache=True)
         self.assertEqual(len(cache), 0)
         self.assertEqual(len(self.con._stmts_to_close), 0)
 
@@ -468,25 +468,25 @@ async def test_prepare_24_max_lifetime(self):
         self.assertEqual(cache.get_max_lifetime(), 142)
         cache.set_max_lifetime(1)
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         state = s._state
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         self.assertIs(s._state, state)
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         self.assertIs(s._state, state)
 
         await asyncio.sleep(1, loop=self.loop)
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         self.assertIsNot(s._state, state)
 
     @tb.with_connection_options(max_cached_statement_lifetime=0.5)
     async def test_prepare_25_max_lifetime_reset(self):
         cache = self.con._stmt_cache
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         state = s._state
 
         # Disable max_lifetime
@@ -495,20 +495,20 @@ async def test_prepare_25_max_lifetime_reset(self):
         await asyncio.sleep(1, loop=self.loop)
 
         # The statement should still be cached (as we disabled the timeout).
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         self.assertIs(s._state, state)
 
     @tb.with_connection_options(max_cached_statement_lifetime=0.5)
     async def test_prepare_26_max_lifetime_max_size(self):
         cache = self.con._stmt_cache
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         state = s._state
 
         # Disable max_lifetime
         cache.set_max_size(0)
 
-        s = await self.con.prepare('SELECT 1')
+        s = await self.con._prepare('SELECT 1', use_cache=True)
         self.assertIsNot(s._state, state)
 
         # Check that nothing crashes after the initial timeout
@@ -518,12 +518,12 @@ async def test_prepare_26_max_lifetime_max_size(self):
     async def test_prepare_27_max_cacheable_statement_size(self):
         cache = self.con._stmt_cache
 
-        await self.con.prepare('SELECT 1')
+        await self.con._prepare('SELECT 1', use_cache=True)
         self.assertEqual(len(cache), 1)
 
         # Test that long and explicitly created prepared statements
         # are not cached.
-        await self.con.prepare("SELECT \'" + "a" * 50 + "\'")
+        await self.con._prepare("SELECT \'" + "a" * 50 + "\'", use_cache=True)
         self.assertEqual(len(cache), 1)
 
         # Test that implicitly created long prepared statements
@@ -532,7 +532,7 @@ async def test_prepare_27_max_cacheable_statement_size(self):
         self.assertEqual(len(cache), 1)
 
         # Test that short prepared statements can still be cached.
-        await self.con.prepare('SELECT 2')
+        await self.con._prepare('SELECT 2', use_cache=True)
         self.assertEqual(len(cache), 2)
 
     async def test_prepare_28_max_args(self):
@@ -593,3 +593,10 @@ async def test_prepare_31_pgbouncer_note(self):
             self.assertTrue('pgbouncer' in e.hint)
         else:
             self.fail('InvalidSQLStatementNameError not raised')
+
+    async def test_prepare_does_not_use_cache(self):
+        cache = self.con._stmt_cache
+
+        # prepare with disabled cache
+        await self.con.prepare('select 1')
+        self.assertEqual(len(cache), 0)