Bi-temporal Data Modeling in PostgreSQL with Range Types & Constraints

-- Ensure the btree_gist extension is available for our constraints
CREATE EXTENSION IF NOT EXISTS btree_gist;

CREATE TABLE product_prices (
    -- A surrogate key for this specific version/row. This is the primary key.
    version_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    -- A stable identifier for the logical entity (the product).
    -- All versions of a single product's price share the same product_id.
    product_id UUID NOT NULL,

    -- The payload data.
    price NUMERIC(10, 2) NOT NULL CHECK (price > 0),
    currency CHAR(3) NOT NULL,

    -- == Bi-temporal Axes ==

    -- Valid Time: When was this price effective in the real world?
    -- Using tsrange for a continuous period. '[)' means inclusive start, exclusive end.
    valid_ts TSRANGE NOT NULL,

    -- Transaction Time: When did the database believe this fact?
    -- This range starts when the record is inserted and extends to 'infinity'.
    -- When corrected/superseded, the upper bound is set to the current time.
    tx_ts TSTZRANGE NOT NULL DEFAULT TSTZRANGE(NOW(), 'infinity', '[)'),

    -- == The Magic: Ensuring Data Integrity at the DB Level ==

    -- This constraint is the cornerstone of our bi-temporal model.
    -- For any given product_id, you cannot have overlapping valid_ts ranges
    -- for records that are *currently active* in transaction time.
    CONSTRAINT no_overlapping_valid_time_for_current_records
    EXCLUDE USING GIST (
        product_id WITH =,
        valid_ts WITH &&
    ) WHERE (tx_ts @> 'infinity'::timestamptz)
);

-- Indexes are critical for performance of bi-temporal queries.
-- A GiST index is needed to support the exclusion constraint and range operators.
CREATE INDEX idx_product_prices_valid_ts_gist ON product_prices USING GIST (product_id, valid_ts);

-- A standard B-tree index on transaction time is also essential for "as-of" queries.
CREATE INDEX idx_product_prices_tx_ts ON product_prices USING GIST (product_id, tx_ts);

CREATE OR REPLACE FUNCTION create_new_product_price(
    p_product_id UUID,
    p_price NUMERIC(10, 2),
    p_currency CHAR(3),
    p_valid_from TIMESTAMPTZ
) RETURNS UUID AS $$
DECLARE
    v_version_id UUID;
BEGIN
    INSERT INTO product_prices (product_id, price, currency, valid_ts)
    VALUES (p_product_id, p_price, p_currency, TSRANGE(p_valid_from, 'infinity', '[)'))
    RETURNING version_id INTO v_version_id;

    RETURN v_version_id;
END;
$$ LANGUAGE plpgsql;

-- Set the initial price for a new product, valid from today onwards.
SELECT create_new_product_price('a1b2c3d4-e5f6-4a7b-8c9d-0e1f2a3b4c5d', 199.99, 'USD', NOW());

CREATE OR REPLACE FUNCTION update_product_price(
    p_product_id UUID,
    p_new_price NUMERIC(10, 2),
    p_effective_from TIMESTAMPTZ
) RETURNS UUID AS $$
DECLARE
    v_current_version RECORD;
    v_new_version_id UUID;
BEGIN
    -- Use a SERIALIZABLE transaction to prevent race conditions where two updates
    -- might try to supersede the same record simultaneously.
    SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;

    -- Find the current active version that we need to supersede.
    -- We lock it to ensure no other transaction can modify it.
    SELECT * INTO v_current_version
    FROM product_prices
    WHERE product_id = p_product_id
      AND valid_ts @> p_effective_from
      AND tx_ts @> 'infinity'::timestamptz
    FOR UPDATE;

    IF NOT FOUND THEN
        RAISE EXCEPTION 'No active price found for product % at effective date %', p_product_id, p_effective_from;
    END IF;

    -- 1. Close the transaction time of the old record.
    -- Its valid time is split at the effective date of the new price.
    UPDATE product_prices
    SET tx_ts = TSTZRANGE(lower(tx_ts), NOW(), '[)')
    WHERE version_id = v_current_version.version_id;

    -- 2. If the old price was valid for a period, we need to re-insert the part
    -- of its validity that occurred *before* the new price took effect.
    IF lower(v_current_version.valid_ts) < p_effective_from THEN
        INSERT INTO product_prices(product_id, price, currency, valid_ts, tx_ts)
        VALUES (
            v_current_version.product_id,
            v_current_version.price,
            v_current_version.currency,
            TSRANGE(lower(v_current_version.valid_ts), p_effective_from, '[)'),
            TSTZRANGE(NOW(), 'infinity', '[)')
        );
    END IF;

    -- 3. Insert the new price record, valid from the effective date onwards.
    INSERT INTO product_prices(product_id, price, currency, valid_ts, tx_ts)
    VALUES (
        p_product_id,
        p_new_price,
        v_current_version.currency, -- Assuming currency doesn't change
        TSRANGE(p_effective_from, 'infinity', '[)'),
        TSTZRANGE(NOW(), 'infinity', '[)')
    )
    RETURNING version_id INTO v_new_version_id;

    RETURN v_new_version_id;
