-- ============================================================
-- Lakeflow DDL Audit Setup Script v1.0
-- ============================================================
--
--  This script sets up the PostgreSQL DDL audit system to track
--  schema changes in real-time. It captures DDL operations and stores them
--  in an audit table for replication by the ingestion pipelines.
--
-- DDL Objects created per database:
--   1. Audit Table: public.lakeflow_ddl_audit_table_<major>_<minor>
--      - Stores DDL events with metadata (user, schema, table, SQL text, WAL position)
--
--   2. Event Trigger Functions:
--      - public.lakeflow_ddl_audit_function_<major>_<minor>: Captures ALTER TABLE events
--      - public.lakeflow_drop_ddl_audit_function_<major>_<minor>: Captures DROP TABLE events
--
--   3. Event Triggers:
--      - lakeflow_ddl_audit_trigger_<major>_<minor>: Fires on ddl_command_end
--      - lakeflow_drop_ddl_audit_trigger_<major>_<minor>: Fires on sql_drop
--
-- Note:
--   1. All these objects are created in the 'public' schema.
--   2. With custom provided slot/pub mode ingestion:
--         - Customer needs to ensure lakeflow_ddl_audit_table_<major>_<minor> is present
--           in the publication for each catalog
--
-- VERSIONING:
--   The script uses semantic versioning (major.minor). Old versions are automatically
--   cleaned up when a new version script is executed.
--
-- PREREQUISITES:
--   User Privileges Required:
--     - Superuser privileges to execute this script, specifically to:
--       * Create event triggers (requires superuser)
--       * Create functions with SECURITY DEFINER
--       * Create tables in the public schema
--       * Grant/revoke privileges on functions
--     - Note: Regular users do NOT need any special privileges to use the audit system
--       once installed (DDL auditing happens automatically via SECURITY DEFINER)
--
-- USAGE:
--  1. Execute this script as a superuser for each database from which you want to ingest tables:
--     psql -h <host> -U <admin/super user> -d <database> -f lakeflow_pg_ddl_change_tracking.sql
--
--  2. Verify setup:
--     SELECT * FROM pg_event_trigger WHERE evtname LIKE 'lakeflow%';
--     SELECT * FROM pg_tables WHERE tablename LIKE 'lakeflow_ddl_audit_table%';
--
--  3. Ensure the audit table is added to the publication, when providing custom
--     slots and publications.
--
-- MAINTENANCE:
--   - The script is idempotent; it can be safely re-run
--   - Old versions are automatically cleaned up (triggers, functions, tables)
--   - Monitor audit table size and implement retention policies as needed
--
-- SECURITY:
--   - Audit functions use SECURITY DEFINER to execute with superuser privileges
--   - This allows any user to trigger DDL auditing without explicit INSERT permissions
--   - Event triggers require superuser to create but run for all users
--
-- ============================================================

-- ============================================================
-- 0. SET VERSION AND REPLICATION USER
-- ============================================================
-- Set the current version of the DDL audit system
-- This allows multiple versions to coexist during rolling upgrades
SET lakeflow.current_major = '1';
SET lakeflow.current_minor = '0';

-- ============================================================
-- 1. DROP OLD OBJECTS (Version Cleanup)
-- ============================================================
-- This section removes event triggers, functions, and audit tables from
-- previous versions. This ensures a clean upgrade path and prevents
-- conflicts when deploying a new version.
--
-- Steps:
--   1. Identify the current version suffix (currently: "1_0")
--   2. Find all lakeflow objects that don't match the current version
--   3. Drop them in dependency order: triggers → functions → tables
-- ============================================================
DO $$
DECLARE
    major_version INT := current_setting('lakeflow.current_major')::INT;
    minor_version INT := current_setting('lakeflow.current_minor')::INT;
    current_version_suffix TEXT := major_version || '_' || minor_version;

    obj RECORD;
BEGIN
    -- Drop old event triggers first (they depend on functions)
    FOR obj IN
        SELECT evtname
        FROM pg_event_trigger
        WHERE (evtname LIKE 'lakeflow_ddl_audit_trigger_%'
               OR evtname LIKE 'lakeflow_drop_ddl_audit_trigger_%')
          AND evtname NOT LIKE '%' || current_version_suffix
    LOOP
        RAISE NOTICE 'Dropping old trigger: %', obj.evtname;
        EXECUTE format('DROP EVENT TRIGGER IF EXISTS %I;', obj.evtname);
    END LOOP;

    -- Drop old functions (they depend on tables)
    FOR obj IN
        SELECT proname
        FROM pg_proc
        WHERE (proname LIKE 'lakeflow_ddl_audit_function_%'
               OR proname LIKE 'lakeflow_drop_ddl_audit_function_%')
          AND proname NOT LIKE '%' || current_version_suffix
    LOOP
        RAISE NOTICE 'Dropping old function: %', obj.proname;
        EXECUTE format('DROP FUNCTION IF EXISTS public.%I();', obj.proname);
    END LOOP;

    -- Drop old audit tables (should be empty after function cleanup)
    FOR obj IN
        SELECT tablename
        FROM pg_tables
        WHERE schemaname = 'public'
          AND tablename LIKE 'lakeflow_ddl_audit_table_%'
          AND tablename NOT LIKE 'lakeflow_ddl_audit_table_' || current_version_suffix
    LOOP
        RAISE NOTICE 'Dropping old audit table: %', obj.tablename;
        EXECUTE format('DROP TABLE IF EXISTS public.%I;', obj.tablename);
    END LOOP;
