From 7f0092ff101661c1cbfd64a9fd3a778078d66f94 Mon Sep 17 00:00:00 2001 From: bramval Date: Fri, 23 Jan 2026 20:48:30 +0100 Subject: [PATCH] v1 --- app/__init__.py | 8 +++ app/cli.py | 5 ++ app/models.py | 4 ++ app/routes/admin.py | 49 ++++++++++++- app/routes/company.py | 89 ++++++++++++++++++++++- app/templates/admin/company_detail.html | 32 +++++++++ app/templates/company/my_company.html | 22 ++++++ app/uploads.py | 47 +++++++++++++ scripts/storage_quota_test.py | 93 +++++++++++++++++++++++++ 9 files changed, 346 insertions(+), 3 deletions(-) create mode 100644 scripts/storage_quota_test.py diff --git a/app/__init__.py b/app/__init__.py index 4fbc87f..a7f8810 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -56,6 +56,14 @@ def create_app(): if "description" not in display_cols: db.session.execute(db.text("ALTER TABLE display ADD COLUMN description VARCHAR(200)")) db.session.commit() + + # Companies: optional per-company storage quota + company_cols = [ + r[1] for r in db.session.execute(db.text("PRAGMA table_info(company)")).fetchall() + ] + if "storage_max_bytes" not in company_cols: + db.session.execute(db.text("ALTER TABLE company ADD COLUMN storage_max_bytes BIGINT")) + db.session.commit() except Exception: db.session.rollback() diff --git a/app/cli.py b/app/cli.py index d57495b..d3ebcbf 100644 --- a/app/cli.py +++ b/app/cli.py @@ -31,6 +31,11 @@ def init_db_command(admin_email: str, admin_pass: str): if "description" not in display_cols: db.session.execute(db.text("ALTER TABLE display ADD COLUMN description VARCHAR(200)")) db.session.commit() + + company_cols = [r[1] for r in db.session.execute(db.text("PRAGMA table_info(company)")).fetchall()] + if "storage_max_bytes" not in company_cols: + db.session.execute(db.text("ALTER TABLE company ADD COLUMN storage_max_bytes BIGINT")) + db.session.commit() except Exception: # Best-effort; if it fails we continue so fresh DBs still work. db.session.rollback() diff --git a/app/models.py b/app/models.py index a94aaba..7dfaa1a 100644 --- a/app/models.py +++ b/app/models.py @@ -12,6 +12,10 @@ class Company(db.Model): name = db.Column(db.String(120), unique=True, nullable=False) created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False) + # Optional per-company storage quota for uploaded media (bytes). + # If NULL or <=0: unlimited. + storage_max_bytes = db.Column(db.BigInteger, nullable=True) + users = db.relationship("User", back_populates="company", cascade="all, delete-orphan") displays = db.relationship("Display", back_populates="company", cascade="all, delete-orphan") playlists = db.relationship("Playlist", back_populates="company", cascade="all, delete-orphan") diff --git a/app/routes/admin.py b/app/routes/admin.py index cd499a3..0c3f011 100644 --- a/app/routes/admin.py +++ b/app/routes/admin.py @@ -6,7 +6,7 @@ from flask import Blueprint, abort, current_app, flash, redirect, render_templat from flask_login import current_user, login_required, login_user from ..extensions import db -from ..uploads import abs_upload_path, ensure_company_upload_dir, is_valid_upload_relpath +from ..uploads import abs_upload_path, ensure_company_upload_dir, get_company_upload_bytes, is_valid_upload_relpath from ..models import Company, Display, DisplaySession, Playlist, PlaylistItem, User bp = Blueprint("admin", __name__, url_prefix="/admin") @@ -75,7 +75,52 @@ def company_detail(company_id: int): company = db.session.get(Company, company_id) if not company: abort(404) - return render_template("admin/company_detail.html", company=company) + + upload_root = current_app.config["UPLOAD_FOLDER"] + used_bytes = get_company_upload_bytes(upload_root, company.id) + + return render_template( + "admin/company_detail.html", + company=company, + storage={ + "used_bytes": used_bytes, + }, + ) + + +@bp.post("/companies//storage") +@login_required +def update_company_storage(company_id: int): + admin_required() + + company = db.session.get(Company, company_id) + if not company: + abort(404) + + raw = (request.form.get("storage_max_mb") or "").strip() + if raw == "": + # Treat empty as unlimited + company.storage_max_bytes = None + db.session.commit() + flash("Storage limit cleared (unlimited).", "success") + return redirect(url_for("admin.company_detail", company_id=company_id)) + + try: + mb = float(raw) + except ValueError: + flash("Invalid storage limit. Please enter a number (MB).", "danger") + return redirect(url_for("admin.company_detail", company_id=company_id)) + + if mb <= 0: + company.storage_max_bytes = None + db.session.commit() + flash("Storage limit cleared (unlimited).", "success") + return redirect(url_for("admin.company_detail", company_id=company_id)) + + company.storage_max_bytes = int(mb * 1024 * 1024) + db.session.commit() + flash("Storage limit updated.", "success") + return redirect(url_for("admin.company_detail", company_id=company_id)) @bp.post("/companies//users") diff --git a/app/routes/company.py b/app/routes/company.py index 284ecc5..fd63d89 100644 --- a/app/routes/company.py +++ b/app/routes/company.py @@ -11,7 +11,13 @@ from werkzeug.utils import secure_filename from PIL import Image, ImageOps from ..extensions import db -from ..uploads import abs_upload_path, ensure_company_upload_dir, get_company_upload_bytes, is_valid_upload_relpath +from ..uploads import ( + abs_upload_path, + compute_storage_usage, + ensure_company_upload_dir, + get_company_upload_bytes, + is_valid_upload_relpath, +) from ..models import Company, Display, DisplaySession, Playlist, PlaylistItem, User from ..email_utils import send_email from ..auth_tokens import make_password_reset_token @@ -194,6 +200,12 @@ def _format_bytes(num: int) -> str: return f"{size:.1f} {units[idx]}" +def _storage_limit_error_message(*, storage_max_human: str | None) -> str: + if storage_max_human: + return f"Storage limit reached. Maximum allowed storage is {storage_max_human}. Please delete items to free space." + return "Storage limit reached. Please delete items to free space." + + @bp.get("/my-company") @login_required def my_company(): @@ -226,6 +238,9 @@ def my_company(): upload_root = current_app.config["UPLOAD_FOLDER"] used_bytes = get_company_upload_bytes(upload_root, company.id) + usage = compute_storage_usage(used_bytes=used_bytes, max_bytes=company.storage_max_bytes) + max_human = _format_bytes(usage["max_bytes"]) if usage.get("max_bytes") else None + users = User.query.filter_by(company_id=company.id, is_admin=False).order_by(User.email.asc()).all() return render_template( @@ -240,6 +255,9 @@ def my_company(): "active_sessions": active_sessions, "storage_bytes": used_bytes, "storage_human": _format_bytes(used_bytes), + "storage_max_bytes": usage.get("max_bytes"), + "storage_max_human": max_human, + "storage_used_percent": usage.get("used_percent"), }, ) @@ -489,7 +507,29 @@ def add_playlist_item(playlist_id: int): position=pos, ) + # Enforce storage quota for uploads (image/video). + # Webpage/YouTube do not consume local storage. + # Note: querying the DB triggers an autoflush by default. Because `item` is not yet in the + # session, SQLAlchemy may emit warnings about relationship operations. We explicitly avoid + # autoflush while checking quota. + with db.session.no_autoflush: + company = db.session.get(Company, current_user.company_id) + if not company: + abort(404) + + upload_root = current_app.config["UPLOAD_FOLDER"] + used_bytes = get_company_upload_bytes(upload_root, company.id) + usage = compute_storage_usage(used_bytes=used_bytes, max_bytes=company.storage_max_bytes) + storage_max_human = _format_bytes(usage["max_bytes"]) if usage.get("max_bytes") else None + if item_type in ("image", "video"): + if usage.get("is_exceeded"): + msg = _storage_limit_error_message(storage_max_human=storage_max_human) + if wants_json: + return _json_error(msg, 403) + flash(msg, "danger") + return redirect(url_for("company.playlist_detail", playlist_id=playlist_id)) + f = request.files.get("file") if not f or not f.filename: if wants_json: @@ -516,6 +556,30 @@ def add_playlist_item(playlist_id: int): current_user.company_id, crop_mode=crop_mode, ) + + # Post-save quota check for images as well. + # (We can't reliably estimate image size before compression.) + if company.storage_max_bytes is not None and int(company.storage_max_bytes or 0) > 0: + try: + used_after = get_company_upload_bytes(upload_root, company.id) + except Exception: + used_after = None + if used_after is not None: + usage_after = compute_storage_usage( + used_bytes=used_after, + max_bytes=company.storage_max_bytes, + ) + if usage_after.get("is_exceeded"): + # Remove the saved file and reject. + try: + _try_delete_upload(item.file_path, upload_root) + except Exception: + pass + msg = _storage_limit_error_message(storage_max_human=storage_max_human) + if wants_json: + return _json_error(msg, 403) + flash(msg, "danger") + return redirect(url_for("company.playlist_detail", playlist_id=playlist_id)) except Exception: if wants_json: return _json_error("Failed to process image upload", 500) @@ -558,6 +622,29 @@ def add_playlist_item(playlist_id: int): save_path = os.path.join(company_dir, unique) f.save(save_path) + # Post-save quota check: clients may not report size reliably. + # If quota is exceeded after saving, delete file and reject. + if company.storage_max_bytes is not None and int(company.storage_max_bytes or 0) > 0: + try: + used_after = get_company_upload_bytes(upload_root, company.id) + except Exception: + used_after = None + if used_after is not None: + usage_after = compute_storage_usage( + used_bytes=used_after, + max_bytes=company.storage_max_bytes, + ) + if usage_after.get("is_exceeded"): + try: + os.remove(save_path) + except OSError: + pass + msg = _storage_limit_error_message(storage_max_human=storage_max_human) + if wants_json: + return _json_error(msg, 403) + flash(msg, "danger") + return redirect(url_for("company.playlist_detail", playlist_id=playlist_id)) + # Safety check: validate using the actual saved file size. # (Some clients/framework layers don't reliably report per-part size.) try: diff --git a/app/templates/admin/company_detail.html b/app/templates/admin/company_detail.html index 80cb8d0..4433fe8 100644 --- a/app/templates/admin/company_detail.html +++ b/app/templates/admin/company_detail.html @@ -59,6 +59,38 @@ +
+
+
+

