diff --git a/tiled/_tests/test_asset_access.py b/tiled/_tests/test_asset_access.py index 498e40cad..43285cd32 100644 --- a/tiled/_tests/test_asset_access.py +++ b/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 diff --git a/tiled/catalog/orm.py b/tiled/catalog/orm.py index 693c0f82f..1e54717f9 100644 --- a/tiled/catalog/orm.py +++ b/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.