END $$;

-- ============================================================
-- 2. CREATE AUDIT TABLE
-- ============================================================
-- Creates the central audit table that stores all captured DDL events.
--
-- Notes:
--   - The table is created in the 'public' schema.
--   - Column names are quoted to preserve case sensitivity
--   - Table is idempotent - won't recreate if it already exists
--   - VARCHAR sizes are chosen to accommodate typical PostgreSQL identifiers
-- ============================================================
DO $$
DECLARE
    major_version INT := current_setting('lakeflow.current_major')::INT;
    minor_version INT := current_setting('lakeflow.current_minor')::INT;
    current_version_suffix TEXT := major_version || '_' || minor_version;

    audit_table_name TEXT := 'lakeflow_ddl_audit_table_' || current_version_suffix;
BEGIN
  IF NOT EXISTS (
        SELECT 1 FROM pg_tables
        WHERE schemaname = 'public'
          AND tablename = audit_table_name
    ) THEN
        RAISE NOTICE 'Creating audit table: %', audit_table_name;
        EXECUTE format($sql$
            CREATE TABLE public.%I (
                "ID"               BIGSERIAL PRIMARY KEY,
                "TIMESTAMP"        TIMESTAMPTZ DEFAULT now(),
                "CURRENT_USER"     VARCHAR(128),  -- User who executed the DDL
                "SCHEMA_NAME"      VARCHAR(128),  -- Schema of the affected table
                "OBJECT_NAME"      VARCHAR(128),  -- Name of the affected table
                "OBJECT_TYPE"      VARCHAR(30),   -- Object type (table, view, etc.)
                "OID"              OID,           -- Object identifier
                "OPERATION_TYPE"   VARCHAR(30),   -- DDL operation (ALTER TABLE, DROP TABLE, etc.)
                "SQL_TXT"          TEXT,          -- Full DDL SQL statement
                "LOGICAL_POSITION" VARCHAR(20)    -- WAL LSN position for ordering
            )
        $sql$, audit_table_name);
    ELSE
        RAISE NOTICE 'Audit table % already exists, skipping creation.', audit_table_name;
    END IF;
END $$;

