Add management command to precreate partitioned tables (#14076)

awx/main/management/commands/cleanup_jobs.py

@@ -17,10 +17,7 @@ from django.utils.timezone import now
 
 # AWX
 from awx.main.models import Job, AdHocCommand, ProjectUpdate, InventoryUpdate, SystemJob, WorkflowJob, Notification
-
-
-def unified_job_class_to_event_table_name(job_class):
-    return f'main_{job_class().event_class.__name__.lower()}'
+from awx.main.utils import unified_job_class_to_event_table_name
 
 
 def partition_table_name(job_class, dt):

awx/main/management/commands/precreate_partitions.py

@@ -0,0 +1,27 @@
+from django.utils.timezone import now
+from django.core.management.base import BaseCommand, CommandParser
+from datetime import timedelta
+from awx.main.utils.common import create_partition, unified_job_class_to_event_table_name
+from awx.main.models import Job, SystemJob, ProjectUpdate, InventoryUpdate, AdHocCommand
+
+
+class Command(BaseCommand):
+    """Command used to precreate database partitions to avoid pg_dump locks"""
+
+    def add_arguments(self, parser: CommandParser) -> None:
+        parser.add_argument('--count', dest='count', action='store', help='The amount of hours of partitions to create', type=int, default=1)
+
+    def _create_partitioned_tables(self, count):
+        tables = list()
+        for model in (Job, SystemJob, ProjectUpdate, InventoryUpdate, AdHocCommand):
+            tables.append(unified_job_class_to_event_table_name(model))
+        start = now()
+        while count > 0:
+            for table in tables:
+                create_partition(table, start)
+                print(f'Created partitions for {table} {start}')
+            start = start + timedelta(hours=1)
+            count -= 1
+
+    def handle(self, **options):
+        self._create_partitioned_tables(count=options.get('count'))

awx/main/utils/common.py

@@ -90,6 +90,7 @@ __all__ = [
     'get_event_partition_epoch',
     'cleanup_new_process',
     'log_excess_runtime',
+    'unified_job_class_to_event_table_name',
 ]
 
 
@@ -1219,3 +1220,7 @@ def log_excess_runtime(func_logger, cutoff=5.0, debug_cutoff=5.0, msg=None, add_
         return _new_func
 
     return log_excess_runtime_decorator
+
+
+def unified_job_class_to_event_table_name(job_class):
+    return f'main_{job_class().event_class.__name__.lower()}'