Fix missing exception catch in create_partition

awx/main/tests/live/tests/test_partitions.py

@@ -0,0 +1,23 @@
+from datetime import timedelta
+
+from django.utils.timezone import now
+from django.db import connection
+
+from awx.main.utils.common import create_partition, table_exists
+
+
+def test_table_when_it_exists():
+    with connection.cursor() as cursor:
+        assert table_exists(cursor, 'main_job')
+
+
+def test_table_when_it_does_not_exists():
+    with connection.cursor() as cursor:
+        assert not table_exists(cursor, 'main_not_a_table_check')
+
+
+def test_create_partition_race_condition(mocker):
+    mocker.patch('awx.main.utils.common.table_exists', return_value=False)
+
+    create_partition('main_jobevent', start=now() - timedelta(days=2))
+    create_partition('main_jobevent', start=now() - timedelta(days=2))

awx/main/utils/common.py

@@ -1096,6 +1096,17 @@ def deepmerge(a, b):
         return b
 
 
+def table_exists(cursor, table_name):
+    cursor.execute(f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = '{table_name}');")
+    row = cursor.fetchone()
+    if row is not None:
+        for val in row:  # should only have 1
+            if val is True:
+                logger.debug(f'Event partition table {table_name} already exists')
+                return True
+    return False
+
+
 def create_partition(tblname, start=None):
     """Creates new partition table for events.  By default it covers the current hour."""
     if start is None:
@@ -1112,13 +1123,8 @@ def create_partition(tblname, start=None):
     try:
         with transaction.atomic():
             with connection.cursor() as cursor:
-                cursor.execute(f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = '{tblname}_{partition_label}');")
-                row = cursor.fetchone()
-                if row is not None:
-                    for val in row:  # should only have 1
-                        if val is True:
-                            logger.debug(f'Event partition table {tblname}_{partition_label} already exists')
-                            return
+                if table_exists(cursor, f"{tblname}_{partition_label}"):
+                    return
 
                 cursor.execute(
                     f'CREATE TABLE {tblname}_{partition_label} (LIKE {tblname} INCLUDING DEFAULTS INCLUDING CONSTRAINTS); '
@@ -1130,9 +1136,11 @@ def create_partition(tblname, start=None):
         cause = e.__cause__
         if cause and hasattr(cause, 'sqlstate'):
             sqlstate = cause.sqlstate
+            if sqlstate is None:
+                raise
             sqlstate_cls = psycopg.errors.lookup(sqlstate)
 
-            if psycopg.errors.DuplicateTable == sqlstate_cls or psycopg.errors.UniqueViolation == sqlstate_cls:
+            if sqlstate_cls in (psycopg.errors.DuplicateTable, psycopg.errors.DuplicateObject, psycopg.errors.UniqueViolation):
                 logger.info(f'Caught known error due to partition creation race: {e}')
             else:
                 logger.error('SQL Error state: {} - {}'.format(sqlstate, sqlstate_cls))
@@ -1141,6 +1149,8 @@ def create_partition(tblname, start=None):
         cause = e.__cause__
         if cause and hasattr(cause, 'sqlstate'):
             sqlstate = cause.sqlstate
+            if sqlstate is None:
+                raise
             sqlstate_str = psycopg.errors.lookup(sqlstate)
             logger.error('SQL Error state: {} - {}'.format(sqlstate, sqlstate_str))
         raise