-- This is, sadly, a grab bag of SPI tests.
--
CREATE OR REPLACE FUNCTION cursor_id() RETURNS text LANGUAGE python AS
$python$
def main():
	return prepare('select 1').rows().cursor_id
$python$;
SELECT cursor_id();
     cursor_id      
--------------------
 <unnamed portal 1>
(1 row)

CREATE OR REPLACE FUNCTION prepare_in_module() RETURNS text LANGUAGE python AS
$python$
import Postgres
assert prepare is Postgres.Statement

ps = prepare("SELECT 1").string

def main():
	return ps
$python$;
SELECT prepare_in_module();
 prepare_in_module 
-------------------
 SELECT 1
(1 row)

CREATE OR REPLACE FUNCTION prepare_and_forget(sql text) RETURNS VOID LANGUAGE python AS
$python$

def main(sql):
	prepare(sql)
$python$;
SELECT prepare_and_forget('SELECT 1');
 prepare_and_forget 
--------------------
 
(1 row)

SELECT prepare_and_forget('SELECT 1;');
 prepare_and_forget 
--------------------
 
(1 row)

SELECT prepare_and_forget('CREATE TABLE foo (i int)');
 prepare_and_forget 
--------------------
 
(1 row)

SELECT prepare_and_forget('SELECT $1::int');
 prepare_and_forget 
--------------------
 
(1 row)

-- Now some failures
SELECT prepare_and_forget('SELECT 1; SELECT 2;');
ERROR:  cannot insert multiple commands into a prepared statement
CONTEXT:  SQL statement "SELECT 1; SELECT 2;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('INSERT INTO foo_does_not_exist VALUES ()');
ERROR:  syntax error at or near ")"
LINE 1: INSERT INTO foo_does_not_exist VALUES ()
                                               ^
QUERY:  INSERT INTO foo_does_not_exist VALUES ()
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('INSERT INTO foo_does_not_exist VALUES ($1)');
ERROR:  relation "foo_does_not_exist" does not exist
LINE 1: INSERT INTO foo_does_not_exist VALUES ($1)
                    ^
QUERY:  INSERT INTO foo_does_not_exist VALUES ($1)
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('INSERT INTO foo_does_not_exist VALUES ($1, $2)');
ERROR:  relation "foo_does_not_exist" does not exist
LINE 1: INSERT INTO foo_does_not_exist VALUES ($1, $2)
                    ^
