From 6976ac9273ee921a23e12260505d78b83d20639a Mon Sep 17 00:00:00 2001
From: Gabriel Muniz <gmuniz@redhat.com>
Date: Mon, 5 Jun 2023 14:20:53 -0400
Subject: [PATCH] Add management command to precreate partitioned tables
 (#14076)

---
 awx/main/management/commands/cleanup_jobs.py  |  5 +---
 .../commands/precreate_partitions.py          | 27 +++++++++++++++++++
 awx/main/utils/common.py                      |  5 ++++
 3 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 awx/main/management/commands/precreate_partitions.py

diff --git a/awx/main/management/commands/cleanup_jobs.py b/awx/main/management/commands/cleanup_jobs.py
index b0de7ae7a3..bd87e1e198 100644
--- a/awx/main/management/commands/cleanup_jobs.py
+++ b/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):
diff --git a/awx/main/management/commands/precreate_partitions.py b/awx/main/management/commands/precreate_partitions.py
new file mode 100644
index 0000000000..e2337ad401
--- /dev/null
+++ b/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'))
diff --git a/awx/main/utils/common.py b/awx/main/utils/common.py
index f874ff2a3f..896f37d779 100644
--- a/awx/main/utils/common.py
+++ b/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()}'