Add command: drop, to calculate partition drops based on retention periods

Author: jcjones

partitionmanager/cli.py

@@ -10,6 +10,8 @@
 import traceback
 import yaml
 
+import partitionmanager.database_helpers
+import partitionmanager.dropper
 import partitionmanager.migrate
 import partitionmanager.sql
 import partitionmanager.stats
@@ -121,10 +123,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:
@@ -465,6 +467,57 @@ def do_stats(conf, metrics=partitionmanager.stats.PrometheusMetrics()):
     return all_results
 
 
+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_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,72 @@
+"""
+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()
+
+    if not len(exact_time_result) == 1:
+        raise partitionmanager.types.NoExactTimeException("No exact timestamp result")
+    if not len(exact_time_result[0]) == 1:
+        raise partitionmanager.types.NoExactTimeException(
+            "Unexpected row count for the timestamp result"
+        )
+    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,110 @@
+"""
+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)
+
+        approx_size = 0
+        for a, b in zip(
+            next_partition.position.as_list(), partition.position.as_list()
+        ):
+            approx_size += a - b
+
+        try:
+            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
+                )
+            )
+
+            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)
+        except partitionmanager.types.NoExactTimeException:
+            log.warning(
+                "Couldn't determine exact times for %s.%s, it is probably droppable too.",
+                table,
+                partition,
+            )
+
+            results[partition.name] = {
+                "oldest_time": "unable to determine",
+                "youngest_time": "unable to determine",
+                "oldest_position": partition.position,
+                "youngest_position": next_partition.position,
+                "oldest_age": "unable to determine",
+                "youngest_age": "unable to determine",
+                "approx_size": approx_size,
+            }
+            droppable.append(partition)
+
+    if droppable:
+        results["drop_query"] = _drop_statement(table, droppable)
+
+    return results

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):
@@ -350,6 +350,9 @@ def __lt__(self, other):
                 return True
         return False
 
+    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
@@ -399,6 +402,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
@@ -609,3 +615,7 @@ class NoEmptyPartitionsAvailableException(Exception):
 
 class DatabaseCommandException(Exception):
     """Raised if the database command failed."""
+
+
+class NoExactTimeException(Exception):
+    """Raised if there's no exact time available for this partition."""

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,