From b69ded8f66208bcf78c25ec5ad467a1d525bea2d Mon Sep 17 00:00:00 2001
From: Sandro Santilli <strk@keybit.net>
Date: Fri, 9 May 2014 12:55:41 +0200
Subject: [PATCH] Let CDB_CheckQuota second parameter override configured user
 quota

Closes #13
Includes automated test
---
 scripts-available/CDB_Quota.sql | 17 +++++++++++------
 test/CDB_QuotaTest.sql          |  6 ++++++
 test/CDB_QuotaTest_expect       |  3 +++
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/scripts-available/CDB_Quota.sql b/scripts-available/CDB_Quota.sql
index 4b5c97d..290c0a2 100644
--- a/scripts-available/CDB_Quota.sql
+++ b/scripts-available/CDB_Quota.sql
@@ -35,12 +35,17 @@ BEGIN
 
   IF dice < pbfact THEN
     RAISE DEBUG 'Checking quota on table % (dice:%, needed:<%)', TG_RELID::text, dice, pbfact;
-    BEGIN
-      qmax := public._CDB_UserQuotaInBytes();
-    EXCEPTION WHEN undefined_function THEN
-      RAISE WARNING 'Missing _CDB_UserQuotaInBytes(), assuming no quota';
-      qmax := 0;
-    END;
+    IF TG_NARGS > 1 THEN
+      RAISE NOTICE 'Using quota specified via trigger parameter';
+      qmax := TG_ARGV[1];
+    ELSE
+      BEGIN
+        qmax := public._CDB_UserQuotaInBytes();
+      EXCEPTION WHEN undefined_function THEN
+        RAISE WARNING 'Missing _CDB_UserQuotaInBytes(), assuming no quota';
+        qmax := 0;
+      END;
+    END IF;
 
     IF qmax = 0 THEN
       RETURN NEW;
diff --git a/test/CDB_QuotaTest.sql b/test/CDB_QuotaTest.sql
index 0c081de..a9590cd 100644
--- a/test/CDB_QuotaTest.sql
+++ b/test/CDB_QuotaTest.sql
@@ -1,6 +1,12 @@
 set client_min_messages to ERROR;
 \set VERBOSITY terse
 CREATE TABLE big(a int);
+-- Try the legacy interface
+-- See https://github.com/CartoDB/cartodb-postgresql/issues/13
+CREATE TRIGGER test_quota BEFORE UPDATE OR INSERT ON big
+      EXECUTE PROCEDURE CDB_CheckQuota(1,1);
+INSERT INTO big VALUES (1); -- allowed, check runs before
+INSERT INTO big VALUES (1); -- disallowed, quota exceeds before
 SELECT CDB_CartodbfyTable('big');
 INSERT INTO big SELECT generate_series(1,1024);
 SELECT CDB_SetUserQuotaInBytes(8);
diff --git a/test/CDB_QuotaTest_expect b/test/CDB_QuotaTest_expect
index 895dc1d..01a1866 100644
--- a/test/CDB_QuotaTest_expect
+++ b/test/CDB_QuotaTest_expect
@@ -1,5 +1,8 @@
 SET
 CREATE TABLE
+CREATE TRIGGER
+INSERT 0 1
+ERROR:  Quota exceeded by 3.9990234375KB
 
 INSERT 0 1024
 0