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 bdd7465b90..dfc93b58ad 100644 --- a/awx/main/utils/common.py +++ b/awx/main/utils/common.py @@ -1127,6 +1127,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: @@ -1143,13 +1154,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); ' @@ -1161,9 +1167,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)) @@ -1172,6 +1180,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