59d3ad8f0527_.py

"""Add churn() function

Revision ID: 59d3ad8f0527
Revises: dc0d42d537a6
Create Date: 2023-10-06 16:10:49.364836

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = "59d3ad8f0527"
down_revision = "dc0d42d537a6"
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()
    conn.execute(
        """
        BEGIN;
        CREATE OR REPLACE FUNCTION churn(vyear integer, vmonth integer DEFAULT NULL) 
        RETURNS TABLE (
            registration_year integer,
            registration_month integer,
            registered_companies bigint,
            first_employee_invitation_date bigint,
            first_mission_validation_by_admin_date bigint
        )
        LANGUAGE plpgsql
        AS $$
        BEGIN
            -- Sanity check   
            IF (vyear < 2000 OR vyear > 9999) THEN
                RAISE EXCEPTION 'The year must be in YYYY format (between 2000 and 9999)';
            ELSIF vmonth IS NOT NULL AND (vmonth < 1 OR vmonth > 12) THEN
                RAISE EXCEPTION 'The month must be in MM format (between 01 and 12)';
            END IF;

            RETURN QUERY (
                SELECT
                    EXTRACT(YEAR FROM cs.company_creation_date)::integer AS registration_year,
                    EXTRACT(MONTH FROM cs.company_creation_date)::integer AS registration_month,
                    COUNT(DISTINCT cs.company_id)::bigint AS registered_companies,
                    COUNT(DISTINCT CASE WHEN EXTRACT(MONTH FROM cs.first_employee_invitation_date) = EXTRACT(MONTH FROM cs.company_creation_date) THEN cs.company_id ELSE NULL END)::bigint AS first_employee_invitation_date,
                    COUNT(DISTINCT CASE WHEN EXTRACT(MONTH FROM cs.first_mission_validation_by_admin_date) = EXTRACT(MONTH FROM cs.company_creation_date) THEN cs.company_id ELSE NULL END)::bigint AS first_mission_validation_by_admin_date
                FROM
                    company_stats AS cs
                WHERE
                    EXTRACT(YEAR FROM cs.company_creation_date) = vyear
                    AND (vmonth IS NULL OR EXTRACT(MONTH FROM cs.company_creation_date) = vmonth)
                GROUP BY
                    registration_year, registration_month
                ORDER BY
                    registration_year, registration_month
            );
        END;
        $$;
        COMMIT;
    """
    )

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.execute("DROP FUNCTION churn(integer, integer)")
    # ### end Alembic commands ###