plpython: Add SPI cursor support

Add a function plpy.cursor that is similar to plpy.execute but uses an
SPI cursor to avoid fetching the entire result set into memory.

Jan Urbański, reviewed by Steve Singer

src/pl/plpython/sql/plpython_spi.sql

@@ -105,3 +105,119 @@ else:
 $$ LANGUAGE plpythonu;
 
 SELECT result_nrows_test();
+
+
+-- cursor objects
+
+CREATE FUNCTION simple_cursor_test() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users")
+does = 0
+for row in res:
+    if row['lname'] == 'doe':
+        does += 1
+return does
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION double_cursor_close() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users")
+res.close()
+res.close()
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_fetch() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users")
+assert len(res.fetch(3)) == 3
+assert len(res.fetch(3)) == 1
+assert len(res.fetch(3)) == 0
+assert len(res.fetch(3)) == 0
+try:
+    # use next() or __next__(), the method name changed in
+    # http://www.python.org/dev/peps/pep-3114/
+    try:
+        res.next()
+    except AttributeError:
+        res.__next__()
+except StopIteration:
+    pass
+else:
+    assert False, "StopIteration not raised"
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_mix_next_and_fetch() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users order by fname")
+assert len(res.fetch(2)) == 2
+
+item = None
+try:
+    item = res.next()
+except AttributeError:
+    item = res.__next__()
+assert item['fname'] == 'rick'
+
+assert len(res.fetch(2)) == 1
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION fetch_after_close() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users")
+res.close()
+try:
+    res.fetch(1)
+except ValueError:
+    pass
+else:
+    assert False, "ValueError not raised"
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION next_after_close() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users")
+res.close()
+try:
+    try:
+        res.next()
+    except AttributeError:
+        res.__next__()
+except ValueError:
+    pass
+else:
+    assert False, "ValueError not raised"
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_fetch_next_empty() RETURNS int AS $$
+res = plpy.cursor("select fname, lname from users where false")
+assert len(res.fetch(1)) == 0
+try:
+    try:
+        res.next()
+    except AttributeError:
+        res.__next__()
+except StopIteration:
+    pass
+else:
+    assert False, "StopIteration not raised"
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_plan() RETURNS SETOF text AS $$
+plan = plpy.prepare(
+    "select fname, lname from users where fname like $1 || '%' order by fname",
+    ["text"])
+for row in plpy.cursor(plan, ["w"]):
+    yield row['fname']
+for row in plpy.cursor(plan, ["j"]):
+    yield row['fname']
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_plan_wrong_args() RETURNS SETOF text AS $$
+plan = plpy.prepare("select fname, lname from users where fname like $1 || '%'",
+                    ["text"])
+c = plpy.cursor(plan, ["a", "b"])
+$$ LANGUAGE plpythonu;
+
+SELECT simple_cursor_test();
+SELECT double_cursor_close();
+SELECT cursor_fetch();
+SELECT cursor_mix_next_and_fetch();
+SELECT fetch_after_close();
+SELECT next_after_close();
+SELECT cursor_fetch_next_empty();
+SELECT cursor_plan();
+SELECT cursor_plan_wrong_args();

src/pl/plpython/sql/plpython_subtransaction.sql

@@ -242,3 +242,55 @@ SELECT pk_violation_inside_subtransaction();
 SELECT * FROM subtransaction_tbl;
 
 DROP TABLE subtransaction_tbl;
+
+-- cursor/subtransactions interactions
+
+CREATE FUNCTION cursor_in_subxact() RETURNS int AS $$
+with plpy.subtransaction():
+    cur = plpy.cursor("select * from generate_series(1, 20) as gen(i)")
+    cur.fetch(10)
+fetched = cur.fetch(10);
+return int(fetched[5]["i"])
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_aborted_subxact() RETURNS int AS $$
+try:
+    with plpy.subtransaction():
+        cur = plpy.cursor("select * from generate_series(1, 20) as gen(i)")
+        cur.fetch(10);
+        plpy.execute("select no_such_function()")
+except plpy.SPIError:
+    fetched = cur.fetch(10)
+    return int(fetched[5]["i"])
+return 0 # not reached
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_plan_aborted_subxact() RETURNS int AS $$
+try:
+    with plpy.subtransaction():
+        plpy.execute('create temporary table tmp(i) '
+                     'as select generate_series(1, 10)')
+        plan = plpy.prepare("select i from tmp")
+        cur = plpy.cursor(plan)
+        plpy.execute("select no_such_function()")
+except plpy.SPIError:
+    fetched = cur.fetch(5)
+    return fetched[2]["i"]
+return 0 # not reached
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION cursor_close_aborted_subxact() RETURNS boolean AS $$
+try:
+    with plpy.subtransaction():
+        cur = plpy.cursor('select 1')
+        plpy.execute("select no_such_function()")
+except plpy.SPIError:
+    cur.close()
+    return True
+return False # not reached
+$$ LANGUAGE plpythonu;
+
+SELECT cursor_in_subxact();
+SELECT cursor_aborted_subxact();
+SELECT cursor_plan_aborted_subxact();
+SELECT cursor_close_aborted_subxact();