olirice / alembic_utils

An alembic/sqlalchemy extension for migrating sql views, functions, triggers, and policies
https://olirice.github.io/alembic_utils
MIT License
222 stars 47 forks source link

Create aggregate function #37

Open xoelop opened 3 years ago

xoelop commented 3 years ago

I'm trying to create an aggregate function to concat arrays even if they're empty, as described here

Running this directly on Postgres works

CREATE AGGREGATE array_accum (anyarray)
(
    sfunc = array_cat,
    stype = anyarray,
    initcond = '{}'
);  

But I don't know how to (or if it's even possible) do that with alembic_utils.

I tried adding this to my functions.py file:

ARRAY_ACCUM_DEFINITION = """
(
    sfunc = array_cat,
    stype = anyarray,
    initcond = '{}'
);  
"""

array_distinct = PGFunction(
    schema='public',
    signature='array_accum(anyarray)',
    definition=ARRAY_ACCUM_DEFINITION
)

But if fails because the SQL code outputted tries to create a normal function, instead of an aggregate one, and I couldn't find any references on how to create aggregate functions

[SQL: CREATE FUNCTION "public"."array_accum"(anyarray) (
    sfunc = array_cat,
    stype = anyarray,
    initcond = '{}'
)]
(Background on this error at: http://sqlalche.me/e/13/f405)

So what would be the right way to do this?

Thanks in advance, and thanks for making this library, is great to have things much cleaner

olirice commented 3 years ago

Aggregate functions aren't currently supported

The create statements are pretty similar, which makes it look like it'd be a small change to add, but aggregate functions' properties are stored in a different view, pg_aggregate vs pg_function, and the properties are significantly different.

Due to those differences, it needs its own class to correctly detect when a replace operation is needed.

I don't currently have a plan to implement them but would consider a PR that doesn't reduce coverage

That functions definition isn't likely to change much so you could consider creating it using raw sql in the migration. Since it won't be registered with the pg_function view, alembic_utils wont attempt to drop it in subsequent revisions that use --autogenerate.

Alternatively, you could avoid needing a custom aggregate by unpacking the arrays and using array_agg

select
    array_agg(unpacked.elem)
from 
    (
        VALUES
        ('{5}'::int[]),
        (null),
        ('{1,2,3}'),
        ('{}')
    ) abc(arr),
    -- Unnest the arrays
    lateral (
        select * from unnest(abc.arr)
    ) unpacked(elem)

/* returns
 --------------
| array_agg    |
| int[]        |
 --------------
| {5, 1, 2, 3} |
*/

Hope that helps

Leaving this open to track changes wrt aggregate functions in the future

xoelop commented 3 years ago

Thanks for the quick reply! I think I'll go with creating it directly on the DB.

That's the 1st aggregate function I ever created, so I'm missing too much of the big picture to be able to create a PR for this.