Begin user quota

2 files changed, 62 insertions, 2 deletions

diff --git a/server/models/user/user-interface.ts b/server/models/user/user-interface.ts
index 0b97a8f6d..8974a9a97 100644
--- a/server/models/user/user-interface.ts
+++ b/server/models/user/user-interface.ts
@@ -11,6 +11,7 @@ export namespace UserMethods {
 
   export type ToFormattedJSON = (this: UserInstance) => FormattedUser
   export type IsAdmin = (this: UserInstance) => boolean
+  export type IsAbleToUploadVideo = (this: UserInstance, videoFile: Express.Multer.File) => Promise<boolean>
 
   export type CountTotal = () => Promise<number>
 
@@ -31,6 +32,7 @@ export interface UserClass {
   isPasswordMatch: UserMethods.IsPasswordMatch,
   toFormattedJSON: UserMethods.ToFormattedJSON,
   isAdmin: UserMethods.IsAdmin,
+  isAbleToUploadVideo: UserMethods.IsAbleToUploadVideo,
 
   countTotal: UserMethods.CountTotal,
   getByUsername: UserMethods.GetByUsername,
@@ -42,11 +44,13 @@ export interface UserClass {
 }
 
 export interface UserAttributes {
+  id?: number
   password: string
   username: string
   email: string
   displayNSFW?: boolean
   role: UserRole
+  videoQuota: number
 }
 
 export interface UserInstance extends UserClass, UserAttributes, Sequelize.Instance<UserAttributes> {
diff --git a/server/models/user/user.ts b/server/models/user/user.ts
index d481fa13c..12a7547f5 100644
--- a/server/models/user/user.ts
+++ b/server/models/user/user.ts
@@ -1,5 +1,6 @@
 import { values } from 'lodash'
 import * as Sequelize from 'sequelize'
+import * as Promise from 'bluebird'
 
 import { getSort } from '../utils'
 import { USER_ROLES } from '../../initializers'
@@ -8,7 +9,8 @@ import {
   comparePassword,
   isUserPasswordValid,
   isUserUsernameValid,
-  isUserDisplayNSFWValid
+  isUserDisplayNSFWValid,
+  isUserVideoQuotaValid
 } from '../../helpers'
 
 import { addMethodsToModel } from '../utils'
@@ -30,6 +32,7 @@ let listForApi: UserMethods.ListForApi
 let loadById: UserMethods.LoadById
 let loadByUsername: UserMethods.LoadByUsername
 let loadByUsernameOrEmail: UserMethods.LoadByUsernameOrEmail
+let isAbleToUploadVideo: UserMethods.IsAbleToUploadVideo
 
 export default function (sequelize: Sequelize.Sequelize, DataTypes: Sequelize.DataTypes) {
   User = sequelize.define<UserInstance, UserAttributes>('User',
@@ -75,6 +78,16 @@ export default function (sequelize: Sequelize.Sequelize, DataTypes: Sequelize.Da
       role: {
         type: DataTypes.ENUM(values(USER_ROLES)),
         allowNull: false
+      },
+      videoQuota: {
+        type: DataTypes.BIGINT,
+        allowNull: false,
+        validate: {
+          videoQuotaValid: value => {
+            const res = isUserVideoQuotaValid(value)
+            if (res === false) throw new Error('Video quota is not valid.')
+          }
+        }
       }
     },
     {
@@ -109,7 +122,8 @@ export default function (sequelize: Sequelize.Sequelize, DataTypes: Sequelize.Da
   const instanceMethods = [
     isPasswordMatch,
     toFormattedJSON,
-    isAdmin
+    isAdmin,
+    isAbleToUploadVideo
   ]
   addMethodsToModel(User, classMethods, instanceMethods)
 
@@ -136,6 +150,7 @@ toFormattedJSON = function (this: UserInstance) {
     email: this.email,
     displayNSFW: this.displayNSFW,
     role: this.role,
+    videoQuota: this.videoQuota,
     createdAt: this.createdAt
   }
 }
@@ -144,6 +159,14 @@ isAdmin = function (this: UserInstance) {
   return this.role === USER_ROLES.ADMIN
 }
 
+isAbleToUploadVideo = function (this: UserInstance, videoFile: Express.Multer.File) {
+  if (this.videoQuota === -1) return Promise.resolve(true)
+
+  return getOriginalVideoFileTotalFromUser(this).then(totalBytes => {
+    return (videoFile.size + totalBytes) < this.videoQuota
+  })
+}
+
 // ------------------------------ STATICS ------------------------------
 
 function associate (models) {
@@ -215,3 +238,36 @@ loadByUsernameOrEmail = function (username: string, email: string) {
   // FIXME: https://github.com/DefinitelyTyped/DefinitelyTyped/issues/18387
   return (User as any).findOne(query)
 }
+
+// ---------------------------------------------------------------------------
+
+function getOriginalVideoFileTotalFromUser (user: UserInstance) {
+  const query = {
+    attributes: [
+      Sequelize.fn('COUNT', Sequelize.col('VideoFile.size'), 'totalVideoBytes')
+    ],
+    where: {
+      id: user.id
+    },
+    include: [
+      {
+        model: User['sequelize'].models.Author,
+        include: [
+          {
+            model: User['sequelize'].models.Video,
+            include: [
+              {
+                model: User['sequelize'].models.VideoFile
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+
+  // FIXME: cast to any because of bad typing...
+  return User.findAll(query).then((res: any) => {
+    return res.totalVideoBytes
+  })
+}