From b72ecc321c315bafe40cc7406e87e088564ab8a9 Mon Sep 17 00:00:00 2001
From: Jon Bergli Heier <snakebite@jvnv.net>
Date: Wed, 24 Jul 2019 09:02:43 +0200
Subject: Add file storage modules

Allows for storing files other places than the local file system.
Currently the local filesystem and S3 are supported.
---
 fbin/db.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

(limited to 'fbin/db.py')

diff --git a/fbin/db.py b/fbin/db.py
index 2bf153b..be79c76 100644
--- a/fbin/db.py
+++ b/fbin/db.py
@@ -4,7 +4,7 @@ import mimetypes
 import os
 
 from flask import current_app
-from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Index, ForeignKey, Boolean, JSON
+from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Index, ForeignKey, Boolean, JSON, BigInteger
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relation, backref
 from sqlalchemy.orm.exc import NoResultFound
@@ -48,6 +48,7 @@ class File(Base):
     id = Column(Integer, primary_key = True)
     hash = Column(String, unique = True, index = True)
     filename = Column(String)
+    size = Column(BigInteger)
     date = Column(DateTime)
     user_id = Column(Integer, ForeignKey('users.id'), nullable = True)
     ip = Column(String)
@@ -55,9 +56,10 @@ class File(Base):
     scanned = Column(Boolean, nullable=False, default=False)
     blocked_reason = Column(JSON)
 
-    def __init__(self, hash, filename, date, user_id = None, ip = None):
+    def __init__(self, hash, filename, size, date, user_id = None, ip = None):
         self.hash = hash
         self.filename = filename
+        self.size = size
         self.date = date
         self.user_id = user_id
         self.ip = ip
@@ -82,6 +84,8 @@ class File(Base):
 
     def get_size(self):
         try:
+            if self.size:
+                return self.size
             return os.path.getsize(self.get_path())
         except OSError:
             return None
-- 
cgit v1.2.3