How to maintain your Prefect database

-- Total database size
SELECT pg_size_pretty(pg_database_size('prefect')) AS database_size;

-- Table sizes with row counts
SELECT 
    schemaname,
    relname AS tablename,
    pg_size_pretty(pg_total_relation_size(schemaname||'.'||relname)) AS total_size,
    to_char(n_live_tup, 'FM999,999,999') AS row_count
FROM pg_stat_user_tables
WHERE schemaname = 'public'
ORDER BY pg_total_relation_size(schemaname||'.'||relname) DESC
LIMIT 20;

-- Check database disk usage
SELECT 
    current_setting('data_directory') AS data_directory,
    pg_size_pretty(pg_database_size('prefect')) AS database_size,
    pg_size_pretty(pg_total_relation_size('public.events')) AS events_table_size,
    pg_size_pretty(pg_total_relation_size('public.log')) AS log_table_size;

-- Check available disk space (requires pg_stat_disk extension or shell access)
-- Run from shell: df -h /path/to/postgresql/data

SELECT
    schemaname,
    relname AS tablename,
    n_live_tup AS live_tuples,
    n_dead_tup AS dead_tuples,
    CASE WHEN n_live_tup > 0 
        THEN round(100.0 * n_dead_tup / n_live_tup, 2)
        ELSE 0
    END AS bloat_percent,
    last_vacuum,
    last_autovacuum
FROM pg_stat_user_tables
WHERE schemaname = 'public' 
    AND n_dead_tup > 1000
ORDER BY bloat_percent DESC;

-- Check index sizes and bloat
SELECT 
    schemaname,
    relname AS tablename,
    indexrelname AS indexname,
    pg_size_pretty(pg_relation_size(indexrelid)) AS index_size,
    idx_scan AS index_scans,
    idx_tup_read AS tuples_read,
    idx_tup_fetch AS tuples_fetched
FROM pg_stat_user_indexes
WHERE schemaname = 'public'
ORDER BY pg_relation_size(indexrelid) DESC
LIMIT 20;

