Initial Commit: Retention-based Partition Dropping

Lacking tests, but can determine which partitions are wholly outside
of the configured `retention_period` for a given table and emit
commands to drop them.

Author: jcjones

partitionmanager/cli.py

@@ -10,6 +10,7 @@
 import traceback
 import yaml
 
+import partitionmanager.dropper
 import partitionmanager.migrate
 import partitionmanager.sql
 import partitionmanager.stats
@@ -121,10 +122,10 @@ def from_yaml_file(self, file):
             for key in data["tables"]:
                 tab = partitionmanager.types.Table(key)
                 tabledata = data["tables"][key]
-                if isinstance(tabledata, dict) and "retention" in tabledata:
-                    tab.set_retention(
+                if isinstance(tabledata, dict) and "retention_period" in tabledata:
+                    tab.set_retention_period(
                         partitionmanager.types.timedelta_from_dict(
-                            tabledata["retention"]
+                            tabledata["retention_period"]
                         )
                     )
                 if isinstance(tabledata, dict) and "partition_period" in tabledata:
@@ -289,6 +290,20 @@ def migrate_cmd(args):
 MIGRATE_PARSER.set_defaults(func=migrate_cmd)
 
 
+def drop_cmd(args):
+    """Calculates drop.
+
+    Helper for argparse.
+    """
+    conf = config_from_args(args)
+
+    return do_find_drops_for_tables(conf)
+
+
+DROP_PARSER = SUBPARSERS.add_parser("drop", help="drop old partitions")
+DROP_PARSER.set_defaults(func=drop_cmd)
+
+
 def do_partition(conf):
     """Produces SQL statements to manage partitions per the supplied configuration.
 
@@ -463,6 +478,46 @@ def do_stats(conf, metrics=partitionmanager.stats.PrometheusMetrics()):
     return all_results
 
 
+def do_find_drops_for_tables(conf):
+    all_results = dict()
+    for table in conf.tables:
+        log = logging.getLogger(f"do_find_drops_for_tables:{table.name}")
+
+        if not table.has_date_query:
+            log.debug(f"Cannot process {table}: no date query specified")
+            continue
+
+        if not table.retention_period:
+            log.debug(f"Cannot process {table}: no retention specified")
+            continue
+
+        try:
+            table_problems = pm_tap.get_table_compatibility_problems(conf.dbcmd, table)
+            if table_problems:
+                log.debug(f"Cannot process {table}: {table_problems}")
+                continue
+
+            map_data = pm_tap.get_partition_map(conf.dbcmd, table)
+            current_position = partitionmanager.database_helpers.get_position_of_table(
+                conf.dbcmd, table, map_data
+            )
+
+            droppable = partitionmanager.dropper.get_droppable_partitions(
+                conf.dbcmd,
+                map_data["partitions"],
+                current_position,
+                conf.curtime,
+                table,
+            )
+
+            all_results[table.name] = droppable
+        except Exception as e:
+            log.warning(f"Error processing table {table.name}")
+            raise e
+
+    return all_results
+
+
 def main():
     """Start here."""
     args = PARSER.parse_args()

partitionmanager/database_helpers.py

@@ -0,0 +1,68 @@
+"""
+Helper functions for database operations
+"""
+
+from datetime import datetime, timezone
+import logging
+
+import partitionmanager.table_append_partition as pm_tap
+import partitionmanager.types
+
+
+def get_position_of_table(database, table, map_data):
+    """ Returns a Position of the table at the current moment. """
+
+    pos_list = pm_tap.get_current_positions(database, table, map_data["range_cols"])
+
+    cur_pos = partitionmanager.types.Position()
+    cur_pos.set_position([pos_list[col] for col in map_data["range_cols"]])
+
+    return cur_pos
+
+
+def calculate_exact_timestamp_via_query(database, table, position_partition):
+    """ Calculates the exact timestamp of a PositionPartition.
+
+    raises ValueError if the position is incalculable
+    """
+
+    log = logging.getLogger(f"calculate_exact_timestamp_via_query:{table.name}")
+
+    if not table.has_date_query:
+        raise ValueError("Table has no defined date query")
+
+    if not isinstance(position_partition, partitionmanager.types.PositionPartition):
+        raise ValueError("Only PositionPartitions are supported")
+
+    if len(position_partition.position) != 1:
+        raise ValueError(
+            "This method is only valid for single-column partitions right now"
+        )
+    arg = position_partition.position.as_sql_input()[0]
+
+    sql_select_cmd = table.earliest_utc_timestamp_query.get_statement_with_argument(arg)
+    log.debug(
+        "Executing %s to derive partition %s at position %s",
+        sql_select_cmd,
+        position_partition.name,
+        position_partition.position,
+    )
+
+    start = datetime.now()
+    exact_time_result = database.run(sql_select_cmd)
+    end = datetime.now()
+
+    assert len(exact_time_result) == 1
+    assert len(exact_time_result[0]) == 1
+    for key, value in exact_time_result[0].items():
+        exact_time = datetime.fromtimestamp(value, tz=timezone.utc)
+        break
+
+    log.debug(
+        "Exact time of %s returned for %s at position %s, query took %s",
+        exact_time,
+        position_partition.name,
+        position_partition.position,
+        (end - start),
+    )
+    return exact_time

partitionmanager/dropper.py

@@ -0,0 +1,88 @@
+"""
+Determine which partitions can be dropped.
+"""
+
+import logging
+
+import partitionmanager.types
+import partitionmanager.tools
+
+
+def _drop_statement(table, partition_list):
+    """ Generate an ALTER TABLE statement to drop these partitions. """
+
+    log = logging.getLogger("get_droppable_partitions")
+
+    partitions = ",".join(map(lambda x: f"`{x.name}`", partition_list))
+
+    alter_cmd = (
+        f"ALTER TABLE `{table.name}` " f"DROP PARTITION IF EXISTS {partitions} ;"
+    )
+
+    log.debug("Yielding %s", alter_cmd)
+
+    return alter_cmd
+
+
+def get_droppable_partitions(
+    database, partitions, current_position, current_timestamp, table
+):
+    """Return a dictionary of partitions which can be dropped and why."""
+    log = logging.getLogger("get_droppable_partitions")
+    results = {}
+    droppable = []
+
+    if not table.retention_period:
+        raise ValueError(f"{table.name} does not have a retention period set")
+
+    if not partitions:
+        return results
+
+    for partition, next_partition in partitionmanager.tools.pairwise(partitions):
+        if next_partition >= current_position:
+            log.debug(
+                "Stopping at %s because current position %s indicates "
+                "subsequent partition is empty",
+                partition,
+                current_position,
+            )
+            break
+
+        if isinstance(next_partition, partitionmanager.types.MaxValuePartition):
+            log.debug("Stopping at %s because we can't handle MaxValuePartitions.")
+            break
+
+        assert isinstance(next_partition, partitionmanager.types.PositionPartition)
+
+        start_time = partitionmanager.database_helpers.calculate_exact_timestamp_via_query(
+            database, table, partition
+        )
+        end_time = partitionmanager.database_helpers.calculate_exact_timestamp_via_query(
+            database, table, next_partition
+        )
+
+        approx_size = 0
+        for a, b in zip(
+            next_partition.position.as_list(), partition.position.as_list()
+        ):
+            approx_size += a - b
+
+        oldest_age = current_timestamp - start_time
+        youngest_age = current_timestamp - end_time
+
+        if youngest_age > table.retention_period:
+            results[partition.name] = {
+                "oldest_time": f"{start_time}",
+                "youngest_time": f"{end_time}",
+                "oldest_position": partition.position,
+                "youngest_position": next_partition.position,
+                "oldest_age": f"{oldest_age}",
+                "youngest_age": f"{youngest_age}",
+                "approx_size": approx_size,
+            }
+            droppable.append(partition)
+
+    if droppable:
+        results["drop_query"] = _drop_statement(table, droppable)
+
+    return results

partitionmanager/table_append_partition.py

@@ -2,11 +2,12 @@
 Design and perform partition management.
 """
 
-from datetime import datetime, timedelta, timezone
+from datetime import timedelta
 import logging
 import operator
 import re
 
+import partitionmanager.database_helpers
 import partitionmanager.types
 import partitionmanager.tools
 
@@ -417,48 +418,15 @@ def _get_rate_partitions_with_queried_timestamps(
 
     The partitions' timestamps are explicitly queried.
     """
-    log = logging.getLogger(
-        f"_get_rate_partitions_with_queried_timestamps:{table.name}"
-    )
 
     if not table.has_date_query:
         raise ValueError("Table has no defined date query")
 
     instant_partitions = list()
 
     for partition in partition_list:
-        if len(partition.position) != 1:
-            raise ValueError(
-                "This method is only valid for single-column partitions right now"
-            )
-        arg = partition.position.as_sql_input()[0]
-
-        sql_select_cmd = table.earliest_utc_timestamp_query.get_statement_with_argument(
-            arg
-        )
-        log.debug(
-            "Executing %s to derive partition %s at position %s",
-            sql_select_cmd,
-            partition.name,
-            partition.position,
-        )
-
-        start = datetime.now()
-        exact_time_result = database.run(sql_select_cmd)
-        end = datetime.now()
-
-        assert len(exact_time_result) == 1
-        assert len(exact_time_result[0]) == 1
-        for key, value in exact_time_result[0].items():
-            exact_time = datetime.fromtimestamp(value, tz=timezone.utc)
-            break
-
-        log.debug(
-            "Exact time of %s returned for %s at position %s, query took %s",
-            exact_time,
-            partition.name,
-            partition.position,
-            (end - start),
+        exact_time = partitionmanager.database_helpers.calculate_exact_timestamp_via_query(
+            database, table, partition
         )
 
         instant_partitions.append(

partitionmanager/types.py

@@ -30,17 +30,17 @@ class Table:
 
     def __init__(self, name):
         self.name = SqlInput(name)
-        self.retention = None
+        self.retention_period = None
         self.partition_period = None
         self.earliest_utc_timestamp_query = None
 
-    def set_retention(self, ret):
+    def set_retention_period(self, ret):
         """
         Sets the retention period as a timedelta for this table
         """
         if not isinstance(ret, timedelta):
             raise ValueError("Must be a timedelta")
-        self.retention = ret
+        self.retention_period = ret
         return self
 
     def set_partition_period(self, dur):
@@ -348,6 +348,9 @@ def __lt__(self, other):
                 return False
         return True
 
+    def __ge__(self, other):
+        return not self < other
+
     def __eq__(self, other):
         if isinstance(other, PositionPartition):
             return self.name == other.name and self._position == other.position
@@ -397,6 +400,9 @@ def __lt__(self, other):
             return False
         return ValueError()
 
+    def __ge__(self, other):
+        return not self < other
+
     def __eq__(self, other):
         if isinstance(other, MaxValuePartition):
             return self.name == other.name and self._count == other.num_columns

partitionmanager/types_test.py

@@ -146,7 +146,7 @@ def test_table(self):
         self.assertEqual(type(Table("name").name), SqlInput)
 
         t = Table("t")
-        self.assertEqual(None, t.retention)
+        self.assertEqual(None, t.retention_period)
 
         self.assertEqual(
             Table("a").set_partition_period(timedelta(days=9)).partition_period,