bookwiz.io / supabase / migrations / 040_fix_usage_period_migration.sql
040_fix_usage_period_migration.sql
Raw
-- Fix duplicate user_usage records when users upgrade to paid plans
-- This migration addresses the issue where upgrading creates a new user_usage record
-- instead of migrating the existing free tier usage

-- Create function to migrate usage from free tier period to subscription period
CREATE OR REPLACE FUNCTION migrate_usage_to_subscription_period(user_uuid UUID)
RETURNS VOID
SECURITY DEFINER
SET search_path = public, pg_catalog, pg_temp
AS $$
DECLARE
    subscription_record RECORD;
    current_time TIMESTAMPTZ := NOW();
    free_tier_period_start TIMESTAMPTZ;
    free_tier_period_end TIMESTAMPTZ;
    existing_usage_record user_usage;
    subscription_period_start TIMESTAMPTZ;
    subscription_period_end TIMESTAMPTZ;
BEGIN
    -- Get the user's active subscription (handle time casting issues safely)
    SELECT * INTO subscription_record
    FROM subscriptions s
    WHERE s.user_id = user_uuid 
    AND s.status IN ('active', 'trialing', 'past_due')
    ORDER BY s.created_at DESC
    LIMIT 1;
    
    -- Check if subscription is actually active by trying to compare dates safely
    IF subscription_record IS NOT NULL AND subscription_record.current_period_end IS NOT NULL THEN
        BEGIN
            -- Try to safely compare the timestamp
            IF subscription_record.current_period_end > current_time THEN
                -- Subscription is active, continue with migration
                NULL;
            ELSE
                -- Subscription is expired, treat as no subscription
                subscription_record := NULL;
            END IF;
        EXCEPTION 
            WHEN OTHERS THEN
                -- If timestamp comparison fails, treat as no subscription
                RAISE LOG 'Failed to compare subscription period_end for user %: %', user_uuid, SQLERRM;
                subscription_record := NULL;
        END;
    END IF;
    
    -- Only proceed if user has an active subscription
    IF subscription_record IS NULL THEN
        RETURN;
    END IF;
    
    -- Calculate what the free tier period would be
    free_tier_period_start := DATE_TRUNC('month', current_time);
    free_tier_period_end := DATE_TRUNC('month', current_time) + INTERVAL '1 month';
    
    -- Get subscription period
    subscription_period_start := subscription_record.current_period_start;
    subscription_period_end := subscription_record.current_period_end;
    
    -- Look for existing free tier usage record
    SELECT * INTO existing_usage_record
    FROM user_usage 
    WHERE user_id = user_uuid 
    AND period_start = free_tier_period_start 
    AND period_end = free_tier_period_end;
    
    -- If we found a free tier record and it's different from the subscription period
    IF existing_usage_record IS NOT NULL AND 
       (existing_usage_record.period_start != subscription_period_start OR 
        existing_usage_record.period_end != subscription_period_end) THEN
        
        -- Check if subscription period record already exists
        IF EXISTS (
            SELECT 1 FROM user_usage 
            WHERE user_id = user_uuid 
            AND period_start = subscription_period_start 
            AND period_end = subscription_period_end
        ) THEN
            -- If subscription record exists, merge the usage and delete the free tier record
            UPDATE user_usage 
            SET 
                advanced_prompts_used = advanced_prompts_used + existing_usage_record.advanced_prompts_used,
                auto_prompts_used = auto_prompts_used + existing_usage_record.auto_prompts_used,
                storage_used_gb = storage_used_gb + existing_usage_record.storage_used_gb,
                custom_templates_created = custom_templates_created + existing_usage_record.custom_templates_created,
                books_created = books_created + existing_usage_record.books_created,
                updated_at = NOW()
            WHERE user_id = user_uuid 
            AND period_start = subscription_period_start 
            AND period_end = subscription_period_end;
            
            -- Delete the free tier record
            DELETE FROM user_usage WHERE id = existing_usage_record.id;
        ELSE
            -- Update the existing free tier record to use subscription periods
            UPDATE user_usage 
            SET 
                period_start = subscription_period_start,
                period_end = subscription_period_end,
                updated_at = NOW()
            WHERE id = existing_usage_record.id;
        END IF;
        
        -- Log the migration
        INSERT INTO usage_logs (user_id, feature_type, action, amount, metadata)
        VALUES (user_uuid, 'advanced_prompts', 'migrate', 0, 
                jsonb_build_object(
                    'migration_type', 'free_to_subscription_period',
                    'old_period_start', free_tier_period_start,
                    'old_period_end', free_tier_period_end,
                    'new_period_start', subscription_period_start,
                    'new_period_end', subscription_period_end
                ));
    END IF;
END;
$$ LANGUAGE plpgsql;

-- Update the get_or_create_current_usage function to handle migration
CREATE OR REPLACE FUNCTION get_or_create_current_usage(user_uuid UUID)
RETURNS user_usage
SECURITY DEFINER
SET search_path = public, pg_catalog, pg_temp
AS $$
DECLARE
    usage_record user_usage;
    period_info RECORD;
BEGIN
    -- First, attempt to migrate any existing free tier usage to subscription period
    PERFORM migrate_usage_to_subscription_period(user_uuid);
    
    -- Get current period (after any migration)
    SELECT * INTO period_info FROM get_current_usage_period(user_uuid);
    
    -- Try to get existing usage record
    SELECT * INTO usage_record 
    FROM user_usage 
    WHERE user_id = user_uuid 
    AND period_start = period_info.period_start 
    AND period_end = period_info.period_end;
    
    -- Create if doesn't exist
    IF usage_record IS NULL THEN
        INSERT INTO user_usage (
            user_id, 
            period_start, 
            period_end
        ) VALUES (
            user_uuid,
            period_info.period_start,
            period_info.period_end
        )
        RETURNING * INTO usage_record;
    END IF;
    
    RETURN usage_record;
END;
$$ LANGUAGE plpgsql;

-- Run a one-time cleanup for any existing duplicate records
-- This handles cases where users have already upgraded and have duplicate records
DO $$
DECLARE
    user_record RECORD;
    subscription_record RECORD;
    current_time TIMESTAMPTZ := NOW();
BEGIN
    -- Find all users with subscriptions (we'll check if they're active in the migration function)
    FOR user_record IN 
        SELECT DISTINCT user_id 
        FROM subscriptions 
        WHERE status IN ('active', 'trialing', 'past_due')
    LOOP
        -- Run migration for each user (the function will handle checking if subscription is active)
        BEGIN
            PERFORM migrate_usage_to_subscription_period(user_record.user_id);
        EXCEPTION 
            WHEN OTHERS THEN
                -- Log any errors but continue with other users
                RAISE LOG 'Failed to migrate usage for user %: %', user_record.user_id, SQLERRM;
        END;
    END LOOP;
END $$;

-- Add helpful comments
COMMENT ON FUNCTION migrate_usage_to_subscription_period(UUID) IS 'Migrates existing free tier usage to subscription billing period when user upgrades';
COMMENT ON FUNCTION get_or_create_current_usage(UUID) IS 'Gets existing usage record for current period or creates new one, with automatic migration for subscription upgrades';