_handle_partitions uses get_or_create_partition (#3330)

Co-authored-by: Michael Skarbek <[email protected]>

koku/koku/pg_partition.py

@@ -1806,7 +1806,7 @@ def _move_partition_data(target_partition, source_partition, conditions):
     return (cp_recs, dl_recs)
 
 
-def get_or_create_partition(part_rec):
+def get_or_create_partition(part_rec, _default_partition=None):
     """
     Get or create an existing partition.
     Creating:
@@ -1837,7 +1837,11 @@ def get_or_create_partition(part_rec):
 
     with schema_context(part_rec["schema_name"]):
         # Find or create the default partition
-        default_partition, default_created = _get_or_create_default_partition(part_rec)
+        if _default_partition:
+            default_partition, default_created = _default_partition, False
+        else:
+            default_partition, default_created = _get_or_create_default_partition(part_rec)
+
         if part_rec.get("partition_parameters", {}).get("default") or "default" in part_rec.get("table_name"):
             return default_partition, default_created
 
@@ -1885,14 +1889,17 @@ def _handle_partitions(self, schema_name, table_names, start_date, end_date):  #
             table_names = [table_names]
 
         for table_name in table_names:
-            tmplpart = PartitionedTable.objects.filter(
-                schema_name=schema_name, partition_of_table_name=table_name, partition_type=PartitionedTable.RANGE
+            default_part = PartitionedTable.objects.filter(
+                schema_name=schema_name,
+                partition_of_table_name=table_name,
+                partition_type=PartitionedTable.RANGE,
+                partition_parameters__default=True,
             ).first()
-            if tmplpart:
+            if default_part:
                 partition_start = start_date.replace(day=1)
                 month_interval = relativedelta(months=1)
                 needed_partition = None
-                partition_col = tmplpart.partition_col
+                partition_col = default_part.partition_col
                 newpart_vals = dict(
                     schema_name=schema_name,
                     table_name=None,
@@ -1913,12 +1920,7 @@ def _handle_partitions(self, schema_name, table_names, start_date, end_date):  #
                     newpart_vals["partition_parameters"]["from"] = str(needed_partition)
                     newpart_vals["partition_parameters"]["to"] = str(needed_partition + month_interval)
                     # Successfully creating a new record will also create the partition
-                    newpart, created = PartitionedTable.objects.get_or_create(
-                        defaults=newpart_vals,
-                        schema_name=schema_name,
-                        partition_of_table_name=table_name,
-                        table_name=partition_name,
-                    )
+                    newpart, created = get_or_create_partition(newpart_vals, _default_partition=default_part)
                     LOG.debug(f"partition = {newpart}")
                     LOG.debug(f"created = {created}")
                     if created:

koku/koku/test_pg_partition.py

@@ -110,6 +110,28 @@ def test_input_scrub(self):
 
             dp.delete()
 
+    def test_call_with_default(self):
+        """Test input dict value scrubbing"""
+        with schema_context(self.SCHEMA_NAME):
+            part_rec = {
+                "schema_name": self.SCHEMA_NAME,
+                "table_name": self.PARTITIONED_TABLE_NAME + "_default",
+                "partition_of_table_name": self.PARTITIONED_TABLE_NAME,
+                "partition_parameters": {"default": True},
+                "active": True,
+            }
+            default_partition, _ = PartitionedTable.objects.get_or_create(
+                schema_name=self.SCHEMA_NAME,
+                table_name=part_rec["table_name"],
+                partition_parameters__default=True,
+                defaults=part_rec,
+            )
+            dp, dc = ppart.get_or_create_partition(part_rec, _default_partition=default_partition)
+            self.assertTrue(dp == default_partition, "get_or_create_partition did not return the default partition")
+            self.assertFalse(dc, "get_or_create_partition with default says it created a default partition")
+
+            dp.delete()
+
     def test_create_default_partition(self):
         """Test that a default partition can be explicitly created"""
         with schema_context(self.SCHEMA_NAME):