import time,psycopg2,random,sys

def query(s,*args,**kwargs):
    print(s)
    try:
        c.execute( s,*args,**kwargs )
        print( "%d rows" % c.rowcount )
    except Exception as e:
        print(e)
        query("ROLLBACK")
        sys.exit(-1)

def drop():
    query( "DROP TABLE IF EXISTS sales CASCADE" );
    query("COMMIT")

def build():
    query( """create table if not exists sales
(
    category_a  smallint,    -- sequential integer values from 0 - 10000
    category_b  varchar(3),  -- 3-digit ids (all numeric, padded with zeros)
    product     varchar(14), -- essentially random 14 character identifiers
    location_id varchar(5),  -- location id, 5-digit number (left padded with zeros)
    units       int,         -- value of interest
    sales       float       -- second value of interest
) partition by range (category_a);""" )
    for part in range(1,51):
        query("""CREATE UNLOGGED TABLE sales_a%03d PARTITION OF sales
    FOR VALUES FROM (%s) TO (%s);"""%(part,part,part+1))

def fill():
    query( """insert into sales
    (category_a, category_b, product, location_id, units, sales)
select cat_a,
       lpad(cat_b::varchar, 3, '0'),
       lpad(product::varchar, 14, '0'),
       lpad(location_id::varchar, 5, '0'),
       (random() * 10000)::int,
       (random() * 100000)::int
from generate_series(1, 50) cat_a
         cross join generate_series(1, 25) cat_b
         cross join generate_series(1, 10) location_id
         cross join generate_series(1, 5000) product;""")

db = psycopg2.connect("dbname=test")
c = db.cursor()
if 1:
    drop()
    build()
    fill()
    query("COMMIT")
    query("ALTER TABLE sales ADD primary key (category_a, category_b, product, location_id)")
    query("COMMIT")
    query("VACUUM ANALYZE sales" )