diff --git a/scripts-available/CDB_Quota.sql b/scripts-available/CDB_Quota.sql index 31eb92c..5cf24d4 100644 --- a/scripts-available/CDB_Quota.sql +++ b/scripts-available/CDB_Quota.sql @@ -36,11 +36,16 @@ BEGIN IF dice < pbfact THEN RAISE DEBUG 'Checking quota on table % (dice:%, needed:<%)', TG_RELID::text, dice, pbfact; BEGIN - qmax := public._CDB_UserQuotaInBytes(); + qmax := cartodb._CDB_UserQuotaInBytes(); EXCEPTION WHEN undefined_function THEN RAISE WARNING 'Missing _CDB_UserQuotaInBytes(), assuming no quota'; - RETURN NEW; + qmax := 0; END; + + IF qmax = 0 THEN + RETURN NEW; + END IF; + SELECT CDB_UserDataSize() INTO quota; IF quota > qmax THEN RAISE EXCEPTION 'Quota exceeded by %KB', (quota-qmax)/1024; @@ -53,3 +58,27 @@ BEGIN END; $$ LANGUAGE 'plpgsql' VOLATILE; + +CREATE OR REPLACE FUNCTION CDB_SetUserQuotaInBytes(bytes int8) +RETURNS int8 AS +$$ +DECLARE + current_quota int8; + sql text; +BEGIN + BEGIN + current_quota := cartodb._CDB_UserQuotaInBytes(); + EXCEPTION WHEN undefined_function THEN + current_quota := 0; + END; + + sql := 'CREATE OR REPLACE FUNCTION cartodb._CDB_UserQuotaInBytes() ' + || 'RETURNS int8 AS $X$ SELECT ' || bytes + || '::int8 $X$ LANGUAGE sql IMMUTABLE'; + EXECUTE sql; + + return current_quota; + +END +$$ +LANGUAGE 'plpgsql' VOLATILE STRICT; diff --git a/test/CDB_QuotaTest.sql b/test/CDB_QuotaTest.sql new file mode 100644 index 0000000..0c081de --- /dev/null +++ b/test/CDB_QuotaTest.sql @@ -0,0 +1,11 @@ +set client_min_messages to ERROR; +\set VERBOSITY terse +CREATE TABLE big(a int); +SELECT CDB_CartodbfyTable('big'); +INSERT INTO big SELECT generate_series(1,1024); +SELECT CDB_SetUserQuotaInBytes(8); +INSERT INTO big VALUES (1); +SELECT CDB_SetUserQuotaInBytes(0); +INSERT INTO big VALUES (1); +DROP TABLE big; +set client_min_messages to NOTICE; diff --git a/test/CDB_QuotaTest_expect b/test/CDB_QuotaTest_expect new file mode 100644 index 0000000..895dc1d --- /dev/null +++ b/test/CDB_QuotaTest_expect @@ -0,0 +1,10 @@ +SET +CREATE TABLE + +INSERT 0 1024 +0 +ERROR: Quota exceeded by 103.9921875KB +8 +INSERT 0 1 +DROP TABLE +SET