CREATE OR REPLACE FUNCTION replaces_self() RETURNS text LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

yield "foo"

yield "bar"

$python$;

SELECT i, replaces_self() FROM generate_series(0, 5) AS g(i);

CREATE OR REPLACE FUNCTION replaces_self(i int) RETURNS text LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main(i):

yield "bar: %d,%d" % (i, (yield "foo: " + str(i))[0])

$python$;

-- getting new args, right?

SELECT i, replaces_self(i) FROM generate_series(0, 5) AS g(i);

CREATE OR REPLACE FUNCTION scount() RETURNS int LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

i=0

while True:

yield i

i = i + 1

$python$;

-- simple count-with-me

SELECT scount() AS sc, i FROM generate_series(0, 20-1) AS g(i);

CREATE OR REPLACE FUNCTION srcount(i int) RETURNS int LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main(i):

n = 0

while True:

n = i + n

i, = (yield n)

$python$;

-- takes a parameter to work properly

SELECT srcount(i) AS sc, i FROM generate_series(0, 20-1) AS g(i);

CREATE OR REPLACE FUNCTION srcount(i int, j int) RETURNS int LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main(i, j):

k = 0

while True:

k = k + 1

i, j = (yield i + j + k)

$python$;

-- takes a couple parameters

SELECT srcount(i, 0) AS sc1, srcount(i, 10) AS sc2, i FROM generate_series(0, 20-1) AS g(i);

-- non-composite stateful function used in materialization

CREATE OR REPLACE FUNCTION mat_stateful() RETURNS text LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

yield 'should be seen'

yield 'should not be seen'

$python$;

-- it's not an SRF, but it's being materialized.

-- This test makes sure that something surprising doesn't happen.

SELECT mat_stateful();

SELECT * FROM mat_stateful();

CREATE OR REPLACE FUNCTION yay_stateful() RETURNS text language python as

$$

from Postgres import Stateful

@Stateful

def main():

yield 'foo'

$$;

CREATE OR REPLACE FUNCTION load_yay_stateful() RETURNS text language python as

$$

def main():

fn = proc('yay_stateful()')

fn.load_module()

return 'success'

$$;

SELECT load_yay_stateful();

CREATE OR REPLACE FUNCTION call_yay_stateful() RETURNS text language python as

$$

import sys

def main():

fn = proc('yay_stateful()')

assert fn() == 'foo'

fn_mod = sys.modules[fn.oidstr]

fn_mod.main()

return 'fail'

$$;

-- should bomb out; invalid context

SELECT call_yay_stateful();

-- CORNERS --

-- not really stateful

CREATE OR REPLACE FUNCTION not_stateful() RETURNS int LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

return "it lied"

$python$;

SELECT not_stateful();

-- not really stateful, but it is an iterator

CREATE OR REPLACE FUNCTION kinda_stateful() RETURNS int LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

return iter("it lied")

$python$;

SELECT kinda_stateful();

-- SRFs --

-- Not that I would recommend it.

CREATE OR REPLACE FUNCTION stateful_srf() RETURNS SETOF text LANGUAGE python AS

$python$

from Postgres import Stateful

@Stateful

def main():

yield ("foo","bar")

yield ("fluffy","bunnies")

$python$;

SELECT stateful_srf() FROM generate_series(1, 2) AS g(i);

-- TRFs --

-- Scary. I know.

CREATE OR REPLACE FUNCTION stateful_trigger() RETURNS TRIGGER LANGUAGE python AS

$python$

from Postgres import NOTICE, Stateful

@Stateful

def before_insert(td, new):

assert new[0] == 'fluffy'

td, new = (yield ('foo',))

assert new[0] == 'bunnies'

yield ('bar',)

assert "fail" is True

@Stateful

def before_update(td, old, new):

NOTICE('foo')

td, old, new = yield new

NOTICE('bar')

yield new

$python$;

DROP TABLE IF EXISTS stateful_trigger_table CASCADE;

CREATE TABLE stateful_trigger_table(t text);

CREATE TRIGGER stateful_trigger_trigger

BEFORE INSERT OR UPDATE ON stateful_trigger_table

FOR EACH ROW

EXECUTE PROCEDURE stateful_trigger();

BEGIN;

INSERT INTO stateful_trigger_table VALUES ('fluffy'), ('bunnies');

SELECT * FROM stateful_trigger_table;

UPDATE stateful_trigger_table SET t = NULL;

SELECT COUNT(t IS NULL) AS count FROM stateful_trigger_table;

COMMIT;

CREATE OR REPLACE FUNCTION too_much_state() RETURNS int LANGUAGE python AS

$$

from Postgres import Stateful

@Stateful

def too():

yield 5

yield 7

@Stateful

def main():

yield 1

yield too()

$$;

SELECT too_much_state() FROM generate_series(1, 10) AS g(i);

CREATE OR REPLACE FUNCTION self_dependent_state() RETURNS int LANGUAGE python AS

$$

from Postgres import Stateful

@Stateful

def main():

yield main()

$$;

SELECT self_dependent_state();