fix(docworker): Fix storage limit in document worker

ds-wizard · Dec 19, 2024 · 1f63e27 · 1f63e27
1 parent 35f2d37
commit 1f63e27
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/packages/dsw-database/dsw/database/database.py b/packages/dsw-database/dsw/database/database.py
@@ -65,7 +65,9 @@ class Database:
                      'FROM document WHERE tenant_uuid = %(tenant_uuid)s) ' \
                      '+ (SELECT COALESCE(SUM(file_size)::bigint, 0) ' \
                      'FROM document_template_asset WHERE tenant_uuid = %(tenant_uuid)s) ' \
-                     'as result;'
+                     '+ (SELECT COALESCE(SUM(file_size)::bigint, 0) ' \
+                     'FROM questionnaire_file WHERE tenant_uuid = %(tenant_uuid)s) ' \
+                     'AS result;'
 
     def __init__(self, cfg: DatabaseConfig, connect: bool = True,
                  with_queue: bool = True):

diff --git a/packages/dsw-document-worker/dsw/document_worker/limits.py b/packages/dsw-document-worker/dsw/document_worker/limits.py
@@ -21,6 +21,8 @@ def check_doc_size(job_id: str, doc_size: int):
     @staticmethod
     def check_size_usage(job_id: str, doc_size: int,
                          used_size: int, limit_size: Optional[int]):
+        limit_size = abs(limit_size) if limit_size is not None else None
+
         if limit_size is None or doc_size + used_size < limit_size:
             return
         remains = limit_size - used_size

diff --git a/packages/dsw-document-worker/dsw/document_worker/worker.py b/packages/dsw-document-worker/dsw/document_worker/worker.py
@@ -56,8 +56,8 @@ def __init__(self, command: PersistentCommand, document_uuid: str):
         self.doc_context = command.body  # type: dict
         self.doc = None  # type: Optional[DBDocument]
         self.final_file = None  # type: Optional[DocumentFile]
-        self.tenant_config = None  # type: Optional[DBTenantConfig]
-        self.tenant_limits = None  # type: Optional[DBTenantLimits]
+        self.tenant_config = self.ctx.app.db.get_tenant_config(self.tenant_uuid)  # type: Optional[DBTenantConfig]
+        self.tenant_limits = self.ctx.app.db.fetch_tenant_limits(self.tenant_uuid)  # type: Optional[DBTenantLimits]
 
     @property
     def safe_doc(self) -> DBDocument: