"""add_budget_stats_functions

Revision ID: 531c40a3f0fb
Revises: ff55991be102
Create Date: 2020-04-17 15:41:52.386361

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '531c40a3f0fb'
down_revision = 'ff55991be102'
branch_labels = None
depends_on = None


def upgrade():
    bind = op.get_bind()
    bind.execute(
        '''
        CREATE OR REPLACE FUNCTION dates_between(start_date date, end_date date) RETURNS int AS $$
        DECLARE
        BEGIN
          RETURN 1 + date_part('day', end_date::TIMESTAMP - start_date::TIMESTAMP)::INT;
        END;
        $$ LANGUAGE plpgsql;
        
        CREATE OR REPLACE FUNCTION days_series_between(start_date date, end_date date) RETURNS SETOF timestamptz AS $$
        DECLARE
        BEGIN
          RETURN QUERY SELECT * FROM generate_series(start_date, end_date, '1 day'::INTERVAL);
        END;
        $$ LANGUAGE plpgsql;
        '''
                 )


def downgrade():
    bind = op.get_bind()

    bind.execute('''DROP FUNCTION dates_between''')
    bind.execute('''DROP FUNCTION days_series_between''')