-- Copyright (c) Microsoft Corporation.
-- Licensed under the PostgreSQL License.

-- pg_durable upgrade: 0.1.1 → 0.2.0
--
-- Run with: ALTER EXTENSION pg_durable UPDATE TO '0.2.0';
--
-- Each schema-changing PR should add its DDL here.
-- See docs/upgrade-testing.md for the upgrade testing plan.

-- Changes:
--   - df.vars: Add per-user scoping via `owner` column + RLS
--     (Implements rls.md Decision 5, Option A)
--
-- Run with: ALTER EXTENSION pg_durable UPDATE TO '0.2.0';

-- ============================================================================
-- 1. Migrate df.vars schema: add owner column, change PK
-- ============================================================================

-- Add the owner column with a default. Existing rows get the current user
-- (the superuser running ALTER EXTENSION). Since vars are ephemeral
-- (set before df.start(), captured at start time), stale rows in this table
-- are unlikely to matter. If they do, admins should reassign ownership
-- manually before upgrading.
ALTER TABLE df.vars ADD COLUMN owner REGROLE NOT NULL DEFAULT current_user::regrole;

-- Change PK from (name) to (owner, name)
ALTER TABLE df.vars DROP CONSTRAINT vars_pkey;
ALTER TABLE df.vars ADD PRIMARY KEY (owner, name);

-- Keep direct INSERT on df.instances/df.nodes, but limit it to the columns
-- df.start() writes. Runtime-owned columns remain protected from direct INSERT.
REVOKE INSERT ON df.instances FROM PUBLIC;
REVOKE INSERT ON df.nodes FROM PUBLIC;
GRANT INSERT (id, label, root_node, submitted_by, database) ON df.instances TO PUBLIC;
GRANT INSERT (id, instance_id, node_type, query, result_name, left_node, right_node, submitted_by, database) ON df.nodes TO PUBLIC;

