CDB-3303 Cartodbfication also supporting optional schemas

scripts-available/CDB_CartodbfyTable.sql

@@ -25,7 +25,7 @@ $$ LANGUAGE plpgsql VOLATILE;
 
 -- Ensure a table is a "cartodb" table
 -- See https://github.com/CartoDB/cartodb/wiki/CartoDB-user-table
-CREATE OR REPLACE FUNCTION CDB_CartodbfyTable(reloid REGCLASS)
+CREATE OR REPLACE FUNCTION CDB_CartodbfyTable(schema_name TEXT, reloid REGCLASS)
 RETURNS void 
 AS $$
 DECLARE
@@ -40,9 +40,14 @@ DECLARE
   exists_geom_cols BOOLEAN[];
 BEGIN
 
+  IF cartodb.schema_exists(schema_name) = false THEN
+    RAISE EXCEPTION 'Invalid schema name "%"', schema_name;
+  END IF;
+
   -- TODO: Check that user quota is set ?
   BEGIN
-    PERFORM public._CDB_UserQuotaInBytes();
+  -- Content will be discarded
+    EXECUTE FORMAT('SELECT "%I"._CDB_UserQuotaInBytes();', schema_name::text) INTO sql;
   EXCEPTION WHEN undefined_function THEN
     RAISE EXCEPTION 'Please set user quota before cartodbfying tables.';
   END;
@@ -450,13 +455,27 @@ cname, rec2.typname;
 
   sql := 'CREATE TRIGGER test_quota BEFORE UPDATE OR INSERT ON '
       || reloid::text
