server/info: implement simple global info

For now it returns total post count and disk usage. To reduce stress on
the file system, the disk usage is cached for an hour.

5 files changed, 63 insertions, 0 deletions

diff --git a/server/szurubooru/api/__init__.py b/server/szurubooru/api/__init__.py
index c694418..fd68915 100644
--- a/server/szurubooru/api/__init__.py
+++ b/server/szurubooru/api/__init__.py
@@ -5,4 +5,5 @@ from szurubooru.api.user_api import UserListApi, UserDetailApi
 from szurubooru.api.tag_api import TagListApi, TagDetailApi, TagMergeApi, TagSiblingsApi
 from szurubooru.api.tag_category_api import TagCategoryListApi, TagCategoryDetailApi
 from szurubooru.api.snapshot_api import SnapshotListApi
+from szurubooru.api.info_api import InfoApi
 from szurubooru.api.context import Context, Request
diff --git a/server/szurubooru/api/info_api.py b/server/szurubooru/api/info_api.py
new file mode 100644
index 0000000..41c1423
--- /dev/null
+++ b/server/szurubooru/api/info_api.py
@@ -0,0 +1,30 @@
+import datetime
+import os
+from szurubooru import config
+from szurubooru.api.base_api import BaseApi
+from szurubooru.func import posts
+
+class InfoApi(BaseApi):
+    def __init__(self):
+        self._cache_time = None
+        self._cache_result = None
+
+    def get(self, ctx):
+        return {
+            'postCount': posts.get_post_count(),
+            'diskUsage': self._get_disk_usage()
+        }
+
+    def _get_disk_usage(self):
+        threshold = datetime.timedelta(hours=1)
+        now = datetime.datetime.now()
+        if self._cache_time and self._cache_time > now - threshold:
+            return self._cache_result
+        total_size = 0
+        for dir_path, _, file_names in os.walk(config.config['data_dir']):
+            for file_name in file_names:
+                file_path = os.path.join(dir_path, file_name)
+                total_size += os.path.getsize(file_path)
+        self._cache_time = now
+        self._cache_result = total_size
+        return total_size
diff --git a/server/szurubooru/app.py b/server/szurubooru/app.py
index 8a5f902..b6c285c 100644
--- a/server/szurubooru/app.py
+++ b/server/szurubooru/app.py
@@ -57,6 +57,7 @@ def create_app():
     tag_siblings_api = api.TagSiblingsApi()
     password_reset_api = api.PasswordResetApi()
     snapshot_list_api = api.SnapshotListApi()
+    info_api = api.InfoApi()
 
     app.add_error_handler(errors.AuthError, _on_auth_error)
     app.add_error_handler(errors.IntegrityError, _on_integrity_error)
@@ -75,5 +76,6 @@ def create_app():
     app.add_route('/tag-siblings/{tag_name}', tag_siblings_api)
     app.add_route('/password-reset/{user_name}', password_reset_api)
     app.add_route('/snapshots/', snapshot_list_api)
+    app.add_route('/info/', info_api)
 
     return app
diff --git a/server/szurubooru/func/posts.py b/server/szurubooru/func/posts.py
new file mode 100644
index 0000000..027cf5e
--- /dev/null
+++ b/server/szurubooru/func/posts.py
@@ -0,0 +1,5 @@
+import sqlalchemy
+from szurubooru import db
+
+def get_post_count():
+    return db.session.query(sqlalchemy.func.count(db.Post.post_id)).one()[0]
diff --git a/server/szurubooru/tests/api/test_info.py b/server/szurubooru/tests/api/test_info.py
new file mode 100644
index 0000000..e9d2c1a
--- /dev/null
+++ b/server/szurubooru/tests/api/test_info.py
@@ -0,0 +1,25 @@
+from szurubooru import api, db
+
+def test_info_api(
+        tmpdir, config_injector, context_factory, post_factory, fake_datetime):
+    directory = tmpdir.mkdir('data')
+    directory.join('test.txt').write('abc')
+    config_injector({'data_dir': str(directory)})
+    db.session.add_all([post_factory(), post_factory()])
+    info_api = api.InfoApi()
+    with fake_datetime('13:00'):
+        assert info_api.get(context_factory()) == {
+            'postCount': 2,
+            'diskUsage': 3,
+        }
+    directory.join('test2.txt').write('abc')
+    with fake_datetime('13:59'):
+        assert info_api.get(context_factory()) == {
+            'postCount': 2,
+            'diskUsage': 3, # still 3 - it's cached
+        }
+    with fake_datetime('14:01'):
+        assert info_api.get(context_factory()) == {
+            'postCount': 2,
+            'diskUsage': 6, # cache expired
+        }