END;
$$ LANGUAGE plpgsql;

-- On 2024-02-15, we decide the price for product 'a1b2...' will change to $219.99 starting 2024-03-01
SELECT update_product_price('a1b2c3d4-e5f6-4a7b-8c9d-0e1f2a3b4c5d', 219.99, '2024-03-01'::timestamptz);

CREATE OR REPLACE FUNCTION delete_product_price(
    p_product_id UUID,
    p_effective_from TIMESTAMPTZ
) RETURNS VOID AS $$
BEGIN
    SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;

    -- This logic is similar to the update. We find the active record,
    -- terminate its transaction time, and re-insert the portion of its
    -- validity before the deletion date.
    WITH current_version AS (
        SELECT version_id, valid_ts, tx_ts, product_id, price, currency
        FROM product_prices
        WHERE product_id = p_product_id
          AND valid_ts @> p_effective_from
          AND tx_ts @> 'infinity'::timestamptz
        FOR UPDATE
    ),
    terminate_old AS (
        UPDATE product_prices
        SET tx_ts = TSTZRANGE(lower(current_version.tx_ts), NOW(), '[)')
        FROM current_version
        WHERE product_prices.version_id = current_version.version_id
    )
    -- Re-insert the valid portion before the deletion date
    INSERT INTO product_prices(product_id, price, currency, valid_ts, tx_ts)
    SELECT
        product_id, price, currency,
        TSRANGE(lower(valid_ts), p_effective_from, '[)'),
        TSTZRANGE(NOW(), 'infinity', '[)')
    FROM current_version
    WHERE lower(valid_ts) < p_effective_from;

    IF NOT FOUND THEN
        RAISE EXCEPTION 'No active price found for product % to delete at effective date %', p_product_id, p_effective_from;
    END IF;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION correct_historical_price(
    p_product_id UUID,
    p_correct_price NUMERIC(10, 2),
    p_valid_range TSRANGE
) RETURNS UUID AS $$
DECLARE
    v_incorrect_version RECORD;
    v_new_version_id UUID;
BEGIN
    SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;

    -- Find the incorrect historical record(s). This query is tricky because we need to
    -- find what we *believed* to be true just before this transaction started.
    -- For simplicity, we'll find the version(s) that were active and covered this valid range.
    -- A more complex implementation might need to handle multiple corrections.
    SELECT * INTO v_incorrect_version
    FROM product_prices
    WHERE product_id = p_product_id
      AND valid_ts && p_valid_range -- It overlaps with the period we are correcting
      AND tx_ts @> 'infinity'::timestamptz -- It is currently believed to be true
    FOR UPDATE;

    IF NOT FOUND THEN
        RAISE EXCEPTION 'No current record found to correct for product % in range %', p_product_id, p_valid_range;
    END IF;

    -- 1. Terminate the incorrect record's transaction time.
    UPDATE product_prices
    SET tx_ts = TSTZRANGE(lower(tx_ts), NOW(), '[)')
    WHERE version_id = v_incorrect_version.version_id;

    -- Note: This is a simplified correction. A full implementation would need to handle
    -- splitting the original range if the correction only applies to a sub-period.
    -- Here we assume the correction replaces the entire old record's validity.

    -- 2. Insert the new, correct record.
    -- Its valid_ts is in the past, but its tx_ts starts now.
    INSERT INTO product_prices(product_id, price, currency, valid_ts, tx_ts)
    VALUES (
        v_incorrect_version.product_id,
        p_correct_price,
        v_incorrect_version.currency,
        p_valid_range,
        TSTZRANGE(NOW(), 'infinity', '[)')
    )
    RETURNING version_id INTO v_new_version_id;

    RETURN v_new_version_id;
END;
$$ LANGUAGE plpgsql;

-- What is the price on 2023-03-15?
SELECT product_id, price, currency
FROM product_prices
WHERE product_id = '...' 
  AND valid_ts @> '2023-03-15'::timestamptz
  AND tx_ts @> 'infinity'::timestamptz; -- Only currently active records

-- Result: $105 (v2)

-- What did we believe the price for 2023-03-15 was, back on 2023-02-01?
SELECT product_id, price, currency
FROM product_prices
WHERE product_id = '...'
  AND valid_ts @> '2023-03-15'::timestamptz
  AND tx_ts @> '2023-02-01'::timestamptz; -- As we knew it on this date

-- Result: $100 (v1), because the correction (v2) hadn't happened yet.

-- Show me the price for the product that was valid on 2023-07-01,
-- as we knew it on 2023-06-15.
SELECT product_id, price, currency
FROM product_prices
WHERE product_id = '...'
  AND valid_ts @> '2023-07-01'::timestamptz -- The 'as-at' part
  AND tx_ts @> '2023-06-15'::timestamptz;   -- The 'as-of' part

-- Result: $110 (v3), because the price update was known on 2023-06-15.