QUERY:  INSERT INTO foo_does_not_exist VALUES ($1, $2)
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
-- restricted commands
SELECT prepare_and_forget('BEGIN;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "BEGIN;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('COMMIT;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "COMMIT;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('ABORT;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "ABORT;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('SAVEPOINT foo;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "SAVEPOINT foo;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('ROLLBACK TO foo;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "ROLLBACK TO foo;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('RELEASE foo;');
ERROR:  cannot execute transaction statements
HINT:  The "xact()" context manager can be used to manage subtransactions.
CONTEXT:  SQL statement "RELEASE foo;"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
SELECT prepare_and_forget('COPY pg_type FROM STDIN');
ERROR:  cannot execute COPY TO STDOUT or FROM STDIN statements
CONTEXT:  SQL statement "COPY pg_type FROM STDIN"
[exception from Python]
Traceback (most recent call last):
   File "public.prepare_and_forget(text)", line 4, in main
    prepare(sql)
 Postgres.Exception

[public.prepare_and_forget(text)]
-- Look at output.
CREATE OR REPLACE FUNCTION prepare_and_str(sql text) RETURNS text LANGUAGE python AS
$python$
def main(sql):
	return str([tuple(map(str, x)) for x in prepare(sql).rows()])
$python$;
SELECT prepare_and_str('SELECT 1');
 prepare_and_str 
-----------------
 [('1',)]
(1 row)

SELECT prepare_and_str('SELECT 1 UNION ALL SELECT 2');
 prepare_and_str  
------------------
 [('1',), ('2',)]
(1 row)

-- Basic I/O statement and coercion from text to target type.
CREATE OR REPLACE FUNCTION args_prepare_and_str(sql text, args text) RETURNS text LANGUAGE python AS
$python$
def main(sql, args):
	params = str(args).split(':')
	if not params:
		params = []
	ps = prepare(sql)
	return str([tuple(map(str, x)) for x in ps.rows(*params)])
$python$;
SELECT args_prepare_and_str('SELECT $1::text, $2::text', 'foo:bar');
 args_prepare_and_str 
----------------------
 [('foo', 'bar')]
(1 row)

SELECT args_prepare_and_str('SELECT $1::text, $2::text, $3::int', 'foo:bar:11');
  args_prepare_and_str  
------------------------
 [('foo', 'bar', '11')]
(1 row)

SELECT args_prepare_and_str('SELECT $1::text, $2::text, $3::int, $4::bigint', 'foo:bar:11:123213424322');
          args_prepare_and_str          
----------------------------------------
 [('foo', 'bar', '11', '123213424322')]
(1 row)

SELECT args_prepare_and_str('SELECT $1::text, $2::text, $3::int, $4::bigint', 'foo:bar:11:FAIL');
ERROR:  invalid input syntax for integer: "FAIL"
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.args_prepare_and_str(text,text)", line 7, in main
    return str([tuple(map(str, x)) for x in ps.rows(*params)])
 Postgres.Exception

[public.args_prepare_and_str(text,text)]
-- SPI works with xact()? failure caused by bad statement.
CREATE OR REPLACE FUNCTION spi_trap_failure() RETURNS text LANGUAGE python AS
$python$
import Postgres

def main():
	try:
		with xact():
			ps = prepare('selekt 1')
	except:
		pass
	else:
		raise ValueError("failed")
	# double check that it's not in error.
	prepare("SELECT 1")
	return "xact rolled back, and exception caught from prepare() failure"
$python$;
-- Check for appropriate teardown with repeats
SELECT spi_trap_failure();
                       spi_trap_failure                        
---------------------------------------------------------------
 xact rolled back, and exception caught from prepare() failure
(1 row)

SELECT spi_trap_failure();
                       spi_trap_failure                        
---------------------------------------------------------------
 xact rolled back, and exception caught from prepare() failure
(1 row)

BEGIN;
SELECT spi_trap_failure();
                       spi_trap_failure                        
---------------------------------------------------------------
 xact rolled back, and exception caught from prepare() failure
(1 row)

SELECT spi_trap_failure();
                       spi_trap_failure                        
---------------------------------------------------------------
 xact rolled back, and exception caught from prepare() failure
(1 row)

ABORT;
-- SPI works with xact()? failure caused at execution time.
--
CREATE OR REPLACE FUNCTION spi_raise_exc() RETURNS int LANGUAGE python AS
$python$
def main():
	raise ValueError("fail pls")
$python$;
CREATE OR REPLACE FUNCTION cursor_spi_trap_failure() RETURNS text LANGUAGE python AS
$python$
def main():
	try:
		with xact():
			ps = prepare('select spi_raise_exc()')
			ps()
	except:
		pass
	else:
		raise ValueError("failed")
	# double check that it's not in error.
	prepare("SELECT 1")
	return "xact rolled back, and exception caught from execution"
$python$;
SELECT cursor_spi_trap_failure();
                cursor_spi_trap_failure                
-------------------------------------------------------
 xact rolled back, and exception caught from execution
(1 row)

-- Test a merge case
DROP TABLE IF EXISTS kvpair;
NOTICE:  table "kvpair" does not exist, skipping
CREATE TABLE kvpair (k int PRIMARY KEY, v text);
NOTICE:  CREATE TABLE / PRIMARY KEY will create implicit index "kvpair_pkey" for table "kvpair"
CREATE OR REPLACE FUNCTION merge_kv(k int, v text) RETURNS text LANGUAGE python AS
$python$

exists = prepare("SELECT TRUE FROM kvpair WHERE k = $1")
dels = prepare("DELETE FROM kvpair WHERE k = $1")
ins = prepare("INSERT INTO kvpair (k, v) VALUES ($1, $2)")
ups = prepare("UPDATE kvpair SET v = $2 WHERE k = $1")

def main(k, v):
	if exists(k):
		if v is None:
			r = dels(k)
			op = 'deleted'
		else:
			r = ups(k, v)
			op = 'updated'
	else:
		r = ins(k, v)
		op = 'inserted'
	return ':'.join((op, str(k), str(v), str(r)))
$python$;
SELECT merge_kv(1, 'one');
           merge_kv           
------------------------------
 inserted:1:one:('INSERT', 1)
(1 row)

SELECT merge_kv(1, 'two');
          merge_kv           
-----------------------------
 updated:1:two:('UPDATE', 1)
(1 row)

SELECT merge_kv(1, 'five');
           merge_kv           
------------------------------
 updated:1:five:('UPDATE', 1)
(1 row)

SELECT merge_kv(2, 'two');
           merge_kv           
------------------------------
 inserted:2:two:('INSERT', 1)
(1 row)

SELECT merge_kv(4, 'f');
          merge_kv          
----------------------------
 inserted:4:f:('INSERT', 1)
(1 row)

SELECT merge_kv(0, 'zero');
           merge_kv            
-------------------------------
 inserted:0:zero:('INSERT', 1)
(1 row)

SELECT merge_kv(4, 'fours');
           merge_kv            
-------------------------------
 updated:4:fours:('UPDATE', 1)
(1 row)

SELECT merge_kv(1, 'one');
          merge_kv           
-----------------------------
 updated:1:one:('UPDATE', 1)
(1 row)

SELECT merge_kv(1, NULL);
           merge_kv           
------------------------------
 deleted:1:None:('DELETE', 1)
(1 row)

SELECT * FROM kvpair ORDER BY k ASC;
 k |   v   
---+-------
 0 | zero
 2 | two
 4 | fours
(3 rows)

-- Cache invalidating? Recovery?
CREATE OR REPLACE FUNCTION return_count() RETURNS bigint LANGUAGE python AS
$python$

countit = prepare("SELECT count(*) FROM kvpair")

def main():
	return countit.first()
$python$;
SELECT return_count();
 return_count 
--------------
            3
(1 row)

DROP TABLE kvpair;
SELECT return_count();
ERROR:  relation "kvpair" does not exist
LINE 1: SELECT return_count();
                             ^
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.return_count()", line 6, in main
    return countit.first()
 Postgres.Exception

[public.return_count()]
CREATE TABLE kvpair (k int PRIMARY KEY, v text);
NOTICE:  CREATE TABLE / PRIMARY KEY will create implicit index "kvpair_pkey" for table "kvpair"
SELECT return_count();
 return_count 
--------------
            0
(1 row)

-- Hold a reference to a cursor in the module.
-- This function provides the basis for cursor lifetime tests.
CREATE OR REPLACE FUNCTION cached_cursor(dump_cur bool) RETURNS bigint LANGUAGE python AS
$python$
cur = prepare("SELECT 1::bigint").rows()

def main(dump_cursor):
	global cur
	if dump_cursor:
		# use .declare() in order to validate cursor_is_closed()
		# is being used for each next().
		cur = prepare("SELECT 1::bigint").declare()
		return 0
	r = next(cur)[0]
	try:
		next(cur)
	except StopIteration:
		pass
	else:
		raise ValueError("next() on exhausted cursor did not raise StopIteration")
	return r
$python$;
SELECT cached_cursor(false);
 cached_cursor 
---------------
             1
(1 row)

-- should fail with cursor is closed
SELECT cached_cursor(false);
ERROR:  cannot use 'read' operation on a closed cursor
HINT:  All Postgres.Cursor objects are closed at the end of a transaction.
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.cached_cursor(boolean)", line 11, in main
    r = next(cur)[0]
 Postgres.Exception

[public.cached_cursor(boolean)]
BEGIN;
-- "reset" has to be inside the xact.
SELECT cached_cursor(true);
 cached_cursor 
---------------
             0
(1 row)

SELECT cached_cursor(false);
 cached_cursor 
---------------
             1
(1 row)

-- should fail with StopIteration
SELECT cached_cursor(false);
ERROR:  function's "main" raised a Python exception
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.cached_cursor(boolean)", line 11, in main
    r = next(cur)[0]
 StopIteration

[public.cached_cursor(boolean)]
ABORT;
-- Rolled back savepoint; current this shouldn't work.
BEGIN;
SAVEPOINT a;
SELECT cached_cursor(true);
 cached_cursor 
---------------
             0
(1 row)

ROLLBACK TO a;
SELECT cached_cursor(false);
ERROR:  cannot use 'read' operation on a closed cursor
HINT:  All Postgres.Cursor objects are closed at the end of a transaction.
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.cached_cursor(boolean)", line 11, in main
    r = next(cur)[0]
 Postgres.Exception

[public.cached_cursor(boolean)]
ABORT;
-- should last across *released* savepoints..
BEGIN;
SAVEPOINT a;
SELECT cached_cursor(true);
 cached_cursor 
---------------
             0
(1 row)

RELEASE a;
SELECT cached_cursor(false);
 cached_cursor 
---------------
             1
(1 row)

ABORT;
-- wasn't started in the savepoint, so it should survive.
BEGIN;
SELECT cached_cursor(true);
 cached_cursor 
---------------
             0
(1 row)

SAVEPOINT a;
SELECT TRUE FROM kvpair LIMIT 1;
 bool 
------
(0 rows)

RELEASE a;
SELECT cached_cursor(false);
 cached_cursor 
---------------
             1
(1 row)

ABORT;
DROP SEQUENCE IF EXISTS kvpair_seq;
NOTICE:  sequence "kvpair_seq" does not exist, skipping
CREATE SEQUENCE kvpair_seq;
ALTER TABLE kvpair ALTER k SET DEFAULT nextval('kvpair_seq')::int4;
-- check DML ... RETURNING
CREATE OR REPLACE FUNCTION check_dml_ret(v text) RETURNS int LANGUAGE python AS
$python$
ins = prepare("INSERT INTO kvpair (v) VALUES ($1) RETURNING k")

def main(val):
	return ins.first(val)
$python$;
SELECT check_dml_ret('foo');
 check_dml_ret 
---------------
             1
(1 row)

SELECT check_dml_ret('bar');
 check_dml_ret 
---------------
             2
(1 row)

SELECT check_dml_ret('bar');
 check_dml_ret 
---------------
             3
(1 row)

SELECT * FROM kvpair;
 k |  v  
---+-----
 1 | foo
 2 | bar
 3 | bar
(3 rows)

--
-- Postgres.Transaction functionality in conjunction with SPI.
-- In _xact.sql, the effects of subtransactions are tested using SQL functions.
-- Here, we perform many of those tests again to validate that SPI is
-- being used properly.
--
DELETE FROM kvpair;
CREATE OR REPLACE FUNCTION x_success_insert() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	with xact():
		ins('SHOULD EXIST')
	return 'success'
$python$;
SELECT x_success_insert();
 x_success_insert 
------------------
 success
(1 row)

SELECT * FROM kvpair;
 k |      v       
---+--------------
 4 | SHOULD EXIST
(1 row)

CREATE OR REPLACE FUNCTION x_fail_insert() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	try:
		with xact():
			ins('SHOULD NOT EXIST')
			raise ValueError
	except ValueError:
		return 'caught valueerror'
	return 'this failed'
$python$;
SELECT x_fail_insert();
   x_fail_insert   
-------------------
 caught valueerror
(1 row)

SELECT * FROM kvpair;
 k |      v       
---+--------------
 4 | SHOULD EXIST
(1 row)

CREATE OR REPLACE FUNCTION x_failed_xact() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (k, v) VALUES ($1, $2)')
key = prepare('SELECT k FROM kvpair ORDER BY k LIMIT 1').first()

def main():
	try:
		with xact():
			ins(key, 'INSERT SHOULD FAIL')
	except:
		return 'success, caught dberror'
	return 'this failed'
$python$;
SELECT x_failed_xact();
      x_failed_xact      
-------------------------
 success, caught dberror
(1 row)

SELECT * FROM kvpair;
 k |      v       
---+--------------
 4 | SHOULD EXIST
(1 row)

CREATE OR REPLACE FUNCTION x_recover_xact() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (k, v) VALUES ($1, $2)')
key = prepare('SELECT k FROM kvpair ORDER BY k LIMIT 1').first()
dins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	try:
		with xact():
			ins(key, 'INSERT SHOULD FAIL')
	except:
		dins('THIS SHOULD ALSO EXIST')
		return 'success'
	return 'this failed'
$python$;
SELECT x_recover_xact();
 x_recover_xact 
----------------
 success
(1 row)

SELECT * FROM kvpair;
 k |           v            
---+------------------------
 4 | SHOULD EXIST
 6 | THIS SHOULD ALSO EXIST
(2 rows)

CREATE OR REPLACE FUNCTION x_dont_recover_xact() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (k, v) VALUES ($1, $2)')
key = prepare('SELECT k FROM kvpair ORDER BY k LIMIT 1').first()
dins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	try:
		with xact():
			dins('THIS SHOULD **NOT** EXIST')
			ins(key, 'INSERT SHOULD FAIL')
	except:
		dins('this succeeds, but will fail when the following error is raised')
		raise
$python$;
SELECT x_dont_recover_xact();
ERROR:  duplicate key value violates unique constraint "kvpair_pkey"
CONTEXT:  SQL statement "INSERT INTO kvpair (k, v) VALUES ($1, $2)"
[exception from Python]
Traceback (most recent call last):
   File "public.x_dont_recover_xact()", line 10, in main
    ins(key, 'INSERT SHOULD FAIL')
 Postgres.Exception

[public.x_dont_recover_xact()]
-- No change expected.
SELECT * FROM kvpair;
 k |           v            
---+------------------------
 4 | SHOULD EXIST
 6 | THIS SHOULD ALSO EXIST
(2 rows)

--
-- Very basic load_rows & load_chunks tests.
--
CREATE OR REPLACE FUNCTION load_rows() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	return ins.load_rows((str(x),) for x in range(200))
$python$;
SELECT load_rows();
 load_rows 
-----------
 200
(1 row)

SELECT count(*) FROM kvpair;
 count 
-------
   202
(1 row)

CREATE OR REPLACE FUNCTION load_chunks() RETURNS text LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	return ins.load_chunks([
		((str(x),) for x in range(25)),
		((str(x),) for x in range(25,50)),
		((str(x),) for x in range(50,75)),
		((str(x),) for x in range(75,100)),
	])
$python$;
SELECT load_chunks();
 load_chunks 
-------------
 100
(1 row)

SELECT count(*) FROM kvpair;
 count 
-------
   302
(1 row)

-- make sure first() works with DML.
CREATE OR REPLACE FUNCTION dml_first() RETURNS int LANGUAGE python AS
$python$
ins = prepare('INSERT INTO kvpair (v) VALUES ($1)')

def main():
	return ins.first('does first work')
$python$;
-- expect 1
SELECT dml_first() = 1;
 ?column? 
----------
 t
(1 row)

--
-- Test cursor configuration restrictions
--
-- Validate that NO SCROLL cursors cannot use read() and seek()
-- read() and seek() are SCROLL interfaces.
--
CREATE OR REPLACE FUNCTION rows_cur_cant_read() RETURNS VOID LANGUAGE python AS
$python$

def main():
	cur = prepare("SELECT 1").rows()
	cur.read(10)
$python$;
SELECT rows_cur_cant_read();
ERROR:  function's "main" raised a Python exception
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.rows_cur_cant_read()", line 5, in main
    cur.read(10)
 RuntimeError: cannot use read method with NO SCROLL cursors

[public.rows_cur_cant_read()]
CREATE OR REPLACE FUNCTION rows_cur_cant_seek() RETURNS VOID LANGUAGE python AS
$python$

def main():
	cur = prepare("SELECT 1").rows()
	cur.seek(10)
$python$;
SELECT rows_cur_cant_seek();
ERROR:  function's "main" raised a Python exception
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.rows_cur_cant_seek()", line 5, in main
    cur.seek(10)
 RuntimeError: cannot seek NO SCROLL cursors

[public.rows_cur_cant_seek()]
CREATE OR REPLACE FUNCTION chunks_cur_cant_read() RETURNS VOID LANGUAGE python AS
$python$

def main():
	cur = prepare("SELECT 1").chunks()
	cur.read(10)
$python$;
SELECT chunks_cur_cant_read();
ERROR:  function's "main" raised a Python exception
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.chunks_cur_cant_read()", line 5, in main
    cur.read(10)
 RuntimeError: cannot use read method with NO SCROLL cursors

[public.chunks_cur_cant_read()]
CREATE OR REPLACE FUNCTION chunks_cur_cant_seek() RETURNS VOID LANGUAGE python AS
$python$

def main():
	cur = prepare("SELECT 1").chunks()
	cur.seek(10)
$python$;
SELECT chunks_cur_cant_seek();
ERROR:  function's "main" raised a Python exception
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.chunks_cur_cant_seek()", line 5, in main
    cur.seek(10)
 RuntimeError: cannot seek NO SCROLL cursors

[public.chunks_cur_cant_seek()]
-- validate that column_names, parameter_types, pg_parameter_types, pg_column_types
-- provide the expected information.
CREATE OR REPLACE FUNCTION spi_metadata() RETURNS text LANGUAGE python AS
$python$
import Postgres

colnames = ('a_text_column', 'a_int4_column', 'a_bytea_column')
pg_types = (Postgres.CONST['TEXTOID'], Postgres.CONST['INT4OID'], Postgres.CONST['BYTEAOID'])
types = tuple(map(Postgres.Type, pg_types))

stmt = prepare('SELECT $1::text AS a_text_column, $2::int4 AS a_int4_column, $3::bytea AS a_bytea_column')

args = ('text', 0, b'bytea')

def main():
	cur = stmt.rows(*args)

	assert cur.column_names == colnames
	assert stmt.column_names == colnames

	assert cur.column_types == types
	assert stmt.column_types == types
	assert stmt.parameter_types == types

	assert stmt.pg_parameter_types == pg_types
	assert stmt.pg_column_types == pg_types
	assert cur.pg_column_types == pg_types

	return 'success'
$python$;
SELECT spi_metadata();
 spi_metadata 
--------------
 success
(1 row)

--
-- Generate a series of positive integers and validate that seek and read work
-- as expected for both possible directions.
CREATE OR REPLACE FUNCTION scrollable() RETURNS text LANGUAGE python AS
$python$

def testScroll(direction = True):
	imin = 0
	imax = 2**16
	if direction:
		ps = prepare("SELECT i FROM generate_series(0, (2^16)::int) AS g(i)")
	else:
		ps = prepare("SELECT i FROM generate_series((2^16)::int, 0, -1) AS g(i)")

	c = ps.declare()
	c.direction = direction
	if not direction:
		c.seek(0)

	assert [x for x, in c.read(10)] == list(range(10))
	assert [x for x, in c.read(10, 'BACKWARD')] == list(range(8, -1, -1))
	c.seek(0, 2)
	# try a non-text direction parameter
	assert [x for x, in c.read(10, False)] == list(range(imax, imax-10, -1))

	# move to end
	c.seek(0, 2)
	assert [x for x, in c.read(100, 'BACKWARD')] == list(range(imax, imax-100, -1))
	# move backwards, relative
	c.seek(-100, 1)
	assert [x for x, in c.read(100, 'BACKWARD')] == list(range(imax-200, imax-300, -1))

	# move abs, again
	c.seek(14000)
	assert [x for x, in c.read(100)] == list(range(14000, 14100))
	# move forwards, relative
	c.seek(100, 1)
	assert [x for x, in c.read(100)] == list(range(14200, 14300))
	# move abs, again
	c.seek(24000)
	assert [x for x, in c.read(200)] == list(range(24000, 24200))
	# move to end and then back some
	c.seek(20, 2)
	assert [x for x, in c.read(200, 'BACKWARD')] == list(range(imax-20, imax-20-200, -1))

	c.seek(0, 2)
	c.seek(-10, 1)
	r1 = c.read(10)
	c.seek(10, 2)
	assert [x for x, in r1] == [x for x, in c.read(10)]

def main():
	testScroll()
	testScroll(direction = False)
	return 'success'
$python$;	
SELECT scrollable();
 scrollable 
------------
 success
(1 row)

-- check constant parameter use and their effect on the interfaces
CREATE OR REPLACE FUNCTION check_const_params() RETURNS text LANGUAGE python AS
$python$
ps = prepare("SELECT $1::int, $2::text, $3::numeric", 123, "text", 1.23)

meths = [
	ps,
	ps.rows,
	ps.chunks,
	ps.declare,
]

def main():
	r = [
		ps.first(),
		ps()[0],
		list(ps.rows())[0],
		list(ps.chunks())[0][0],
		ps.declare().read(1)[0],
	]
	for x in r:
		assert tuple(x) == (123, "text", "1.23")

	for x in meths:
		try:
			# params have already been specified.
			x(321, "wrong", "1.23")
		except TypeError:
			pass
	##
	# both load_rows and load_chunks should throw a TypeError
	try:
		ps.load_rows([(321,"wrong","3.12")])
	except TypeError:
		pass
	try:
		ps.load_chunks([[(321,"wrong","3.12")]])
	except TypeError:
		pass

	return 'success'
$python$;
SELECT check_const_params();
 check_const_params 
--------------------
 success
(1 row)

-- Check Postgres.eval()
CREATE OR REPLACE FUNCTION pg_eval(text) RETURNS text LANGUAGE python AS
$python$
from Postgres import eval as main
$python$;
SELECT pg_eval('''some text''::text');
  pg_eval  
-----------
 some text
(1 row)

SELECT pg_eval('1');
 pg_eval 
---------
 1
(1 row)

SELECT pg_eval('1::int4');
 pg_eval 
---------
 1
(1 row)

SELECT pg_eval('-123::int4');
 pg_eval 
---------
 -123
(1 row)

SELECT pg_eval('-123::int4 + 1000');
 pg_eval 
---------
 877
(1 row)

--
-- Check read-only restriction. Make sure to exercise a few combinations..
-- Notably, calling DML is different then getting rows, so be sure to
-- exercise the cursor_open paths with INSERT ... RETURNING
-- first, make sure that the flag is properly present in the module execution context
CREATE OR REPLACE FUNCTION do_mod_in_mod_stable() RETURNS text STABLE LANGUAGE python AS
$python$
ps = prepare("insert into kvpair (v) VALUES ('fail')")
ps()
$python$;
SELECT do_mod_in_mod_stable();
ERROR:  INSERT is not allowed in a non-volatile function
CONTEXT:  SQL statement "insert into kvpair (v) VALUES ('fail')"
[exception from Python]
Traceback (most recent call last):
   File "public.do_mod_in_mod_stable()", line 3, in <module>
    ps()
 Postgres.Exception

[public.do_mod_in_mod_stable() while loading function module]
CREATE OR REPLACE FUNCTION do_mod_in_mod_stable_rows() RETURNS text STABLE LANGUAGE python AS
$python$
ps = prepare("insert into kvpair (v) VALUES ('fail') RETURNING k, v")
r = list(ps.rows())
$python$;
SELECT do_mod_in_mod_stable_rows();
ERROR:  INSERT is not allowed in a non-volatile function
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.do_mod_in_mod_stable_rows()", line 3, in <module>
    r = list(ps.rows())
 Postgres.Exception

[public.do_mod_in_mod_stable_rows() while loading function module]
-- and finally, in main.
CREATE OR REPLACE FUNCTION do_mod_stable() RETURNS text STABLE LANGUAGE python AS
$python$
ps = prepare("insert into kvpair (v) VALUES ('fail')")
def main():
	ps()
	return 'fail'
$python$;
SELECT do_mod_stable();
ERROR:  INSERT is not allowed in a non-volatile function
CONTEXT:  SQL statement "insert into kvpair (v) VALUES ('fail')"
[exception from Python]
Traceback (most recent call last):
   File "public.do_mod_stable()", line 4, in main
    ps()
 Postgres.Exception

[public.do_mod_stable()]
CREATE OR REPLACE FUNCTION do_mod_stable_rows() RETURNS text STABLE LANGUAGE python AS
$python$
ps = prepare("insert into kvpair (v) VALUES ('fail') RETURNING k, v")
def main():
	r = list(ps.rows())
	return 'fail'
$python$;
SELECT do_mod_stable_rows();
ERROR:  INSERT is not allowed in a non-volatile function
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.do_mod_stable_rows()", line 4, in main
    r = list(ps.rows())
 Postgres.Exception

[public.do_mod_stable_rows()]
-- Just check one IMMUTABLE as we really just want to make sure
-- that the contrast from VOLATILE is being made. We already
-- checked the combinations with STABLE, and those
-- code paths don't change due to IMMUTABLE.
CREATE OR REPLACE FUNCTION do_mod_immutable_rows() RETURNS text IMMUTABLE LANGUAGE python AS
$python$
ps = prepare("insert into kvpair (v) VALUES ('fail') RETURNING k, v")
def main():
	r = list(ps.rows())
	return 'fail'
$python$;
SELECT do_mod_immutable_rows();
ERROR:  INSERT is not allowed in a non-volatile function
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.do_mod_immutable_rows()", line 4, in main
    r = list(ps.rows())
 Postgres.Exception

[public.do_mod_immutable_rows()]
-- chunksize tests; the only thing we can validate is that
-- chunks() yields the appropriate size
CREATE OR REPLACE FUNCTION check_chunksize() RETURNS text LANGUAGE python AS
$python$
ps = prepare('select i FROM generate_series(0, 2000) AS g(i)')
def main():
	c = ps.chunks()
	c.chunksize = 33
	assert len(next(c)) == 33
	c.chunksize = 51
	assert len(next(c)) == 51
	c.chunksize = 1
	assert len(next(c)) == 1
	c.chunksize = 5 
	assert len(next(c)) == 5
	assert len(next(c)) == 5
	assert len(next(c)) == 5
	return 'success(' + str(next(c)[0][0]) + ')'
$python$;
SELECT check_chunksize();
 check_chunksize 
-----------------
 success(100)
(1 row)

-- Postgres.execute --
-- Make sure that subsequent queries can reference objects created
-- by earlier statements(Postgres.execute)
CREATE OR REPLACE FUNCTION check_pg_execute() RETURNS text LANGUAGE python AS
$python$
import Postgres
def main():
	r = Postgres.execute("""
CREATE TEMP TABLE ttt (i int);
INSERT INTO ttt VALUES (1);
""")
	r = Postgres.execute("""
""")
	assert r == None
	assert sqleval('select i from ttt') == 1
	r = Postgres.execute("DROP TABLE ttt")
	assert r == None
	return 'success'
$python$;
SELECT check_pg_execute();
 check_pg_execute 
------------------
 success
(1 row)

DROP TABLE IF EXISTS ttmp CASCADE;
NOTICE:  table "ttmp" does not exist, skipping
CREATE TABLE ttmp (i int);
CREATE OR REPLACE FUNCTION insert_into_t(i int) RETURNS text LANGUAGE SQL AS
'INSERT INTO ttmp VALUES($1); SELECT ''foo''::text;';
CREATE OR REPLACE FUNCTION check_pg_execute_readall() RETURNS text LANGUAGE python AS
$python$
import Postgres

def main():
	r = Postgres.execute("""
SELECT insert_into_t(i) FROM generate_series(1, 100) g(i);
""")
	assert r == None
	return 'success'
$python$;
SELECT check_pg_execute_readall();
 check_pg_execute_readall 
--------------------------
 success
(1 row)

SELECT COUNT(*) FROM ttmp;
 count 
-------
   100
(1 row)

CREATE OR REPLACE FUNCTION check_recursive_execute() RETURNS text LANGUAGE python AS
$python$
import Postgres

sql = "SELECT check_recursive_execute()"
init = None
first = True

def main():
	global init, first
	if first is True:
		first = False
		r = Postgres.execute(sql)
	else:
		init = 'success'
	return str(init)
$python$;
SELECT check_recursive_execute();
 check_recursive_execute 
-------------------------
 success
(1 row)

-- Test empty first
CREATE OR REPLACE FUNCTION check_empty() RETURNS text LANGUAGE python AS
$python$
stmt = prepare('select 1 limit 0')
def main():
	assert stmt.first() is None
	assert stmt() == []
	return 'success'
$python$;
SELECT check_empty();
 check_empty 
-------------
 success
(1 row)

-- test column() cursors
CREATE OR REPLACE FUNCTION check_column() RETURNS SETOF text LANGUAGE python AS
$python$
stmt = prepare('SELECT i FROM generate_series(1, 10) AS g(i)').column

def main():
	yield 'begin'
	for x in stmt():
		yield str(x.__class__.typname) + ': ' + str(x)
	yield 'end'
$python$;
SELECT check_column();
 check_column 
--------------
 begin
 int4: 1
 int4: 2
 int4: 3
 int4: 4
 int4: 5
 int4: 6
 int4: 7
 int4: 8
 int4: 9
 int4: 10
 end
(12 rows)

-- make sure column-configured cursors are being properly
-- respected by the PL's materialization process.
-- the internal iterator buffer holds the individual columns, not records,
-- so the existing srf_cursor_materialize function won't be able to consume
-- buffer as it would a .rows() iterator. For now, it punts it to the generic
-- materialization mechanism.
CREATE OR REPLACE FUNCTION check_column_materialize() RETURNS SETOF text LANGUAGE python AS
$python$
stmt = prepare("SELECT i FROM (VALUES ('a'::text),('b'),('c'),('d')) AS g(i)")

def main():
	return stmt.column()
$python$;
SELECT check_column_materialize();
 check_column_materialize 
--------------------------
 a
 b
 c
 d
(4 rows)

SELECT * FROM check_column_materialize();
 check_column_materialize 
--------------------------
 a
 b
 c
 d
(4 rows)

CREATE OR REPLACE FUNCTION check_column_materialize_record(OUT text, OUT int) RETURNS SETOF record LANGUAGE python AS
$python$
stmt = prepare("SELECT i FROM (VALUES ('a'::text),('b'),('c'),('d')) AS g(i)")

def main():
	return stmt.column()
$python$;
-- fail.
SELECT check_column_materialize_record();
ERROR:  input of anonymous composite types is not implemented
SELECT * FROM check_column_materialize_record();
ERROR:  could not materialize result from returned iterable
CONTEXT:  [exception from Python]
TypeError: cannot normalize row from 'Postgres.types.text'

[public.check_column_materialize_record() while materializing]
-- check handling of VOID
CREATE OR REPLACE FUNCTION return_void() RETURNS VOID LANGUAGE python AS
'main = lambda: None';
CREATE OR REPLACE FUNCTION select_voids() RETURNS text LANGUAGE python AS
$python$
voids = list(map(tuple, prepare("""
	SELECT return_void(), return_void() UNION ALL
	SELECT return_void(), return_void()
""")()))

def main():
	assert voids == [(None,None), (None,None)]
	return 'success'
$python$;
SELECT select_voids();
 select_voids 
--------------
 success
(1 row)

-- comparable composites
CREATE OR REPLACE FUNCTION compare_statements(s1 text, s2 text) RETURNS BOOLEAN LANGUAGE python AS
$python$
def main(a,b):
	return prepare(a)() == prepare(b)()
$python$;
SELECT compare_statements('SELECT 1 AS a', 'SELECT 1 AS b');
 compare_statements 
--------------------
 t
(1 row)

SELECT compare_statements('SELECT 1 AS a', 'SELECT 1 AS a');
 compare_statements 
--------------------
 t
(1 row)

SELECT compare_statements('SELECT ''foo''::text AS a', 'SELECT 1 AS b');
ERROR:  cannot compare dissimilar column types text and integer at record column 1
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.compare_statements(text,text)", line 3, in main
    return prepare(a)() == prepare(b)()
 Postgres.Exception

[public.compare_statements(text,text)]
SELECT compare_statements('SELECT ''foo''::text AS b, 1 AS a', 'SELECT 1 AS a, ''foo''::text AS b');
ERROR:  cannot compare dissimilar column types text and integer at record column 1
CONTEXT:  [exception from Python]
Traceback (most recent call last):
   File "public.compare_statements(text,text)", line 3, in main
    return prepare(a)() == prepare(b)()
 Postgres.Exception

[public.compare_statements(text,text)]
SELECT compare_statements('SELECT NULL::int AS a', 'SELECT 1::int AS a');
 compare_statements 
--------------------
 f
(1 row)

-- invalidated anonymous composites
DROP TABLE IF EXISTS it_gets_changed;
NOTICE:  table "it_gets_changed" does not exist, skipping
CREATE TABLE it_gets_changed (c1 text, c2 text);
CREATE OR REPLACE FUNCTION invalidated_anonymous_composites() RETURNS text LANGUAGE python AS
$python$
import Postgres
original = prepare("SELECT ROW('foo','bar')::it_gets_changed")
original_type = original.output

def main():
	stmt = prepare("SELECT ROW('foo','bar')::it_gets_changed")
	if stmt.output is not original.output:
		Postgres.WARNING("output was invalidated")
	return ' . '.join(map(str, stmt.first()))
$python$;
SELECT invalidated_anonymous_composites();
 invalidated_anonymous_composites 
----------------------------------
 foo . bar
(1 row)

ALTER TABLE it_gets_changed ADD COLUMN c3 text;
ALTER TABLE it_gets_changed DROP COLUMN c3;
-- expecting an invalidated WARNING here.
SELECT invalidated_anonymous_composites();
WARNING:  output was invalidated
 invalidated_anonymous_composites 
----------------------------------
 foo . bar
(1 row)

