From 7df278ef1fb9de3e772b8ba761e53d24cb6a33f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Mon, 18 Apr 2016 11:03:21 +0200
Subject: [PATCH] Added manage.py cmd to refresh all file links of a certain
 backend.

---
 pillar/application/modules/file_storage.py | 28 ++++++++++++++++++++++
 pillar/manage.py                           | 16 +++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/pillar/application/modules/file_storage.py b/pillar/application/modules/file_storage.py
index 4c2b94b1..cecf1de1 100644
--- a/pillar/application/modules/file_storage.py
+++ b/pillar/application/modules/file_storage.py
@@ -457,6 +457,34 @@ def refresh_links_for_project(project_uuid, chunk_size, expiry_seconds):
 
     log.info('Refreshed %i links', min(chunk_size, to_refresh.count()))
 
+
+def refresh_links_for_backend(backend_name, chunk_size, expiry_seconds):
+    from flask import current_app
+
+    # Retrieve expired links.
+    files_collection = current_app.data.driver.db['files']
+
+    now = datetime.datetime.now(tz=bson.tz_util.utc)
+    expire_before = now + datetime.timedelta(seconds=expiry_seconds)
+    log.info('Limiting to links that expire before %s', expire_before)
+
+    to_refresh = files_collection.find(
+        {'$or': [{'backend': backend_name, 'link_expires': None},
+                 {'backend': backend_name, 'link_expires': {'$lt': expire_before}},
+                 {'backend': backend_name, 'link': None}]
+        }).sort([('link_expires', pymongo.ASCENDING)]).limit(chunk_size)
+
+    if to_refresh.count() == 0:
+        log.info('No links to refresh.')
+        return
+
+    for file_doc in to_refresh:
+        log.debug('Refreshing links for file %s', file_doc['_id'])
+        _generate_all_links(file_doc, now)
+
+    log.info('Refreshed %i links', min(chunk_size, to_refresh.count()))
+
+
 def setup_app(app, url_prefix):
     app.on_pre_GET_files += on_pre_get_files
     app.on_post_POST_files += post_POST_files
diff --git a/pillar/manage.py b/pillar/manage.py
index dbc2b44c..fd02621f 100755
--- a/pillar/manage.py
+++ b/pillar/manage.py
@@ -761,6 +761,22 @@ def refresh_project_links(project, chunk_size=50, quiet=False):
     file_storage.refresh_links_for_project(project, chunk_size, 2 * 3600)
 
 
+@manager.command
+def refresh_backend_links(backend_name, chunk_size=50, quiet=False):
+    """Refreshes all file links that are using a certain storage backend."""
+
+    if quiet:
+        import logging
+        from application import log
+
+        logging.getLogger().setLevel(logging.WARNING)
+        log.setLevel(logging.WARNING)
+
+    chunk_size = int(chunk_size)  # CLI parameters are passed as strings
+    from application.modules import file_storage
+    file_storage.refresh_links_for_backend(backend_name, chunk_size, 2 * 3600)
+
+
 @manager.command
 def expire_all_project_links(project_uuid):
     """Expires all file links for a certain project without refreshing.