Fix missing exception catch in `create_partition` (#15691)

Fix error creating partition due to uncaught exception

the primary fix is to simply add an exception class
  to those caught in the except block

This also adds live tests for the general scenario
  although this does not hit the new exception type

2 files changed, 41 insertions, 8 deletions

diff --git a/awx/main/tests/live/tests/test_partitions.py b/awx/main/tests/live/tests/test_partitions.py
new file mode 100644
index 0000000000..f395eda8e1
--- /dev/null
+++ b/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))
diff --git a/awx/main/utils/common.py b/awx/main/utils/common.py
index 7d4439f51a..8b482f6ed0 100644
--- a/awx/main/utils/common.py
+++ b/awx/main/utils/common.py
@@ -1138,6 +1138,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:
@@ -1154,13 +1165,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); '
@@ -1172,9 +1178,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))
@@ -1183,6 +1191,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