cleanup old logs improvements

Author: OriHoch

hasadna_k8s/storage/cleanup_old_logs.py

@@ -15,19 +15,37 @@ def delete_old_files(root_path, cutoff_date, dry_run=True):
                 if dry_run:
                     print(f'remove: {file_path}')
                 else:
-                    os.remove(file_path)
+                    try:
+                        os.remove(file_path)
+                    except:
+                        print(f'WARNING: could not remove file {file_path}')
         if not os.listdir(root) and root != root_path:
             if dry_run:
                 print(f'rmtree: {root}')
             else:
-                shutil.rmtree(root)
+                try:
+                    shutil.rmtree(root)
+                except:
+                    print(f'WARNING: could not remove directory {root}')
+
+
+def find_log_paths_airflow_logs(root):
+    if os.path.basename(root) == 'logs':
+        parent = os.path.dirname(root)
+        if os.path.exists(f'{parent}/airflow.cfg'):
+            return True
+    return False
 
 
 def find_log_paths(root_path, log_path_prefixes):
     for root, dirs, files in os.walk(root_path):
         for log_path_prefix in log_path_prefixes:
-            if root.endswith(log_path_prefix):
-                yield root
+            if log_path_prefix == 'airflow:logs':
+                if find_log_paths_airflow_logs(root):
+                    yield root
+            else:
+                if root.endswith(log_path_prefix):
+                    yield root
 
 
 def main(path, dry_run=True, log_path_prefixes=None):

tests/old_logs_paths/airflow-home/logs/.gitkeep

@@ -0,0 +1,35 @@
+import os
+from glob import glob
+from hasadna_k8s.storage.cleanup_old_logs import find_log_paths, main as cleanup_old_logs
+
+
+def test_find_log_paths():
+    base = os.path.dirname(__file__)
+    assert list(find_log_paths(os.path.dirname(__file__), ["airflow-home/logs", "airflow:logs"])) == [
+        f'{base}/old_logs_paths/foobar/logs',
+        f'{base}/old_logs_paths/airflow-home/logs',
+    ]
+
+
+def test_cleanup_old_logs():
+    base = os.path.dirname(__file__)
+    for f in glob(f'{base}/old_logs_paths/**/.gitkeep', recursive=True):
+        os.utime(f, None)
+    for fn in [
+        f'{base}/old_logs_paths/airflow-home/logs/old_log.log',
+        f'{base}/old_logs_paths/foobar/logs/old_log.log',
+        f'{base}/old_logs_paths/somethingelse/old_log.log',
+    ]:
+        with open(fn, 'w') as f:
+            f.write('test log content')
+        # Set modification time to 40 days ago
+        old_time = os.path.getmtime(fn) - (40 * 24 * 60 * 60)
+        os.utime(fn, (old_time, old_time))
+    cleanup_old_logs(base, dry_run=False, log_path_prefixes=["airflow-home/logs", "airflow:logs"])
+    for fn in [
+        f'{base}/old_logs_paths/airflow-home/logs/old_log.log',
+        f'{base}/old_logs_paths/foobar/logs/old_log.log',
+    ]:
+        assert not os.path.exists(fn)
+    assert os.path.exists(f'{base}/old_logs_paths/somethingelse/old_log.log')
+    os.remove(f'{base}/old_logs_paths/somethingelse/old_log.log')