Postgres stored procedures in python.

dbgod.py

from inspect import getfullargspec, getsourcelines

type_map = {
    int: 'integer'
}

tmpl = """
CREATE FUNCTION {func_name} ({args})
  RETURNS {return_t}
AS $$
{body}
$$ LANGUAGE plpythonu;
"""


class StoredProcedure(object):
    def __init__(self, f):
        self.func = f
        self.__generate_schema()

    def __generate_schema(self):
        function_name = self.func.__name__
        spec = getfullargspec(self.func)

        # Parse the arguments
        psql_args = []
        for arg in spec.args:
            if not arg in spec.annotations:
                raise SyntaxWarning(
                    'Parameter %s requires a type annotation' % arg)

            arg_t = spec.annotations[arg]
            if not arg_t in type_map:
                raise ValueError(
                    "Type '%s' not supported for parameter %s" % (arg_t, arg))

            psql_args.append('%s %s' % (arg, type_map[arg_t]))

        # Parse the return type
        if not 'return' in spec.annotations:
            raise SyntaxWarning(
                'A return type is required for function %s' % function_name)

        return_t = spec.annotations['return']
        if not return_t in type_map:
            raise ValueError(
                "Type '%s' not supported for return type of function %s" % (
                    arg_t, function_name))

        source = ''.join(getsourcelines(self.func)[0][1:])

        self._schema = tmpl.format(
            func_name=function_name,
            args=', '.join(psql_args),
            return_t=type_map[return_t],
            body=source
        )

    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)


def stored_procedure(f):
    return StoredProcedure(f)


# Usage
# @stored_procedure
# def pymax(a:int, b:int) -> int:
#     if a > b:
#         return a
#     return b
#

# TODO: create the plpy module to run the queries locally, and perhaps restrict global usage.