Improve the migration script, removing the need for --legacy

migrations/versions/939a08e1d6e5_.py

@@ -11,12 +11,73 @@ revision = '939a08e1d6e5'
 down_revision = '7e246705da6a'
 
 from alembic import op
+from flask import current_app
+from flask_sqlalchemy import SQLAlchemy
+from pathlib import Path
 import sqlalchemy as sa
 
+import os
+import time
+
+""" For a file of a given size, determine the largest allowed lifespan of that file
+
+Based on the current app's configuration:  Specifically, the MAX_CONTENT_LENGTH, as well
+as FHOST_{MIN,MAX}_EXPIRATION.
+
+This lifespan may be shortened by a user's request, but no files should be allowed to
+expire at a point after this number.
+
+Value returned is a duration in milliseconds.
+"""
+def get_max_lifespan(filesize: int) -> int:
+    min_exp = current_app.config.get("FHOST_MIN_EXPIRATION", 30 * 24 * 60 * 60 * 1000)
+    max_exp = current_app.config.get("FHOST_MAX_EXPIRATION", 365 * 24 * 60 * 60 * 1000)
+    max_size = current_app.config.get("MAX_CONTENT_LENGTH", 256 * 1024 * 1024)
+    return min_exp + int((-max_exp + min_exp) * (filesize / max_size - 1) ** 3)
+
+db = SQLAlchemy(current_app.__weakref__())
+
+# Representations of the original and updated File tables
+class File(db.Model):
+    id = db.Column(db.Integer, primary_key = True)
+    sha256 = db.Column(db.String, unique = True)
+    ext = db.Column(db.UnicodeText)
+    mime = db.Column(db.UnicodeText)
+    addr = db.Column(db.UnicodeText)
+    removed = db.Column(db.Boolean, default=False)
+    nsfw_score = db.Column(db.Float)
+UpdatedFile = sa.table('file',
+    # We only need to describe the columns that are relevent to us
+    sa.column('id', db.Integer),
+    sa.column('expiration', db.BigInteger)
+)
 
 def upgrade():
     op.add_column('file', sa.Column('expiration', sa.BigInteger()))
 
+    storage = Path(current_app.config["FHOST_STORAGE_PATH"])
+    current_time = time.time() * 1000;
+
+    # List of file hashes which have not expired yet
+    # This could get really big for some servers
+    unexpired_files = set(os.listdir(storage))
+
+    # Calculate an expiration date for all existing files
+    files = File.query\
+        .where(
+            sa.not_(File.removed)
+        ).all()
+    for file in files:
+        if file.sha256 in unexpired_files:
+            file_path = storage / file.sha256
+            stat = os.stat(file_path)
+            max_age = get_max_lifespan(stat.st_size) # How long the file is allowed to live, in ms
+            file_birth = stat.st_mtime * 1000 # When the file was created, in ms
+            op.execute(
+                sa.update(UpdatedFile)
+                    .where(UpdatedFile.c.id == file.id)
+                    .values({'expiration': int(file_birth + max_age)})
+            )
 
 def downgrade():
     op.drop_column('file', 'expiration')