Storage limit

+
+
+
+ Used: {{ storage.used_bytes }} bytes +
+ +
+
+ + +
Set to 0 or empty to disable the limit.
+
+
+ +
+
+
+
+
+
diff --git a/app/templates/company/my_company.html b/app/templates/company/my_company.html index 1283043..1b94858 100644 --- a/app/templates/company/my_company.html +++ b/app/templates/company/my_company.html @@ -44,6 +44,28 @@
Storage used
{{ stats['storage_human'] }}
({{ stats['storage_bytes'] }} bytes)
+ {% if stats.get('storage_max_human') %} +
+ Limit: {{ stats['storage_max_human'] }} + {% if stats.get('storage_used_percent') is not none %} + — Used: {{ stats['storage_used_percent'] }}% + {% endif %} +
+ {% if stats.get('storage_used_percent') is not none %} +
+
+
+ {% endif %} + {% else %} +
Limit: Unlimited
+ {% endif %}
diff --git a/app/uploads.py b/app/uploads.py index 52be7f7..90ffc90 100644 --- a/app/uploads.py +++ b/app/uploads.py @@ -56,6 +56,53 @@ def get_company_upload_bytes(upload_root: str, company_id: int | None) -> int: return int(total) +def compute_storage_usage(*, used_bytes: int, max_bytes: int | None): + """Compute storage usage info. + + Args: + used_bytes: current usage + max_bytes: quota; if None or <=0: unlimited + + Returns dict: + { + "max_bytes": int|None, + "used_bytes": int, + "is_limited": bool, + "is_exceeded": bool, + "used_ratio": float|None, # 0..1 when limited + "used_percent": int|None, # rounded percent when limited + "remaining_bytes": int|None, + } + """ + + used = max(0, int(used_bytes or 0)) + mx = None if max_bytes is None else int(max_bytes) + if mx is None or mx <= 0: + return { + "max_bytes": None, + "used_bytes": used, + "is_limited": False, + "is_exceeded": False, + "used_ratio": None, + "used_percent": None, + "remaining_bytes": None, + } + + ratio = used / mx if mx > 0 else 1.0 + percent = int(round(ratio * 100.0)) + return { + "max_bytes": mx, + "used_bytes": used, + "is_limited": True, + "is_exceeded": used >= mx, + # Keep percent un-clamped so the UI can show e.g. 132% when exceeded. + # Clamp only the ratio (for progress bars, etc.). + "used_ratio": max(0.0, min(1.0, ratio)), + "used_percent": max(0, percent), + "remaining_bytes": max(0, mx - used), + } + + def is_valid_upload_relpath(file_path: str | None) -> bool: """True if file_path looks like a path we manage under /static. diff --git a/scripts/storage_quota_test.py b/scripts/storage_quota_test.py new file mode 100644 index 0000000..8b1105f --- /dev/null +++ b/scripts/storage_quota_test.py @@ -0,0 +1,93 @@ +import io +import os +import sys +import tempfile + + +# Ensure repo root is on sys.path when running as a script. +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +from app import create_app +from app.extensions import db +from app.models import Company, Playlist, User + + +def main(): + # Use a temporary SQLite DB so this doesn't touch your real instance DB. + fd, path = tempfile.mkstemp(prefix="rssfeed-test-", suffix=".sqlite") + os.close(fd) + try: + os.environ["DATABASE_URL"] = f"sqlite:///{path}" + os.environ["SECRET_KEY"] = "test-secret" + + app = create_app() + app.config["TESTING"] = True + + with app.app_context(): + db.create_all() + + c = Company(name="TestCo") + # 1 byte quota so any non-empty upload should be rejected. + c.storage_max_bytes = 1 + db.session.add(c) + db.session.commit() + + u = User(username="test@example.com", email="test@example.com", is_admin=False, company_id=c.id) + u.set_password("passw0rd123") + db.session.add(u) + db.session.commit() + + p = Playlist(company_id=c.id, name="Playlist") + db.session.add(p) + db.session.commit() + + client = app.test_client() + + # Login + res = client.post( + "/auth/login", + data={"email": "test@example.com", "password": "passw0rd123"}, + follow_redirects=False, + ) + if res.status_code not in (302, 303): + raise SystemExit(f"Login failed: {res.status_code}") + + # Try to upload a small image payload. + data = { + "item_type": "image", + "title": "Small", + "duration_seconds": "10", + "crop_mode": "none", + "response": "json", + # Not a real image; should fail processing OR quota depending on Pillow. + # Use a tiny valid PNG header so Pillow can parse it. + "file": ( + io.BytesIO( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00\x01\x01\x01\x00\x18\xdd\x8d\x18\x00\x00\x00\x00IEND\xaeB`\x82" + ), + "tiny.png", + ), + } + url = f"/company/playlists/{p.id}/items" + res = client.post(url, data=data, content_type="multipart/form-data") + + if res.status_code != 403: + raise SystemExit(f"Expected 403 for quota exceeded, got: {res.status_code} {res.data!r}") + js = res.get_json(silent=True) or {} + if js.get("ok") is not False: + raise SystemExit(f"Unexpected response: {js}") + if "Storage limit" not in (js.get("error") or ""): + raise SystemExit(f"Unexpected error message: {js}") + + print("OK: storage quota prevents uploads when exceeded") + finally: + try: + os.remove(path) + except OSError: + pass + + +if __name__ == "__main__": + main()