feat(document-worker): Enforce storage limit from the command body

packages/dsw-database/CHANGELOG.md

@@ -21,15 +21,14 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   (instead of `organization_id` + `template_id`, `coordinates` is `<id>:<version>`), the
   `{id, minVersion, maxVersion}` package patterns and `workspace_uuid`, and its
   `metamodel_version` is read as `"major.minor"`; `project` gains `workspace_uuid`
-- `fetch_tenant_limits` returns `None` (no limit) when there is no `tenant_limit_bundle` table,
-  as in FAIR Wizard 5.0, which computes limits from tenant plans
 - `fetch_project_simple` returns `None` for a missing project
 - `DBLocale` has the Wizard 5.0 shape of the `locale` table (`code`, `name`, `pot_version`,
   `enabled`, `default_locale`, `official`, timestamps); `organization_id`, `locale_id`,
   `version` and the `id` property are gone
 
 ### Removed
 
+- `fetch_tenant_limits` and `DBTenantLimits`: the document worker takes the storage limit from the command body instead of `tenant_limit_bundle`
 - `get_default_locale`: mails no longer fall back to the default UI locale, but to the default
   mail template locale
 

packages/dsw-database/dsw/database/database.py

@@ -59,8 +59,6 @@ class Database:
                                   'WHERE d.project_uuid = %s AND d.tenant_uuid = %s;')
     SELECT_PROJECT_SIMPLE = ('SELECT p.* FROM project p '
                              'WHERE p.uuid = %s AND p.tenant_uuid = %s;')