-      || ' EXECUTE PROCEDURE public.CDB_CheckQuota(1)';
+      || ' EXECUTE PROCEDURE public.CDB_CheckQuota(1, ''-1'', '''
+      || schema_name::text
+      || ''')';
   EXECUTE sql;
 
   sql := 'CREATE TRIGGER test_quota_per_row BEFORE UPDATE OR INSERT ON '
       || reloid::text
-      || ' FOR EACH ROW EXECUTE PROCEDURE public.CDB_CheckQuota(0.001)';
+      || ' FOR EACH ROW EXECUTE PROCEDURE public.CDB_CheckQuota(0.001, ''-1'', '''
+      || schema_name::text
+      || ''')';
   EXECUTE sql;
  
 END;
 $$ LANGUAGE PLPGSQL;
+
+
+CREATE OR REPLACE FUNCTION CDB_CartodbfyTable(reloid REGCLASS)
+RETURNS void
+AS $$
+BEGIN
+  PERFORM public.CDB_CartodbfyTable('public', reloid);
+END;
+$$
+LANGUAGE PLPGSQL;

scripts-available/CDB_ExtensionUtils.sql

@@ -15,6 +15,6 @@ $$ language 'plpgsql' VOLATILE;
 CREATE OR REPLACE FUNCTION cartodb.schema_exists(schema_name text)
 RETURNS boolean AS
 $$
-  SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = schema_name);
+  SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = schema_name::text);
 $$
 language sql VOLATILE;

scripts-available/CDB_Quota.sql

@@ -27,34 +27,41 @@ $$
 $$
 LANGUAGE 'sql' VOLATILE;
 
-
-CREATE OR REPLACE FUNCTION CDB_CheckQuota(schema_name text)
+-- Triggers cannot have declared arguments: pbfact float8, qmax int8, schema_name text
+CREATE OR REPLACE FUNCTION CDB_CheckQuota()
 RETURNS trigger AS
 $$
 DECLARE
   pbfact float8;
   qmax int8;
+  schema_name text;
   dice float8;
   quota float8;
 BEGIN
-
+  IF TG_NARGS = 3 THEN
+    schema_name := TG_ARGV[2];
+    IF cartodb.schema_exists(schema_name) = false THEN
+      RAISE EXCEPTION 'Invalid schema name "%"', schema_name;
+    END IF;
+  ELSE
+    schema_name := 'public';
+  END IF;
+  -- Hack to support old versions of CDB_CheckQuota with 2 params but without schema_name
+  IF TG_NARGS >= 2 AND TG_ARGV[1] <> '-1' THEN
+    qmax := TG_ARGV[1];
+  ELSE
+    BEGIN
+      EXECUTE FORMAT('SELECT "%I"._CDB_UserQuotaInBytes();', schema_name) INTO qmax;
+      EXCEPTION WHEN undefined_function THEN
+      RAISE EXCEPTION 'Missing "%"._CDB_UserQuotaInBytes()', schema_name;
+    END;
+  END IF;
   pbfact := TG_ARGV[0];
-  dice := random();
 
-  -- RAISE DEBUG 'CDB_CheckQuota enter: pbfact=% dice=%', pbfact, dice;
+  dice := random();
 
   IF dice < pbfact THEN
     RAISE DEBUG 'Checking quota on table % (dice:%, needed:<%)', TG_RELID::text, dice, pbfact;
-    BEGIN
-      qmax := public._CDB_UserQuotaInBytes(schema_name);
-    EXCEPTION WHEN undefined_function THEN
-      IF TG_NARGS > 1 THEN
-        RAISE NOTICE 'Using quota specified via trigger parameter';
-        qmax := TG_ARGV[1];
-      ELSE
-        RAISE EXCEPTION 'Missing _CDB_UserQuotaInBytes(), and no quota provided as parameter';
-      END IF;
-    END;
 
     IF qmax = 0 THEN
       RETURN NEW;
@@ -62,10 +69,9 @@ BEGIN
 
     SELECT public.CDB_UserDataSize(schema_name) INTO quota;
     IF quota > qmax THEN
-        RAISE EXCEPTION 'Quota exceeded by %KB', (quota-qmax)/1024;
+      RAISE EXCEPTION 'Quota exceeded by %KB', (quota-qmax)/1024;
     ELSE RAISE DEBUG 'User quota in bytes: % < % (max allowed)', quota, qmax;
     END IF;
-  -- ELSE RAISE DEBUG 'Not checking quota on table % (dice:%, needed:<%)', TG_RELID::text, dice, pbfact;
   END IF;
 
   RETURN NEW;
@@ -74,15 +80,7 @@ $$
 LANGUAGE 'plpgsql' VOLATILE;
 
 
-CREATE OR REPLACE FUNCTION CDB_CheckQuota()
-RETURNS trigger AS
-$$
-  SELECT public.CDB_CheckQuota('public');
-$$
-LANGUAGE 'plpgsql' VOLATILE;
-
-
-CREATE OR REPLACE FUNCTION CDB_SetUserQuotaInBytes(bytes int8, schema_name text)
+CREATE OR REPLACE FUNCTION CDB_SetUserQuotaInBytes(schema_name text, bytes int8)
 RETURNS int8 AS
 $$
 DECLARE
@@ -90,17 +88,17 @@ DECLARE
   schema_ok boolean;
   sql text;
 BEGIN
-  IF cartodb.schema_exists(schema_name) = false THEN
-    RAISE EXCEPTION 'Invalid schema name "%"', schema_name;
+  IF cartodb.schema_exists(schema_name::text) = false THEN
+    RAISE EXCEPTION 'Invalid schema name "%"', schema_name::text;
   END IF;
 
   BEGIN
-    EXECUTE FORMAT('current_quota := %I._CDB_UserQuotaInBytes();', schema_name);
+    EXECUTE FORMAT('SELECT "%I"._CDB_UserQuotaInBytes();', schema_name::text) INTO current_quota;
   EXCEPTION WHEN undefined_function THEN
     current_quota := 0;
   END;
 
-  sql := 'CREATE OR REPLACE FUNCTION ' || schema_name || '._CDB_UserQuotaInBytes() '
+  sql := 'CREATE OR REPLACE FUNCTION "' || schema_name::text || '"._CDB_UserQuotaInBytes() '
     || 'RETURNS int8 AS $X$ SELECT ' || bytes
     || '::int8 $X$ LANGUAGE sql IMMUTABLE';
   EXECUTE sql;
@@ -114,6 +112,8 @@ LANGUAGE 'plpgsql' VOLATILE STRICT;
 CREATE OR REPLACE FUNCTION CDB_SetUserQuotaInBytes(bytes int8)
 RETURNS int8 AS
 $$
-  SELECT public.CDB_SetUserQuotaInBytes(bytes, 'public');
+BEGIN
+  return public.CDB_SetUserQuotaInBytes('public', bytes);
+END;
 $$
 LANGUAGE 'plpgsql' VOLATILE STRICT;

test/CDB_CartodbfyTableTest.sql

@@ -30,7 +30,7 @@ BEGIN
     END IF;
   END LOOP;
 
-  PERFORM CDB_CartodbfyTable(tabname); 
+  PERFORM CDB_CartodbfyTable('public', tabname);
 
   sql := 'INSERT INTO ' || tabname::text || '(the_geom) values ( CDB_LatLng(2,1) ) RETURNING cartodb_id';
   EXECUTE sql INTO STRICT id;
@@ -132,7 +132,7 @@ LANGUAGE 'plpgsql';
 
 -- table with single non-geometrical column
 CREATE TABLE t AS SELECT 1::int as a;
-SELECT CDB_CartodbfyTable('t'); -- should fail
+SELECT CDB_CartodbfyTable('public', 't'); -- should fail
 SELECT CDB_SetUserQuotaInBytes(0); -- Set user quota to infinite
 SELECT CDB_CartodbfyTableCheck('t', 'single non-geometrical column');
 DROP TABLE t;

test/CDB_QuotaTest.sql

@@ -4,7 +4,7 @@ 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);
+      EXECUTE PROCEDURE CDB_CheckQuota(1, 1, 'public');
 INSERT INTO big VALUES (1); -- allowed, check runs before
 INSERT INTO big VALUES (1); -- disallowed, quota exceeds before
 SELECT CDB_SetUserQuotaInBytes(0);