🗃️ Add an arbitrary limit on assets that can be placed in a tiled collection

tiled/_tests/test_asset_access.py

@@ -1,5 +1,7 @@
 import hashlib
+import os
 from pathlib import Path
+from unittest import mock
 
 import pandas
 import pytest
@@ -124,3 +126,14 @@ def test_do_not_expose_raw_assets(tmpdir):
         client.write_array([1, 2, 3], key="x")
         with fail_with_status_code(HTTP_403_FORBIDDEN):
             client["x"].raw_export(tmpdir / "exported")
+
+
+@mock.patch.dict(os.environ, {"TILED_ASSET_LIMIT": "3"})
+def test_asset_limit(client):
+    # Recreate Trigger insertions
+    client.write_array([1, 2, 3])
+    client.write_array([4, 5, 6])
+    client.write_array([7, 8, 9])
+    # Failover
+    client.write_array([10, 11, 12])
+    # Clear Trigger insertions

tiled/catalog/orm.py

@@ -1,3 +1,4 @@
+import os
 from typing import List
 
 from sqlalchemy import (
@@ -289,6 +290,54 @@ def unique_parameter_num_null_check(target, connection, **kw):
         )
 
 
+@event.listens_for(DataSourceAssetAssociation.__table__, "after_create")
+def asset_limit_check(target, connection, **kw):
+    # Function to enforce an arbitrary limit on the number of assets
+    # associated with data sources, assists in pagination.
+    # Triggers cannot define hard limits so if there are concurrent inserts
+    # it is possible to exceed the limit slightly.
+    DEFAULT_LIMIT = int(os.getenv("TILED_ASSET_LIMIT", "100000") or "100000")
+
+    sqliteString = f"""
+CREATE TRIGGER assets_exceed_set_limit
+BEFORE INSERT ON data_source_asset_association
+WHEN (SELECT COUNT(*) FROM data_source_asset_association) >= {DEFAULT_LIMIT}
+BEGIN
+    SELECT RAISE(ABORT, 'Hard limit on number of associated assets exceeded : {DEFAULT_LIMIT}');
+END"""
+
+    postgresqlString = f"""
+CREATE OR REPLACE FUNCTION assets_exceed_limit()
+RETURNS TRIGGER AS
+$$
+BEGIN
+    IF (SELECT count(*) FROM data_source_asset_association) > {DEFAULT_LIMIT}
+    THEN
+        RAISE EXCEPTION 'Hard limit on number of associated assets exceeded : {DEFAULT_LIMIT}';
+    END IF;
+    RETURN NEW;
+END;
+$$
+LANGUAGE plpgsql;"""
+
+    if connection.engine.dialect.name == "sqlite":
+        connection.execute(
+            text(sqliteString),
+        )
+    elif connection.engine.dialect.name == "postgresql":
+        connection.execute(
+            text(postgresqlString),
+        )
+        connection.execute(
+            text(
+                """
+CREATE TRIGGER assets_exceed_set_limit
+BEFORE INSERT ON data_source_asset_association
+FOR EACH ROW EXECUTE PROCEDURE assets_exceed_limit();"""
+            )
+        )
+
+
 @event.listens_for(Node.__table__, "after_create")
 def create_index_metadata_tsvector_search(target, connection, **kw):
     # This creates a ts_vector based metadata search index for fulltext.