-    SELECT_TENANT_LIMIT = ('SELECT uuid, storage FROM tenant_limit_bundle '
-                           'WHERE uuid = %(tenant_uuid)s LIMIT 1;')
     UPDATE_DOCUMENT_STATE = 'UPDATE document SET state = %s, worker_log = %s WHERE uuid = %s;'
     UPDATE_DOCUMENT_RETRIEVED = 'UPDATE document SET retrieved_at = %s, state = %s WHERE uuid = %s;'
     UPDATE_DOCUMENT_FINISHED = ('UPDATE document SET finished_at = %s, state = %s, '
@@ -176,29 +174,6 @@ def fetch_document(self, document_uuid: str, tenant_uuid: str) -> model.DBDocume
                 return None
             return model.DBDocument.from_dict_row(result[0])
 
-    @tenacity.retry(
-        reraise=True,
-        wait=tenacity.wait_exponential(multiplier=RETRY_QUERY_MULTIPLIER),
-        stop=tenacity.stop_after_attempt(RETRY_QUERY_TRIES),
-        before=tenacity.before_log(LOG, logging.DEBUG),
-        after=tenacity.after_log(LOG, logging.DEBUG),
-    )
-    @_in_savepoint
-    def fetch_tenant_limits(self, tenant_uuid: str) -> model.DBTenantLimits | None:
-        # FAIR Wizard 5.0 computes limits from tenant plans and has no such table;
-        # the backend enforces them, so the worker has no storage limit of its own then
-        if not self._check_table_exists(table_name='tenant_limit_bundle'):
-            return None
-        with self.conn_query.new_cursor(use_dict=True) as cursor:
-            cursor.execute(
-                query=self.SELECT_TENANT_LIMIT,
-                params={'tenant_uuid': tenant_uuid},
-            )
-            result = cursor.fetchall()
-            if len(result) != 1:
-                return None
-            return model.DBTenantLimits.from_dict_row(result[0])
-
     @tenacity.retry(
         reraise=True,
         wait=tenacity.wait_exponential(multiplier=RETRY_QUERY_MULTIPLIER),

packages/dsw-database/dsw/database/model.py

@@ -301,19 +301,6 @@ def from_dict_row(data: dict):
         )
 
 
-@dataclasses.dataclass
-class DBTenantLimits:
-    tenant_uuid: str
-    storage: int | None
-
-    @staticmethod
-    def from_dict_row(data: dict):
-        return DBTenantLimits(
-            tenant_uuid=str(data['uuid']),
-            storage=data['storage'],
-        )
-
-
 @dataclasses.dataclass
 class DBSubmission:
     TABLE_NAME = 'submission'

packages/dsw-database/tests/test_savepoint.py

@@ -171,16 +171,6 @@ def test_without_savepoint_the_retry_only_sees_the_aborted_transaction(database)
         database.flaky_without_savepoint()
 
 
-@integration
-def test_missing_tenant_limit_table_means_no_limit(database):
-    with database.conn_query.new_cursor() as cursor:
-        cursor.execute("SELECT to_regclass('tenant_limit_bundle') IS NULL;")
-        if not cursor.fetchone()[0]:
-            pytest.skip('the test database has a tenant_limit_bundle table')
-    assert database.fetch_tenant_limits('00000000-0000-0000-0000-000000000000') is None
-    assert database.conn_query.connection.info.transaction_status != INERROR
-
-
 @integration
 def test_committing_query_leaves_a_usable_connection(database):
     with database.conn_query.savepoint(), database.conn_query.new_cursor() as cursor:

packages/dsw-document-worker/CHANGELOG.md

@@ -23,7 +23,8 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - `extras.questionnaire`, the deprecated alias of `extras.project`, is removed; for a KM editor preview (no project) `extras.project` is `null` and `extras.submissions` empty
 - Configuration `documentContext.defaultAppTitleShort` and `documentContext.defaultIllustrationsColor` (`DOCUMENT_CONTEXT_DEFAULT_APP_TITLE_SHORT`, `DOCUMENT_CONTEXT_DEFAULT_ILLUSTRATIONS_COLOR`) are removed with the context keys they filled; `config.example.yml` documents the `documentContext` section under its real name
 - Template references are `<id>:<version>` (Wizard 5.0), so `templates[].ids` prefixes in the configuration change from e.g. `dsw:` to `dsw.`
-- Without a `tenant_limit_bundle` table (FAIR Wizard 5.0) the worker applies no storage limit of its own; the backend enforces the plan
+- The tenant storage limit comes only from `limits.storage` of the `generateDocument` command body, as resolved by the backend; the key is removed before rendering, so it is not part of the document context, and without it (or with `null`) there is no storage limit. The worker no longer reads `tenant_limit_bundle`
+- The maximum size of a generated document is no longer experimental and is limited by default: `documents.maxSize` (`DOCUMENTS_MAX_SIZE`) replaces `experimental.maxDocumentSize` (`EXPERIMENTAL_MAX_DOCUMENT_SIZE`) and defaults to 5 MB (`5000000` bytes); set it to `null` (or a blank variable) for no limit
 - `generatePotFile` reads the single template `id` and `version` from the command body (Wizard 5.0) and stores `document-templates/<uuid>/<id>_<version>.pot`, the name the backend reads; a body queued before 5.0 (`organizationId` + `templateId`) is read as the joined id the 5.0 migration gave the template
 - The `jinja2.ext.i18n` extension is always enabled for Jinja-powered steps
 - `extras.project` (and the deprecated `extras.questionnaire`) provide `knowledge_model_package_uuid` instead of `knowledge_package_uuid`

packages/dsw-document-worker/config.example.yml

@@ -54,6 +54,7 @@ templates:
 #documents:
 #  naming:
 #    strategy: sanitize # uuid|slugify|sanitize
+#  maxSize: 5000000 # bytes per generated document, null for no limit
 
 #externals:
 #  pandoc:

packages/dsw-document-worker/dsw/document_worker/config.py

@@ -44,6 +44,12 @@ class _DocumentsKeys(ConfigKeysContainer):
         default='sanitize',
         cast=cast_str,
     )
+    max_size = ConfigKey(
+        yaml_path=['documents', 'maxSize'],
+        var_names=['DOCUMENTS_MAX_SIZE'],
+        default=5_000_000,
+        cast=cast_optional_int,
+    )
 
 
 class _ExperimentalKeys(ConfigKeysContainer):
@@ -53,12 +59,6 @@ class _ExperimentalKeys(ConfigKeysContainer):
         default=None,
         cast=cast_optional_int,
     )
-    max_doc_size = ConfigKey(
-        yaml_path=['experimental', 'maxDocumentSize'],
-        var_names=['EXPERIMENTAL_MAX_DOCUMENT_SIZE'],
-        default=None,
-        cast=cast_optional_int,
-    )
 
 
 class _SecurityKeys(ConfigKeysContainer):
