From 06054ceeafdc79284093d8379871971d6f84f54e Mon Sep 17 00:00:00 2001
From: Goffi <goffi@goffi.org>
Date: Wed, 5 May 2021 15:37:21 +0200
Subject: [PATCH] core (memory, sqlite): new `fileGetUsedSpace` method

---
 sat/memory/memory.py | 22 ++++++++++++++++++++--
 sat/memory/sqlite.py |  9 +++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/sat/memory/memory.py b/sat/memory/memory.py
index 55f97be9..9783ab9d 100644
--- a/sat/memory/memory.py
+++ b/sat/memory/memory.py
@@ -1456,8 +1456,8 @@ class Memory(object):
             client,
             owner: Optional[jid.JID],
             peer_jid: Optional[jid.JID],
-            file_id: Optional[str],
-            parent: Optional[str]
+            file_id: Optional[str] = None,
+            parent: Optional[str] = None
     ) -> jid.JID:
         """Get owner to use for a file operation
 
@@ -1709,6 +1709,24 @@ class Memory(object):
             extra=extra,
         )
 
+    async def fileGetUsedSpace(
+        self,
+        client,
+        peer_jid: jid.JID,
+        owner: Optional[jid.JID] = None
+    ) -> int:
+        """Get space taken by all files owned by an entity
+
+        @param peer_jid: entity requesting the size
+        @param owner: entity owning the file to check. If None, will be determined by
+            getFilesOwner
+        @return: size of total space used by files of this owner
+        """
+        owner = self.getFilesOwner(client, owner, peer_jid)
+        if peer_jid.userhostJID() != owner and client.profile not in self.admins:
+            raise exceptions.PermissionError("You are not allowed to check this size")
+        return await self.storage.fileGetUsedSpace(client, owner)
+
     def fileUpdate(self, file_id, column, update_cb):
         """Update a file column taking care of race condition
 
diff --git a/sat/memory/sqlite.py b/sat/memory/sqlite.py
index 2d04f480..a9c61fbd 100644
--- a/sat/memory/sqlite.py
+++ b/sat/memory/sqlite.py
@@ -986,6 +986,15 @@ class SqliteStorage(object):
         d.addErrback(lambda failure: log.error(_("Can't save file metadata for [{profile}]: {reason}".format(profile=client.profile, reason=failure))))
         return d
 
+    async def fileGetUsedSpace(self, client, owner):
+        """Get space used by owner of file"""
+        query = "SELECT SUM(size) FROM files WHERE owner=? AND type='file' AND profile_id=?"
+        ret = await self.dbpool.runQuery(
+            query,
+            (owner.userhost(), self.profiles[client.profile])
+        )
+        return ret[0][0]
+
     def _fileUpdate(self, cursor, file_id, column, update_cb):
         query = 'SELECT {column} FROM files where id=?'.format(column=column)
         for i in range(5):