import psycopg2
from tabletypes import new_composite_type

if __name__ == '__main__':
    conn = psycopg2.connect('dbname=postgres user=postgres')
    cur = conn.cursor()

    # Create a simple test database
    cur.execute("""CREATE temp TABLE test1 (
        id integer primary key,
        label varchar
    );""")

    cur.execute("""CREATE temp TABLE test2 (
        id serial primary key,
        label varchar,
        test_id integer references test1(id)
    );""")

    cur.execute("""INSERT INTO test1 (id, label) VALUES
            (1, 'test1'),
            (2, 'test2');""")
    cur.execute("""INSERT INTO test2 (label, test_id) VALUES
            ('testa', 1),
            ('testb', 1),
            ('testc', 2),
            ('testd', 2);""")

    # Fetch the oids for these composite types
    cur.execute("""select oid from pg_type where typname =
        'test1';""")
    test_1_oid = cur.fetchone()[0]

    # And register its
    TEST1_TYPE = new_composite_type(test_1_oid, 'TEST1',
            ['id', 'label'],
            [20, 1043])

    cur.execute("""select oid from pg_type where typname =
    'test2';""")
    test_2_oid = cur.fetchone()[0]
    TEST2_TYPE = new_composite_type(test_2_oid, 'TEST2',
            ['id', 'label', 'test_id'],
            [20, 1043, 20])

    cur.execute("""select oid from pg_type where typname =
    '_test2';""")
    test_2_arryoid = cur.fetchone()[0]

    TEST2_ARRAY_TYPE = psycopg2.extensions.new_array_type((test_2_arryoid,),
            'TEST2_ARRAY', TEST2_TYPE)

    psycopg2.extensions.register_type(TEST1_TYPE)
    psycopg2.extensions.register_type(TEST2_TYPE)
    psycopg2.extensions.register_type(TEST2_ARRAY_TYPE)

    # Actual use case
    cur.execute("""select test1, array_agg(test2) as test2s
        from test1 inner join test2 on test1.id = test2.test_id
        group by test1;""")

    for value in cur.fetchall():
        print value[0].label
        for test2 in value[1]:
            print '\t%s' % test2.label