@@ -171,15 +171,16 @@ class DocWorkerConfigKeys(ConfigKeys):
 @dataclasses.dataclass
 class DocumentsConfig(ConfigModel):
     naming_strategy: str
+    max_size: int | None
 
-    def __init__(self, naming_strategy: str):
+    def __init__(self, naming_strategy: str, max_size: int | None):
         self.naming_strategy = consts.DocumentNamingStrategy.get(naming_strategy)
+        self.max_size = max_size
 
 
 @dataclasses.dataclass
 class ExperimentalConfig(ConfigModel):
     job_timeout: int | None
-    max_doc_size: int | None
 
 
 @dataclasses.dataclass
@@ -301,6 +302,7 @@ def __init__(self):
     def documents(self) -> DocumentsConfig:
         return DocumentsConfig(
             naming_strategy=self.get(self.keys.documents.naming_strategy),
+            max_size=self.get(self.keys.documents.max_size),
         )
 
     @property
@@ -324,7 +326,6 @@ def templates(self) -> TemplatesConfig:
     def experimental(self) -> ExperimentalConfig:
         return ExperimentalConfig(
             job_timeout=self.get(self.keys.experimental.job_timeout),
-            max_doc_size=self.get(self.keys.experimental.max_doc_size),
         )
 
     @property

packages/dsw-document-worker/dsw/document_worker/limits.py

@@ -6,11 +6,25 @@
 from .exceptions import JobError
 
 
+# Key of the generateDocument command body with the tenant limits resolved by
+# the backend, e.g. {"storage": 3000000}; without it (or with a null storage)
+# there is no storage limit. It is not part of the document context, so it is
+# removed before rendering and never reaches templates
+COMMAND_LIMITS_KEY = 'limits'
+
+
 class LimitsEnforcer:
 
+    @staticmethod
+    def take_storage_limit(doc_context: dict) -> int | None:
+        limits = doc_context.pop(COMMAND_LIMITS_KEY, None)
+        if not isinstance(limits, dict):
+            return None
+        return limits.get('storage')
+
     @staticmethod
     def check_doc_size(job_id: str, doc_size: int):
-        max_size = Context.get().app.cfg.experimental.max_doc_size
+        max_size = Context.get().app.cfg.doc.max_size
         if max_size is None or doc_size <= max_size:
             return
         raise JobError(

packages/dsw-document-worker/dsw/document_worker/worker.py

@@ -66,11 +66,11 @@ def __init__(self, command: PersistentCommand, document_uuid: str):
         self.format: Format | None = None
         self.tenant_uuid = self.ctx.tenant_uuid = command.tenant_uuid
         self.doc_uuid: str = document_uuid
-        self.doc_context: dict = command.body
+        self.doc_context: dict = dict(command.body)
         self.doc: DBDocument | None = None
         self.final_file: DocumentFile | None = None
         self.template_config: TemplateConfig | None = None
-        self.tenant_limits = self.ctx.app.db.fetch_tenant_limits(self.tenant_uuid)
+        self.storage_limit = LimitsEnforcer.take_storage_limit(self.doc_context)
 
     @property
     def safe_doc(self) -> DBDocument:
@@ -264,13 +264,12 @@ def build_document(self):
             job_id=self.doc_uuid,
             doc_size=final_file.byte_size,
         )
-        limit_size = None if self.tenant_limits is None else self.tenant_limits.storage
         used_size = self.ctx.app.db.get_currently_used_size(tenant_uuid=self.tenant_uuid)
         LimitsEnforcer.check_size_usage(
             job_id=self.doc_uuid,
             doc_size=final_file.byte_size,
             used_size=used_size,
-            limit_size=limit_size,
+            limit_size=self.storage_limit,
         )
         # finalize
         self.final_file = final_file

packages/dsw-document-worker/tests/test_limits.py