-- Standard VACUUM (doesn't lock table)
VACUUM ANALYZE flow_run;
VACUUM ANALYZE task_run;
VACUUM ANALYZE log;

-- VACUUM FULL (rebuilds table, requires exclusive lock)
-- WARNING: This COMPLETELY LOCKS the table - no reads or writes!
-- Can take HOURS on large tables. Only use as last resort.
VACUUM FULL flow_run;

-- Better alternative: pg_repack (if installed)
-- Rebuilds tables online without blocking
-- pg_repack -t flow_run -d prefect

-- Show autovacuum settings
SHOW autovacuum;
SHOW autovacuum_vacuum_scale_factor;
SHOW autovacuum_vacuum_threshold;

-- Check when tables were last vacuumed
SELECT 
    schemaname,
    relname AS tablename,
    last_vacuum,
    last_autovacuum,
    vacuum_count,
    autovacuum_count
FROM pg_stat_user_tables
WHERE schemaname = 'public'
ORDER BY last_autovacuum NULLS FIRST;

-- For high-volume events table (INSERT/DELETE heavy)
ALTER TABLE events SET (
    autovacuum_vacuum_scale_factor = 0.05,  -- Default is 0.2
    autovacuum_vacuum_threshold = 1000,
    autovacuum_analyze_scale_factor = 0.02  -- Keep stats current
);

-- For state tables (INSERT-heavy)
ALTER TABLE flow_run_state SET (
    autovacuum_vacuum_scale_factor = 0.1,
    autovacuum_analyze_scale_factor = 0.05
);

-- For frequently updated tables
ALTER TABLE flow_run SET (
    autovacuum_vacuum_scale_factor = 0.1,
    autovacuum_vacuum_threshold = 500
);

# Enable both event and flow run vacuum
export PREFECT_SERVER_SERVICES_DB_VACUUM_ENABLED="events,flow_runs"

# Or use prefect config
prefect config set PREFECT_SERVER_SERVICES_DB_VACUUM_ENABLED="events,flow_runs"

export PREFECT_SERVER_SERVICES_DB_VACUUM_ENABLED="events,flow_runs"
export PREFECT_SERVER_SERVICES_DB_VACUUM_LOOP_SECONDS=604800
export PREFECT_SERVER_SERVICES_DB_VACUUM_RETENTION_PERIOD=2592000
export PREFECT_SERVER_SERVICES_DB_VACUUM_BATCH_SIZE=200

import asyncio
from datetime import datetime, timedelta, timezone
from prefect import flow, task, get_run_logger
from prefect.client.orchestration import get_client
from prefect.client.schemas.filters import FlowRunFilter, FlowRunFilterState, FlowRunFilterStateType, FlowRunFilterStartTime
from prefect.client.schemas.objects import StateType
from prefect.exceptions import ObjectNotFound

@task
async def delete_old_flow_runs(
    days_to_keep: int = 30,
    batch_size: int = 100
):
    """Delete completed flow runs older than specified days."""
    logger = get_run_logger()

    async with get_client() as client:
        cutoff = datetime.now(timezone.utc) - timedelta(days=days_to_keep)

        # Create filter for old completed flow runs
        # Note: Using start_time because created time filtering is not available
        flow_run_filter = FlowRunFilter(
            start_time=FlowRunFilterStartTime(before_=cutoff),
            state=FlowRunFilterState(
                type=FlowRunFilterStateType(
                    any_=[StateType.COMPLETED, StateType.FAILED, StateType.CANCELLED]
                )
            )
        )

        # Get flow runs to delete
        flow_runs = await client.read_flow_runs(
            flow_run_filter=flow_run_filter,
            limit=batch_size
        )

        deleted_total = 0

        while flow_runs:
            batch_deleted = 0
            failed_deletes = []

            # Delete each flow run through the API
            for flow_run in flow_runs:
                try:
                    await client.delete_flow_run(flow_run.id)
                    deleted_total += 1
                    batch_deleted += 1
                except ObjectNotFound:
                    # Already deleted (e.g., by concurrent cleanup) - treat as success
                    deleted_total += 1
                    batch_deleted += 1
                except Exception as e:
                    logger.warning(f"Failed to delete flow run {flow_run.id}: {e}")
                    failed_deletes.append(flow_run.id)

                # Rate limiting - adjust based on your API capacity
                if batch_deleted % 10 == 0:
                    await asyncio.sleep(0.5)

            logger.info(f"Deleted {batch_deleted}/{len(flow_runs)} flow runs (total: {deleted_total})")
            if failed_deletes:
                logger.warning(f"Failed to delete {len(failed_deletes)} flow runs")

            # Get next batch
            flow_runs = await client.read_flow_runs(
                flow_run_filter=flow_run_filter,
                limit=batch_size
            )

            # Delay between batches to avoid overwhelming the API
            await asyncio.sleep(1.0)

        logger.info(f"Retention complete. Total deleted: {deleted_total}")

@flow(name="database-retention")
async def retention_flow():
    """Run database retention tasks."""
    await delete_old_flow_runs(
        days_to_keep=30,
        batch_size=100
    )

# Direct SQL only deletes what's defined by database foreign keys
# Logs and artifacts may be orphaned without proper cleanup
async with asyncpg.connect(connection_url) as conn:
    await conn.execute("""
        DELETE FROM flow_run 
        WHERE created < $1 
        AND state_type IN ('COMPLETED', 'FAILED', 'CANCELLED')
    """, cutoff)

# Set retention to 2 days (as environment variable)
export PREFECT_EVENTS_RETENTION_PERIOD="2d"

# Or in your prefect configuration
prefect config set PREFECT_EVENTS_RETENTION_PERIOD="2d"

-- Event table size and row count
SELECT 
    pg_size_pretty(pg_total_relation_size('public.events')) AS total_size,
    to_char(count(*), 'FM999,999,999') AS row_count,
    min(occurred) AS oldest_event,
    max(occurred) AS newest_event
FROM events;

SELECT 
    count(*) AS total_connections,
    count(*) FILTER (WHERE state = 'active') AS active,
    count(*) FILTER (WHERE state = 'idle') AS idle,
    (SELECT setting::int FROM pg_settings WHERE name = 'max_connections') AS max_connections
FROM pg_stat_activity;