Test script that used triggers to populate hierarchical

test-hierarchy-search-in-trigger.py

from pprint import pprint
import psycopg2
import uuid
from contextlib import closing

setup_sql = """
create table data (id int primary key, state jsonb, isc bool);
create table stage (id int primary key, state jsonb, isc bool);

create or replace function find_cid(fid int, fstate jsonb, fisc bool)
  returns int
as $$
declare
  fpid int;
  p_id int;
  p_state jsonb;
  p_isc bool;
begin
  if fisc then
     return fid;
  end if;
  fpid = (fstate->>'p')::int;
  if fpid is null then return null; end if;
  select id, state, isc from data where id = fpid
  into p_id, p_state, p_isc;
  return find_cid(p_id, p_state, p_isc);
end
$$ language plpgsql stable;

create or replace function populate_cid_triggerf() returns trigger
as $$
declare
  cid int;
begin
  cid = find_cid(NEW.id, NEW.state, NEW.isc);
  if cid is not null then
    NEW.state := NEW.state || ('{"cid": ' || cid::text || '}')::jsonb;
  end if;
  return NEW;
end
$$ language plpgsql;

create trigger populate_cid_trigger
  before insert or update
  on data
  for each row
  execute procedure populate_cid_triggerf();
"""

def l(cursor, query):
    cursor.execute(query)
    return list(cursor)

bconn = psycopg2.connect('')
bconn.autocommit = True
dbname = 'd'+uuid.uuid1().hex
print(dbname)
bcursor = bconn.cursor()
bcursor.execute("create database " + dbname)
try:
    with closing(psycopg2.connect("dbname=" + dbname)) as conn:
        with closing(conn.cursor()) as cursor:
            with conn:
                cursor.execute(setup_sql)


            def ex(sql):
                print(sql)
                cursor.execute(sql)
                return cursor

            def save(*inserts):
                with conn:
                    for insert in inserts:
                        ex(insert)
                    ex("DELETE FROM data WHERE id IN (SELECT id FROM stage)")
                    ex("""
                    INSERT INTO data (id, state, isc)
                    SELECT id, state, isc
                    FROM stage order by id""")
                    ex("truncate stage")

            with conn:
                save("""insert into stage (id, state, isc) values (0, '{}', false)""")
            with conn:
                save("""insert into stage (id, state, isc) values (1, '{"p": 0}', true)""",
                     """insert into stage (id, state, isc) values (2, '{"p": 1}', false)""")
            with conn:
                save("""insert into stage (id, state, isc) values (1, '{"p": 0, "z": 2}', true)""",
                     """insert into stage (id, state, isc) values (3, '{"p": 1}', false)""")
            with conn:
                save("""insert into stage (id, state, isc) values (4, '{"p": 1}', false)""",
                     """insert into stage (id, state, isc) values (1, '{"p": 0, "z": 2}', true)""");
            with conn:
                save("""insert into stage (id, state, isc) values (5, '{"p": 2}', false)""",
                     """insert into stage (id, state, isc) values (1, '{"p": 0, "z": 3}', true)""")
            with conn:
                save("""insert into stage (id, state, isc) values (6, '{"p": 5}', false)""",
                     """insert into stage (id, state, isc) values (1, '{"p": 0, "z": 4}', true)""")
            with conn:
                save("""insert into stage (id, state, isc) values (1, '{"p": 0, "z": 5}', true)""",
                     """insert into stage (id, state, isc) values (7, '{"p": 6}', false)""",
                     """insert into stage (id, state, isc) values (71, '{"p": 7}', false)""",
                     """insert into stage (id, state, isc) values (72, '{"p": 7}', false)""")
            with conn:
                save("""insert into stage (id, state, isc) values (81, '{"p": 8}', false)""",
                     """insert into stage (id, state, isc) values (82, '{"p": 8}', false)""",
                     """insert into stage (id, state, isc) values (8, '{"p": 6}', false)""",
                     """insert into stage (id, state, isc) values (1, '{"p": 0, "z": 5}', true)""")

            pprint(l(cursor, "select * from data order by id"))

finally:
    bcursor.execute("drop database " + dbname)