@@ -0,0 +1,65 @@
+import types
+
+import pytest
+
+from dsw.document_worker.config import DocumentsConfig
+from dsw.document_worker.context import Context
+from dsw.document_worker.exceptions import JobError
+from dsw.document_worker.limits import COMMAND_LIMITS_KEY, LimitsEnforcer
+
+
+@pytest.fixture
+def with_max_size(tmp_path):
+    original = Context._instance
+
+    def initialize(max_size: int | None):
+        config = types.SimpleNamespace(
+            doc=DocumentsConfig(naming_strategy='sanitize', max_size=max_size),
+        )
+        Context.initialize(db=None, s3=None, config=config, workdir=tmp_path)
+
+    yield initialize
+    Context._instance = original
+
+
+def test_limit_from_command_body():
+    doc_context = {'document': {}, COMMAND_LIMITS_KEY: {'storage': 3_000_000}}
+    assert LimitsEnforcer.take_storage_limit(doc_context) == 3_000_000
+    assert COMMAND_LIMITS_KEY not in doc_context
+
+
+@pytest.mark.parametrize('doc_context', [
+    {},
+    {COMMAND_LIMITS_KEY: None},
+    {COMMAND_LIMITS_KEY: {}},
+    {COMMAND_LIMITS_KEY: {'storage': None}},
+])
+def test_no_limit_from_command_body(doc_context):
+    assert LimitsEnforcer.take_storage_limit(doc_context) is None
+    assert COMMAND_LIMITS_KEY not in doc_context
+
+
+def test_size_usage_within_limit():
+    LimitsEnforcer.check_size_usage('job', doc_size=100, used_size=800, limit_size=1_000)
+    LimitsEnforcer.check_size_usage('job', doc_size=10**9, used_size=10**9, limit_size=None)
+
+
+def test_size_usage_over_limit():
+    with pytest.raises(JobError, match='No space left'):
+        LimitsEnforcer.check_size_usage('job', doc_size=300, used_size=800, limit_size=1_000)
+
+
+def test_doc_size_within_max_size(with_max_size):
+    with_max_size(5_000_000)
+    LimitsEnforcer.check_doc_size('job', doc_size=5_000_000)
+
+
+def test_doc_size_over_max_size(with_max_size):
+    with_max_size(5_000_000)
+    with pytest.raises(JobError, match='exceeded size limit'):
+        LimitsEnforcer.check_doc_size('job', doc_size=5_000_001)
+
+
+def test_doc_size_without_max_size(with_max_size):
+    with_max_size(None)
+    LimitsEnforcer.check_doc_size('job', doc_size=10**12)

packages/dsw-tdk/CHANGELOG.md

@@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - New `pot` command creating a POT file with translatable strings of the template project
-- New `render` command rendering a document from the local template project and a document context (JSON file), optionally with a PO file and project files, without any DSW instance; rendering is done by `dsw-templating`, the same engine as the document worker uses, and only the files selected by `_tdk.files` are available to it, as on the server; the context gets the document worker's defaults (`config` service information and branding fallbacks, missing `extras`), which can be overridden with `-D NAME=VALUE` or the worker's `DOCUMENT_CONTEXT_*` environment variables; `--secret` / `--allow-requests` provide the `secrets` / `requests` globals and `--pandoc-filters` / `--pandoc-templates` the directories of the worker's Pandoc filters and templates
+- New `render` command rendering a document from the local template project and a document context (JSON file), optionally with a PO file and project files, without any Wizard instance; rendering is done by `dsw-templating`, the same engine as the document worker uses, and only the files selected by `_tdk.files` are available to it, as on the server; the context gets the document worker's defaults (`config` service information and branding fallbacks, missing `extras`), which can be overridden with `-D NAME=VALUE` or the worker's `DOCUMENT_CONTEXT_*` environment variables; `--secret` / `--allow-requests` provide the `secrets` / `requests` globals and `--pandoc-filters` / `--pandoc-templates` the directories of the worker's Pandoc filters and templates
 - `all` extra (`pip install 'dsw-tdk[all]'`) with the optional dependencies of rendering steps (PDF, Excel, RDF, HTTP)
 - `language` field in `template.json` (prompted by `dsw-tdk new`, defaults to `en`)
 - `put --workspace UUID` (or `DSW_WORKSPACE`) creates a new draft in a workspace; by default it is created for the whole organization