Begin tests for user quota

[github/Chocobozzz/PeerTube.git] / server / models / user / user.ts

diff --git a/server/models/user/user.ts b/server/models/user/user.ts
index e1b9339883c78c3c7b6b8e400153ea2839656a60..9bf13ad243e5098cef1fd806332344bd3d64567e 100644 (file)
--- 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) {
@@ -198,7 +221,7 @@ loadById = function (id: number) {
 loadByUsername = function (username: string) {
   const query = {
     where: {
-      username: username
+      username
     }
   }
 
@@ -212,5 +235,42 @@ loadByUsernameOrEmail = function (username: string, email: string) {
     }
   }
 
-  return User.findOne(query)
+  // 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('User.Author.Video.VideoFile.size'), 'totalVideoBytes')
+    ],
+    where: {
+      id: user.id
+    },
+    include: [
+      {
+        model: User['sequelize'].models.Author,
+        required: true,
+        include: [
+          {
+            model: User['sequelize'].models.Video,
+            required: true,
+            include: [
+              {
+                model: User['sequelize'].models.VideoFile,
+                required: true
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+
+  // FIXME: cast to any because of bad typing...
+  return User.findAll(query).then((res: any) => {
+    return res.totalVideoBytes
+  })
 }