-- ============================================================
-- 3. CREATE FUNCTIONS + TRIGGERS
-- ============================================================
-- This section creates event trigger functions and their associated triggers
-- to capture DDL operations in real-time.
--
-- Overview:
--   1. ALTER TABLE Handler:
--      - Function: lakeflow_ddl_audit_function_<version>
--      - Trigger: lakeflow_ddl_audit_trigger_<version>
--      - Fires on: ddl_command_end event
--      - Captures: All ALTER TABLE operations
--
--   2. DROP TABLE Handler:
--      - Function: lakeflow_drop_ddl_audit_function_<version>
--      - Trigger: lakeflow_drop_ddl_audit_trigger_<version>
--      - Fires on: sql_drop event
--      - Captures: All DROP TABLE operations
-- ============================================================
DO $$
DECLARE
    major_version INT := current_setting('lakeflow.current_major')::INT;
    minor_version INT := current_setting('lakeflow.current_minor')::INT;
    current_version_suffix TEXT := major_version || '_' || minor_version;

    audit_table_name TEXT := 'lakeflow_ddl_audit_table_' || current_version_suffix;
    alter_function_name TEXT := 'lakeflow_ddl_audit_function_' || current_version_suffix;
    alter_trigger_name TEXT := 'lakeflow_ddl_audit_trigger_' || current_version_suffix;
    drop_function_name TEXT := 'lakeflow_drop_ddl_audit_function_' || current_version_suffix;
    drop_trigger_name TEXT := 'lakeflow_drop_ddl_audit_trigger_' || current_version_suffix;
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_proc WHERE proname = alter_function_name) THEN
    RAISE NOTICE 'Creating function: %', alter_function_name;
    EXECUTE format($sql$
        CREATE FUNCTION public.%I()
        RETURNS event_trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = pg_catalog
        AS $f$
        DECLARE
            command_record  RECORD;
            ddl_sql         TEXT := current_query();
        BEGIN
            FOR command_record IN
                SELECT * FROM pg_event_trigger_ddl_commands()
                WHERE object_type IN ('table', 'table column')
            LOOP
                BEGIN
                    INSERT INTO public.%I (
                        "CURRENT_USER", "SCHEMA_NAME", "OBJECT_NAME", "OBJECT_TYPE", "OID",
                        "OPERATION_TYPE", "SQL_TXT", "LOGICAL_POSITION"
                    )
                    VALUES (
                        CURRENT_USER,
                        command_record.schema_name,
                        command_record.object_identity,
                        command_record.object_type,
                        command_record.objid,
                        tg_tag,
                        ddl_sql,
                        pg_current_wal_lsn()::text
                    );
                EXCEPTION
                    WHEN OTHERS THEN
                        -- Log the error but don't fail the DDL operation
                        RAISE WARNING 'Failed to audit DDL operation: %%', SQLERRM;
                END;
            END LOOP;
        END;
        $f$
    $sql$, alter_function_name, audit_table_name);
        -- Revoke public execute privilege for security
        EXECUTE format('REVOKE ALL ON FUNCTION public.%I() FROM PUBLIC;', alter_function_name);
    ELSE
        RAISE NOTICE 'Function % already exists, skipping.', alter_function_name;
    END IF;

    ----------------------------------------------------------------
    -- Trigger for ALTER TABLE
    ----------------------------------------------------------------
    IF NOT EXISTS (SELECT 1 FROM pg_event_trigger WHERE evtname = alter_trigger_name) THEN
    RAISE NOTICE 'Creating trigger: %', alter_trigger_name;
    EXECUTE format(
        'CREATE EVENT TRIGGER %I ON ddl_command_end
         WHEN TAG IN (''ALTER TABLE'')
         EXECUTE FUNCTION public.%I();',
        alter_trigger_name, alter_function_name
    );
    ELSE
        RAISE NOTICE 'Trigger % already exists, skipping.', alter_trigger_name;
    END IF;

    ----------------------------------------------------------------
    -- Function for DROP TABLE
    ----------------------------------------------------------------
    IF NOT EXISTS (SELECT 1 FROM pg_proc WHERE proname = drop_function_name) THEN
        RAISE NOTICE 'Creating function: %', drop_function_name;
        EXECUTE format($sql$
            CREATE FUNCTION public.%I()
            RETURNS event_trigger
            LANGUAGE plpgsql
            SECURITY DEFINER
            SET search_path = pg_catalog
            AS $f$
            DECLARE
                dropped_object RECORD;
            BEGIN
                FOR dropped_object IN
                    SELECT * FROM pg_event_trigger_dropped_objects()
                    WHERE object_type IN ('table')
                LOOP
                    BEGIN
                        INSERT INTO public.%I (
                            "CURRENT_USER", "SCHEMA_NAME", "OBJECT_NAME", "OBJECT_TYPE",
                            "OID", "OPERATION_TYPE", "SQL_TXT", "LOGICAL_POSITION"
                        )
                        VALUES (
                            CURRENT_USER,
                            dropped_object.schema_name,
                            dropped_object.object_name,
                            dropped_object.object_type,
                            dropped_object.objid,
                            tg_tag,
                            current_query(),
                            pg_current_wal_lsn()::text
                        );
                    EXCEPTION
                        WHEN OTHERS THEN
                            -- Log the error but don't fail the DDL operation
                            RAISE WARNING 'Failed to audit DROP TABLE operation: %%', SQLERRM;
                    END;
                END LOOP;
            END;
            $f$
        $sql$, drop_function_name, audit_table_name);
        -- Revoke public execute privilege for security
        EXECUTE format('REVOKE ALL ON FUNCTION public.%I() FROM PUBLIC;', drop_function_name);
    ELSE
        RAISE NOTICE 'Function % already exists, skipping.', drop_function_name;
    END IF;

    ----------------------------------------------------------------
    -- Trigger for DROP TABLE
    ----------------------------------------------------------------
    IF NOT EXISTS (SELECT 1 FROM pg_event_trigger WHERE evtname = drop_trigger_name) THEN
        RAISE NOTICE 'Creating trigger: %', drop_trigger_name;
        EXECUTE format(
            'CREATE EVENT TRIGGER %I ON sql_drop
             WHEN TAG IN (''DROP TABLE'')
             EXECUTE FUNCTION public.%I();',
            drop_trigger_name, drop_function_name
        );
    ELSE
        RAISE NOTICE 'Trigger % already exists, skipping.', drop_trigger_name;
    END IF;

    ----------------------------------------------------------------
    -- Done
    ----------------------------------------------------------------
    RAISE NOTICE 'DDL audit objects successfully created for version %', current_version_suffix;
END $$;

-- ============================================================
-- SETUP COMPLETE
-- ============================================================
-- For support or questions, refer to the Lakeflow documentation.
-- ============================================================