-- Enforce a df.start-shaped graph for new direct writes without blocking
-- upgrades on legacy malformed rows that may already exist.
ALTER TABLE df.instances
    ADD CONSTRAINT instances_id_format_chk
        CHECK (id ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT instances_root_node_format_chk
        CHECK (root_node ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT instances_status_chk
        CHECK (status IN ('pending', 'running', 'completed', 'failed', 'cancelled')) NOT VALID,
    -- Supports the composite FK from df.nodes that ties node identity to the instance row.
    ADD CONSTRAINT instances_identity_key
        UNIQUE (id, submitted_by);

ALTER TABLE df.nodes
    ADD CONSTRAINT nodes_instance_id_present_chk
        CHECK (instance_id IS NOT NULL) NOT VALID,
    ADD CONSTRAINT nodes_submitted_by_present_chk
        CHECK (submitted_by IS NOT NULL) NOT VALID,
    ADD CONSTRAINT nodes_id_format_chk
        CHECK (id ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT nodes_instance_id_format_chk
        CHECK (instance_id ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT nodes_left_node_format_chk
        CHECK (left_node IS NULL OR left_node ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT nodes_right_node_format_chk
        CHECK (right_node IS NULL OR right_node ~ '^[0-9a-f]{8}$') NOT VALID,
    ADD CONSTRAINT nodes_node_type_chk
        CHECK (node_type IN ('SQL', 'THEN', 'IF', 'JOIN', 'LOOP', 'BREAK', 'RACE', 'SLEEP', 'WAIT_SCHEDULE', 'HTTP', 'SIGNAL')) NOT VALID,
    ADD CONSTRAINT nodes_result_name_chk
        CHECK (result_name IS NULL OR result_name ~ '^[A-Za-z_][A-Za-z0-9_]*$') NOT VALID,
    ADD CONSTRAINT nodes_status_chk
        CHECK (status IN ('pending', 'running', 'completed', 'failed')) NOT VALID,
    ADD CONSTRAINT nodes_result_status_chk
        CHECK (result IS NULL OR status IN ('completed', 'failed')) NOT VALID,
    ADD CONSTRAINT nodes_structure_chk
        CHECK (
            CASE
                WHEN node_type IN ('SQL', 'SLEEP', 'WAIT_SCHEDULE', 'BREAK', 'HTTP', 'SIGNAL')
                    THEN left_node IS NULL AND right_node IS NULL AND query IS NOT NULL
                WHEN node_type = 'THEN'
                    THEN left_node IS NOT NULL AND right_node IS NOT NULL AND query IS NULL
                WHEN node_type = 'IF'
                    THEN left_node IS NOT NULL AND right_node IS NOT NULL AND query IS NOT NULL
                WHEN node_type = 'LOOP'
                    THEN left_node IS NOT NULL AND right_node IS NULL
                WHEN node_type = 'JOIN'
                    THEN left_node IS NOT NULL AND right_node IS NOT NULL
                WHEN node_type = 'RACE'
                    THEN left_node IS NOT NULL AND right_node IS NOT NULL AND query IS NULL
                ELSE FALSE
            END
        ) NOT VALID,
    ADD CONSTRAINT nodes_instance_node_key
        UNIQUE (instance_id, id);

ALTER TABLE df.nodes
    ADD CONSTRAINT nodes_instance_identity_fkey
        FOREIGN KEY (instance_id, submitted_by)
        REFERENCES df.instances (id, submitted_by)
        DEFERRABLE INITIALLY DEFERRED NOT VALID,
    ADD CONSTRAINT nodes_left_node_same_instance_fkey
        FOREIGN KEY (instance_id, left_node)
        REFERENCES df.nodes (instance_id, id)
        DEFERRABLE INITIALLY DEFERRED NOT VALID,
    ADD CONSTRAINT nodes_right_node_same_instance_fkey
        FOREIGN KEY (instance_id, right_node)
        REFERENCES df.nodes (instance_id, id)
        DEFERRABLE INITIALLY DEFERRED NOT VALID;

ALTER TABLE df.instances
    ADD CONSTRAINT instances_root_node_same_instance_fkey
        FOREIGN KEY (id, root_node)
        REFERENCES df.nodes (instance_id, id)
        DEFERRABLE INITIALLY DEFERRED NOT VALID;

-- ============================================================================
-- 2. Enable RLS on df.vars
-- ============================================================================

ALTER TABLE df.vars ENABLE ROW LEVEL SECURITY;

CREATE POLICY vars_user_isolation ON df.vars
    FOR ALL
    USING (owner = current_user::regrole)
    WITH CHECK (owner = current_user::regrole);

DROP POLICY instances_user_isolation ON df.instances;
CREATE POLICY instances_user_isolation ON df.instances
    FOR ALL
    USING (submitted_by = current_user::regrole)
    WITH CHECK (submitted_by = current_user::regrole);

DROP POLICY nodes_user_isolation ON df.nodes;
CREATE POLICY nodes_user_isolation ON df.nodes
    FOR ALL
    USING (submitted_by = current_user::regrole)
    WITH CHECK (submitted_by = current_user::regrole);

-- ============================================================================
-- 3. Harden PL/pgSQL and SQL helper functions with SET search_path
--    (Defense-in-depth: all calls are already schema-qualified, but this
--    prevents future edits from accidentally introducing unqualified refs.)
-- ============================================================================

CREATE OR REPLACE FUNCTION df.as_op(fut text, name text) RETURNS text AS $$
    SELECT df.as(fut, name);
$$ LANGUAGE SQL IMMUTABLE SET search_path = pg_catalog, df, pg_temp;

CREATE OR REPLACE FUNCTION df.if_then_op(condition text, then_branch text) RETURNS text AS $$
DECLARE
    cond_fut jsonb;
    then_fut jsonb;
    result_obj jsonb;
BEGIN
    cond_fut := df.ensure_durofut(condition)::jsonb;
    then_fut := df.ensure_durofut(then_branch)::jsonb;
    result_obj := jsonb_build_object(
        '_partial_if', true,
        'condition', cond_fut,
        'then_branch', then_fut
    );
    RETURN result_obj::text;
END;
$$ LANGUAGE plpgsql IMMUTABLE SET search_path = pg_catalog, df, pg_temp;

CREATE OR REPLACE FUNCTION df.if_else_op(partial_if text, else_branch text) RETURNS text AS $$
DECLARE
    partial jsonb;
    else_fut text;
    cond_text text;
    then_text text;
BEGIN
    partial := partial_if::jsonb;
    IF partial->>'_partial_if' IS NULL THEN
        RAISE EXCEPTION 'Invalid if-then-else: left side of !> must be a ?> expression';
    END IF;
    cond_text := partial->'condition'::text;
    then_text := partial->'then_branch'::text;
    else_fut := df.ensure_durofut(else_branch);
    RETURN df.if(cond_text, then_text, else_fut);
END;
$$ LANGUAGE plpgsql IMMUTABLE SET search_path = pg_catalog, df, pg_temp;

CREATE OR REPLACE FUNCTION df.ensure_durofut(val text) RETURNS text AS $$
DECLARE
    node_type_val text;
BEGIN
    BEGIN
        node_type_val := (val::jsonb)->>'node_type';
        IF node_type_val IS NOT NULL THEN
            IF node_type_val NOT IN ('SQL', 'THEN', 'IF', 'JOIN', 'LOOP', 'BREAK', 'RACE', 'SLEEP', 'WAIT_SCHEDULE', 'HTTP', 'SIGNAL') THEN
                RAISE EXCEPTION 'Unknown node_type ''%''. Valid types: SQL, THEN, IF, JOIN, LOOP, BREAK, RACE, SLEEP, WAIT_SCHEDULE, HTTP, SIGNAL', node_type_val;
            END IF;
            RETURN val;
        END IF;
    EXCEPTION WHEN invalid_text_representation THEN
        NULL;
    WHEN raise_exception THEN
        RAISE;
    WHEN OTHERS THEN
        NULL;
    END;
    RETURN df.sql(val);
END;
$$ LANGUAGE plpgsql IMMUTABLE SET search_path = pg_catalog, df, pg_temp;

CREATE OR REPLACE FUNCTION df.loop_prefix_op(body text) RETURNS text AS $$
    SELECT df.loop(body);
$$ LANGUAGE SQL IMMUTABLE SET search_path = pg_catalog, df, pg_temp;

-- ============================================================================
-- 4. Add df.if_rows() — branches on whether a named result has rows
-- ============================================================================

CREATE FUNCTION df."if_rows"(
    "result_name" TEXT,
    "then_branch" TEXT,
    "else_branch" TEXT
) RETURNS TEXT
STRICT
LANGUAGE c
AS 'MODULE_PATHNAME', 'if_rows_fn_wrapper';

-- ============================================================================
-- 5. Drop login_role column (user-isolation simplification)
--    See docs/user-isolation.md for design.
--    login_role existed in v0.1.1 on both df.instances (NOT NULL) and df.nodes
--    (nullable). The new model uses only submitted_by.
--
--    We also need to update the column-level INSERT grants since they were
--    set above with the old column list. The REVOKE/GRANT will be re-issued
--    below to reflect the final column set.
-- ============================================================================

-- Drop the login_role columns. Existing data (if any) in these columns is
-- discarded — the new .so no longer reads or writes login_role.
ALTER TABLE df.instances DROP COLUMN IF EXISTS login_role;
ALTER TABLE df.nodes DROP COLUMN IF EXISTS login_role;

-- ============================================================================
-- 6. Add df.grant_usage() and df.revoke_usage() helpers for role privilege management
--
-- df.grant_usage(role, include_http => false, with_grant => false) grants all
-- standard df privileges but NOT df.http() by default.
-- Pass include_http => true to opt in to HTTP.
-- Pass with_grant => true to allow the role to delegate access to others.
-- Authorization is enforced by PostgreSQL's native WITH GRANT OPTION
-- mechanism (see src/lib.rs for details).
--
-- Upgrade note: step 7 below documents the PUBLIC EXECUTE grant on
-- df.http() that v0.1.1 issued via 'GRANT EXECUTE ON ALL FUNCTIONS'.
-- The upgrade intentionally does NOT revoke it (to avoid breaking running
-- workflows).  If any roles should lose HTTP access, run:
--   REVOKE EXECUTE ON FUNCTION df.http(text,text,text,jsonb,integer) FROM <role>;
-- ============================================================================

CREATE OR REPLACE FUNCTION df.grant_usage(
    p_role TEXT,
    include_http boolean DEFAULT false,
    with_grant boolean DEFAULT false
)
RETURNS VOID
LANGUAGE plpgsql
SET search_path = pg_catalog, df, pg_temp
AS $fn$
DECLARE
    grant_opt TEXT := '';
    func_sig TEXT;
    -- Explicit list of df.* functions to grant.  Sensitive functions
    -- (df.http, df.grant_usage, df.revoke_usage) are excluded from this
    -- list and granted conditionally below.
    func_sigs TEXT[] := ARRAY[
        -- DSL functions
        'df.sql(text)',
        'df.seq(text, text)',
        'df.as(text, text)',
        'df.sleep(bigint)',
        'df.wait_for_schedule(text)',
        'df.loop(text, text)',
        'df.break(text)',
        'df.if(text, text, text)',
        'df.if_rows(text, text, text)',
        'df.join(text, text)',
        'df.join3(text, text, text)',
        'df.race(text, text)',
        'df.wait_for_signal(text, integer)',
        'df.signal(text, text, text)',
        'df.start(text, text, text)',
        'df.setvar(text, text)',
        'df.getvar(text)',
        'df.unsetvar(text)',
        'df.clearvars()',
        -- Monitoring functions
        'df.status(text)',
        'df.result(text)',
        'df.cancel(text, text)',
        'df.wait_for_completion(text, integer)',
        'df.run(text)',
        'df.list_instances(text, integer)',
        'df.instance_info(text)',
        'df.instance_nodes(text, integer)',
        'df.instance_executions(text, integer)',
        'df.metrics()',
        -- Internal helpers (operators, version, etc.)
        'df.as_op(text, text)',
        'df.if_then_op(text, text)',
        'df.if_else_op(text, text)',
        'df.ensure_durofut(text)',
        'df.loop_prefix_op(text)',
        'df.version()',
        'df.debug_connection()',
        'df.explain(text)',
        'df.target_database()'
    ];
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = p_role) THEN
        RAISE EXCEPTION 'role "%" does not exist', p_role;
    END IF;

    IF with_grant THEN
        grant_opt := ' WITH GRANT OPTION';
    END IF;

    -- Schema access
    EXECUTE format('GRANT USAGE ON SCHEMA df TO %I', p_role) || grant_opt;

    -- Grant EXECUTE on each standard function explicitly.
    FOREACH func_sig IN ARRAY func_sigs LOOP
        EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO %I', func_sig, p_role) || grant_opt;
    END LOOP;

    -- df.http() — opt-in because it makes outbound network requests.
    IF include_http THEN
        EXECUTE format('GRANT EXECUTE ON FUNCTION df.http(text, text, text, jsonb, integer) TO %I', p_role) || grant_opt;
    END IF;

    -- Admin helpers — only for delegated administrators.
    IF with_grant THEN
        EXECUTE format('GRANT EXECUTE ON FUNCTION df.grant_usage(text, boolean, boolean) TO %I', p_role) || grant_opt;
        EXECUTE format('GRANT EXECUTE ON FUNCTION df.revoke_usage(text) TO %I', p_role) || grant_opt;
    END IF;

    -- Table privileges
    EXECUTE format('GRANT SELECT ON df.instances TO %I', p_role) || grant_opt;
    EXECUTE format('GRANT UPDATE (status, updated_at) ON df.instances TO %I', p_role) || grant_opt;
    EXECUTE format('GRANT SELECT ON df.nodes TO %I', p_role) || grant_opt;
    EXECUTE format('GRANT INSERT (id, label, root_node, submitted_by, database) ON df.instances TO %I', p_role) || grant_opt;
    EXECUTE format('GRANT INSERT (id, instance_id, node_type, query, result_name, left_node, right_node, submitted_by, database) ON df.nodes TO %I', p_role) || grant_opt;
    EXECUTE format('GRANT SELECT, INSERT, UPDATE, DELETE ON df.vars TO %I', p_role) || grant_opt;

    RAISE NOTICE 'pg_durable: granted df usage privileges to "%"', p_role;
END;
$fn$;

CREATE OR REPLACE FUNCTION df.revoke_usage(p_role TEXT)
RETURNS VOID
LANGUAGE plpgsql
SET search_path = pg_catalog, df, pg_temp
AS $fn$
DECLARE
    func_oid oid;
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = p_role) THEN
        RAISE EXCEPTION 'role "%" does not exist', p_role;
    END IF;

    -- Prevent accidentally revoking your own access.  pg_has_role checks
    -- both direct identity (current_user = p_role) and inherited membership
    -- (current_user is a member of p_role), so revoking a parent role that
    -- the caller depends on is also caught.
    -- Superusers are exempt: pg_has_role returns true for all roles when the
    -- caller is a superuser, and superusers can always re-grant themselves.
    IF NOT EXISTS (
        SELECT 1
        FROM pg_roles
        WHERE rolname = current_user
          AND rolsuper
    )
       AND pg_has_role(current_user, p_role, 'MEMBER') THEN
        RAISE EXCEPTION 'cannot revoke df privileges from "%" because the current role ("%") is a member of it — use a different administrator', p_role, current_user;
    END IF;

    -- CASCADE: if the target role granted sub-grants (via WITH GRANT OPTION),
    -- CASCADE ensures those dependent privileges are also revoked.
    -- Column-level revokes must match the column-level grants from grant_usage().
    EXECUTE format('REVOKE SELECT, INSERT, UPDATE, DELETE ON df.vars FROM %I CASCADE', p_role);
    EXECUTE format('REVOKE INSERT (id, instance_id, node_type, query, result_name, left_node, right_node, submitted_by, database) ON df.nodes FROM %I CASCADE', p_role);
    EXECUTE format('REVOKE SELECT ON df.nodes FROM %I CASCADE', p_role);
    EXECUTE format('REVOKE INSERT (id, label, root_node, submitted_by, database) ON df.instances FROM %I CASCADE', p_role);
    EXECUTE format('REVOKE UPDATE (status, updated_at) ON df.instances FROM %I CASCADE', p_role);
    EXECUTE format('REVOKE SELECT ON df.instances FROM %I CASCADE', p_role);

    FOR func_oid IN
        SELECT p.oid FROM pg_proc p
        JOIN pg_namespace n ON p.pronamespace = n.oid
        WHERE n.nspname = 'df'
    LOOP
        BEGIN
            EXECUTE format('REVOKE EXECUTE ON FUNCTION %s FROM %I CASCADE', func_oid::regprocedure, p_role);
        EXCEPTION WHEN insufficient_privilege THEN
            NULL;
        END;
    END LOOP;

    EXECUTE format('REVOKE USAGE ON SCHEMA df FROM %I CASCADE', p_role);

    RAISE NOTICE 'pg_durable: revoked df usage privileges granted by "%" from "%"', current_user, p_role;
END;
$fn$;

-- ============================================================================
-- 7. Note on df.http() PUBLIC privilege
--    v0.1.1 issued GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA df TO PUBLIC,
--    which included df.http().  New installs (0.2.0+) revoke this at
--    CREATE EXTENSION time, but upgraded installs retain the PUBLIC grant.
--
--    If you want to adopt opt-in HTTP permissions after upgrading, run:
--      REVOKE EXECUTE ON FUNCTION df.http(text,text,text,jsonb,integer) FROM PUBLIC;
--    Then use df.grant_usage(role, include_http => true) to re-grant to roles
--    that should have HTTP access.
-- ============================================================================