continueAfterTimeout is not supported by batch

Author: kislyuk

aurora_data_api/__init__.py

@@ -215,24 +215,23 @@ def _format_parameter_set(self, parameters):
 
     def _get_database_error(self, origin_error):
         error_msg = getattr(origin_error, "response", {}).get("Error", {}).get("Message", "")
-
-        # MySql error
-        if error_msg.startswith("Database error code"):
-            code, msg = (s.split(": ", 1)[1] for s in error_msg.split(". ", 1))
-            return DatabaseError(MySQLErrorCodes(int(code)), msg)
-        
-        # Postgresql error
-        elif error_msg.startswith("ERROR"):
-            msg, pos, code = postgresql_error_reg.match(error_msg).groups()
-            return DatabaseError(PostgreSQLErrorCodes(code), msg, pos)
-            
-        else:
-            return DatabaseError(origin_error)
+        try:
+            if error_msg.startswith("Database error code"):  # MySQL error
+                code, msg = (s.split(": ", 1)[1] for s in error_msg.split(". ", 1))
+                return DatabaseError(MySQLErrorCodes(int(code)), msg)
+            elif error_msg.startswith("ERROR: "):  # Postgresql error
+                msg, pos, code = postgresql_error_reg.match(error_msg).groups()
+                return DatabaseError(PostgreSQLErrorCodes(code), msg, pos)
+        except Exception:
+            pass
+        return DatabaseError(origin_error)
 
     def execute(self, operation, parameters=None):
         self._current_response, self._iterator, self._paging_state = None, None, None
         execute_statement_args = dict(self._prepare_execute_args(operation),
                                       includeResultMetadata=True)
+        if self._continue_after_timeout is not None:
+            execute_statement_args["continueAfterTimeout"] = self._continue_after_timeout
         if parameters:
             execute_statement_args["parameters"] = self._format_parameter_set(parameters)
         logger.debug("execute %s", reprlib.repr(operation.strip()))

aurora_data_api/postgresql_error_codes.py

@@ -251,4 +251,4 @@ class PostgreSQLErrorCodes(Enum):
     ER_ASSERT_FAILURE = 'P0004'
     ER_INTERNAL_ERROR = 'XX000'
     ER_DATA_CORRUPTED = 'XX001'
-    ER_INDEX_CORRUPTED = 'XX002'
+    ER_INDEX_CORRUPTED = 'XX002'

test/test.py

@@ -118,6 +118,8 @@ def test_iterators(self):
                         break
                     self.assertIn(len(fm), [1001, 46])
 
+    @unittest.skip("This test now fails because the API was changed to terminate and delete the transaction when the "
+                   "data returned by the statement exceeds the limit, making automated recovery impossible.")
     def test_pagination_backoff(self):
         if self.using_mysql:
             return
@@ -161,24 +163,27 @@ def test_rowcount(self):
     def test_continue_after_timeout(self):
         if os.environ.get("TEST_CONTINUE_AFTER_TIMEOUT", "False") != "True":
             self.skipTest("TEST_CONTINUE_AFTER_TIMEOUT env var is not 'True'")
-        
+
         if self.using_mysql:
             self.skipTest("Not implemented for MySQL")
 
         try:
             with aurora_data_api.connect(database=self.db_name) as conn, conn.cursor() as cur:
                 with self.assertRaisesRegex(conn._client.exceptions.ClientError, "StatementTimeoutException"):
-                    cur.execute("INSERT INTO aurora_data_api_test(name) SELECT 'continue_after_timeout' FROM (SELECT pg_sleep(50)) q")
+                    cur.execute(("INSERT INTO aurora_data_api_test(name) SELECT 'continue_after_timeout'"
+                                 "FROM (SELECT pg_sleep(50)) q"))
                 with self.assertRaisesRegex(aurora_data_api.DatabaseError, "current transaction is aborted"):
                     cur.execute("SELECT COUNT(*) FROM aurora_data_api_test WHERE name = 'continue_after_timeout'")
 
             with aurora_data_api.connect(database=self.db_name) as conn, conn.cursor() as cur:
                 cur.execute("SELECT COUNT(*) FROM aurora_data_api_test WHERE name = 'continue_after_timeout'")
                 self.assertEqual(cur.fetchone(), (0,))
-                
-            with aurora_data_api.connect(database=self.db_name, continue_after_timeout=True) as conn, conn.cursor() as cur:
+
+            with aurora_data_api.connect(database=self.db_name,
+                                         continue_after_timeout=True) as conn, conn.cursor() as cur:
                 with self.assertRaisesRegex(conn._client.exceptions.ClientError, "StatementTimeoutException"):
-                    cur.execute("INSERT INTO aurora_data_api_test(name) SELECT 'continue_after_timeout' FROM (SELECT pg_sleep(50)) q")
+                    cur.execute(("INSERT INTO aurora_data_api_test(name) SELECT 'continue_after_timeout' "
+                                 "FROM (SELECT pg_sleep(50)) q"))
                 cur.execute("SELECT COUNT(*) FROM aurora_data_api_test WHERE name = 'continue_after_timeout'")
                 self.assertEqual(cur.fetchone(), (1,))
         finally: