From 3725b4bbcb7d4d3ef5d5903d819f65e600105330 Mon Sep 17 00:00:00 2001
From: Joar Wandborg <joar@wandborg.se>
Date: Fri, 19 May 2017 16:28:29 +0200
Subject: [PATCH] Add wal-e threshold logging (#440)

Add error log for threshold, so that you can see if thresholds were the reason for ExitCode.FAIL.
---
 patroni/scripts/wale_restore.py | 72 ++++++++++++++++++++++++++++++---
 requirements.txt                |  1 +
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/patroni/scripts/wale_restore.py b/patroni/scripts/wale_restore.py
index aa36dca3..f5490ffb 100755
--- a/patroni/scripts/wale_restore.py
+++ b/patroni/scripts/wale_restore.py
@@ -25,6 +25,7 @@
 #               restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
 import csv
 from collections import namedtuple
+import humanize
 import logging
 import os
 from enum import IntEnum
@@ -65,6 +66,22 @@ def get_major_version(data_dir):
     return 0.0
 
 
+def repr_size(n_bytes):
+    return humanize.naturalsize(n_bytes, binary=True)
+
+
+def size_as_bytes(size_, prefix):
+    si_prefixes = ['K', 'M', 'G', 'T'', P', 'E', 'Z', 'Y']
+
+    prefix = prefix.upper()
+
+    assert prefix in si_prefixes
+
+    exponent = si_prefixes.index(prefix) + 1
+
+    return size_ * 1024.0 ** -exponent
+
+
 WALEConfig = namedtuple(
     'WALEConfig',
     [
@@ -139,7 +156,7 @@ class WALERestore(object):
         """ determine whether it makes sense to use S3 and not pg_basebackup """
 
         threshold_megabytes = self.wal_e.threshold_mb
-        threshold_backup_size_percentage = self.wal_e.threshold_pct
+        threshold_percent = self.wal_e.threshold_pct
 
         try:
             cmd = self.wal_e.cmd + ['backup-list', '--detail', 'LATEST']
@@ -169,7 +186,7 @@ class WALERestore(object):
             return None
 
         try:
-            backup_size = backup_info['expanded_size_bytes']
+            backup_size = int(backup_info['expanded_size_bytes'])
             backup_start_segment = backup_info['wal_segment_backup_start']
             backup_start_offset = backup_info['wal_segment_offset_backup_start']
         except KeyError:
@@ -187,7 +204,7 @@ class WALERestore(object):
         # construct the LSN from the segment and offset
         backup_start_lsn = '{0}/{1}'.format(lsn_segment, lsn_offset)
 
-        diff_in_bytes = int(backup_size)
+        diff_in_bytes = backup_size
         attempts_no = 0
         while True:
             if self.master_connection:
@@ -226,8 +243,53 @@ class WALERestore(object):
 
         # if the size of the accumulated WAL segments is more than a certan percentage of the backup size
         # or exceeds the pre-determined size - pg_basebackup is chosen instead.
-        return (diff_in_bytes < int(threshold_megabytes) * 1048576) and\
-            (diff_in_bytes < int(backup_size) * float(threshold_backup_size_percentage) / 100)
+        is_size_thresh_ok = diff_in_bytes < int(threshold_megabytes) * 1048576
+        threshold_pct_bytes = backup_size * threshold_percent / 100.0
+        is_percentage_thresh_ok = float(diff_in_bytes) < int(threshold_pct_bytes)
+        are_thresholds_ok = is_size_thresh_ok and is_percentage_thresh_ok
+
+        class Size(object):
+            def __init__(self, n_bytes, prefix=None):
+                self.n_bytes = n_bytes
+                self.prefix = prefix
+
+            def __repr__(self):
+                if self.prefix is not None:
+                    n_bytes = size_as_bytes(self.n_bytes, self.prefix)
+                else:
+                    n_bytes = self.n_bytes
+                return repr_size(n_bytes)
+
+        class HumanContext(object):
+            def __init__(self, items):
+                self.items = items
+
+            def __repr__(self):
+                return ', '.join('{}={!r}'.format(key, value)
+                                 for key, value in self.items)
+
+        human_context = HumanContext([
+            ('threshold_size', Size(threshold_megabytes, 'M')),
+            ('threshold_percent', threshold_percent),
+            ('threshold_percent_size', Size(threshold_pct_bytes)),
+            ('backup_size', Size(backup_size)),
+            ('backup_diff', Size(diff_in_bytes)),
+            ('is_size_thresh_ok', is_size_thresh_ok),
+            ('is_percentage_thresh_ok', is_percentage_thresh_ok),
+        ])
+
+        if not are_thresholds_ok:
+            logger.error(
+                'wal-e backup size diff is over threshold, falling back '
+                'to other means of restore. %r',
+                human_context
+            )
+        else:
+            logger.info(
+                'Thresholds are OK, using wal-e basebackup. %r',
+                human_context
+            )
+        return are_thresholds_ok
 
     def fix_subdirectory_path_if_broken(self, dirname):
         # in case it is a symlink pointing to a non-existing location, remove it and create the actual directory
diff --git a/requirements.txt b/requirements.txt
index 99c4c80a..2fd2fa51 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,4 @@ tzlocal
 python-dateutil
 enum34
 cdiff
+humanize==0.5.1
\ No newline at end of file