Database

import json
import os
import threading

import psycopg

from iii import InitOptions, ObserveRejection, register_worker

worker = register_worker(
    os.environ.get("III_URL"),
    InitOptions(worker_name="database-worker", namespace="orders"),
)

ROW_CHANGED_ARGS = {
    "type": "object",
    "additionalProperties": False,
    "required": ["db"],
    "properties": {
        "db": {"type": "string"},
        "table": {"type": "string"},
        "ops": {
            "type": "array",
            "items": {"type": "string", "enum": ["insert", "update", "delete"]},
        },
    },
}

ROW_CHANGE = {
    "type": "object",
    "required": ["db", "table", "op", "row"],
    "properties": {
        "db": {"type": "string"},
        "table": {"type": "string"},
        "op": {"type": "string"},
        "row": {"type": "object"},
    },
}

pools = {"primary": psycopg.connect(os.environ["PRIMARY_DATABASE_URL"], autocommit=True)}


def row_key(db: str, table: str, op: str) -> str:
    return f"{db}|{table}|{op}"


# the tables this worker installed a trigger on, and how many observers want them
holds: dict[str, int] = {}


# the observer's table, or every table of its database
def tables_of(pool: psycopg.Connection, table: str | None) -> list[str]:
    if table is not None:
        return [table]
    rows = pool.execute(
        "SELECT tablename FROM pg_tables WHERE schemaname = 'public'"
    ).fetchall()
    return [row[0] for row in rows]


# set by id, before the function is registered, so the hooks run for every observer
@worker.on_observer_added("database::on-row-changed")
def install(obs):
    db, table = obs.args["db"], obs.args.get("table")
    pool = pools.get(db)
    if pool is None:
        raise ObserveRejection("UNKNOWN_DB", f"no database '{db}'")
    for name in tables_of(pool, table):
        key = f"{db}|{name}"
        holds[key] = holds.get(key, 0) + 1
        if holds[key] > 1:
            continue
        pool.execute(
            f"CREATE OR REPLACE TRIGGER iii_{name}_changed "
            f"AFTER INSERT OR UPDATE OR DELETE ON {name} "
            "FOR EACH ROW EXECUTE FUNCTION iii_notify_row_changed()"
        )


@worker.on_observer_removed("database::on-row-changed")
def release(obs):
    db, table = obs.args["db"], obs.args.get("table")
    pool = pools.get(db)
    if pool is None:
        return
    for name in tables_of(pool, table):
        key = f"{db}|{name}"
        holds[key] = holds.get(key, 1) - 1
        if holds[key] > 0:
            continue
        del holds[key]
        pool.execute(f"DROP TRIGGER IF EXISTS iii_{name}_changed ON {name}")


# an observer gets the changes of its database, table and operation
worker.on_observation_match(
    "database::on-row-changed",
    observer=lambda obs: [
        row_key(obs.args["db"], obs.args.get("table", "*"), op)
        for op in obs.args.get("ops", ["*"])
    ],
    call=lambda call: [
        row_key(call.input["db"], table, op)
        for table in (call.input["table"], "*")
        for op in (call.input["op"], "*")
    ],
)

row_changed = worker.register_function(
    "database::on-row-changed",
    lambda change: change,
    request_format=ROW_CHANGE,
    response_format=ROW_CHANGE,
    args_format=ROW_CHANGED_ARGS,
    private=True,
)

# each notification of the database becomes a call, when an observer wants it
listener = psycopg.connect(os.environ["PRIMARY_DATABASE_URL"], autocommit=True)
listener.execute("LISTEN iii_row_changed")


def on_notifications() -> None:
    for notification in listener.notifies():
        change = json.loads(notification.payload)
        if not row_changed.is_observed({"input": change}):
            continue
        worker.call({"function_id": "database::on-row-changed", "payload": change})


threading.Thread(target=on_notifications, daemon=True).start()

use std::collections::HashMap;
use std::sync::Arc;

use iii_sdk::protocol::CallRequest;
use iii_sdk::{
    register_worker, AddedHook, Call, Error, InitOptions, Keys, Match, ObserveRejection,
    Observer, RegisterFunction, RemovedHook,
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::postgres::{PgListener, PgPool};
use tokio::sync::Mutex;

let url = std::env::var("III_URL").expect("III_URL must be set");
let worker = register_worker(
    &url,
    InitOptions {
        namespace: Some("orders".into()),
        ..Default::default()
    },
);

#[derive(Deserialize, JsonSchema)]
struct RowChangedArgs {
    db: String,
    table: Option<String>,
    ops: Option<Vec<String>>,
}

#[derive(Serialize, Deserialize, JsonSchema)]
struct RowChange {
    db: String,
    table: String,
    op: String,
    row: Value,
}

fn row_key(db: &str, table: &str, op: &str) -> String {
    format!("{db}|{table}|{op}")
}

// the keys of an observer: its database, its table and the operations it names
fn observer_keys(obs: &Observer) -> Keys<'_> {
    let args: RowChangedArgs = obs.args().expect("validated against args_format");
    let table = args.table.unwrap_or_else(|| "*".into());
    let ops = args.ops.unwrap_or_else(|| vec!["*".into()]);
    ops.iter()
        .map(|op| row_key(&args.db, &table, op))
        .collect::<Vec<_>>()
        .into()
}

// the keys of a change
fn call_keys(call: &Call) -> Keys<'_> {
    let field = |name: &str| call.input[name].as_str().unwrap_or_default().to_owned();
    let (db, table, op) = (field("db"), field("table"), field("op"));
    [table.as_str(), "*"]
        .into_iter()
        .flat_map(|t| [op.as_str(), "*"].map(|o| row_key(&db, t, o)))
        .collect::<Vec<_>>()
        .into()
}

// the observer's table, or every table of its database
async fn tables_of(pool: &PgPool, table: Option<&str>) -> Result<Vec<String>, ObserveRejection> {
    let Some(table) = table else {
        let rows: Vec<(String,)> =
            sqlx::query_as("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
                .fetch_all(pool)
                .await
                .map_err(|e| ObserveRejection::new("DB_ERROR", e.to_string()))?;
        return Ok(rows.into_iter().map(|(name,)| name).collect());
    };
    Ok(vec![table.to_string()])
}

#[derive(Clone)]
struct Capture {
    pools: Arc<HashMap<String, PgPool>>,
    // the tables this worker installed a trigger on, and how many observers want them
    holds: Arc<Mutex<HashMap<String, usize>>>,
}
impl AddedHook for Capture {
    async fn on_added(&self, obs: Observer) -> Result<(), ObserveRejection> {
        let args: RowChangedArgs = obs.args()?;
        let Some(pool) = self.pools.get(&args.db) else {
            return Err(ObserveRejection::new(
                "UNKNOWN_DB",
                format!("no database '{}'", args.db),
            ));
        };
        for name in tables_of(pool, args.table.as_deref()).await? {
            let key = format!("{}|{}", args.db, name);
            let first = {
                let mut holds = self.holds.lock().await;
                let count = holds.entry(key).or_insert(0);
                *count += 1;
                *count == 1
            };
            if !first {
                continue;
            }
            sqlx::query(&format!(
                "CREATE OR REPLACE TRIGGER iii_{name}_changed AFTER INSERT OR UPDATE OR DELETE ON {name} \
                 FOR EACH ROW EXECUTE FUNCTION iii_notify_row_changed()"
            ))
            .execute(pool)
            .await
            .map_err(|e| ObserveRejection::new("DB_ERROR", e.to_string()))?;
        }
        Ok(())
    }
}
impl RemovedHook for Capture {
    async fn on_removed(&self, obs: Observer) {
        let Ok(args) = obs.args::<RowChangedArgs>() else {
            return;
        };
        let Some(pool) = self.pools.get(&args.db) else {
            return;
        };
        let Ok(tables) = tables_of(pool, args.table.as_deref()).await else {
            return;
        };
        for name in tables {
            let key = format!("{}|{}", args.db, name);
            let last = {
                let mut holds = self.holds.lock().await;
                let count = holds.entry(key.clone()).or_insert(1);
                *count -= 1;
                if *count > 0 {
                    return;
                }
                holds.remove(&key);
                true
            };
            if !last {
                continue;
            }
            let _ = sqlx::query(&format!(
                "DROP TRIGGER IF EXISTS iii_{name}_changed ON {name}"
            ))
            .execute(pool)
            .await;
        }
    }
}

let capture = Capture {
    pools: Arc::new(HashMap::from([(
        "primary".to_string(),
        PgPool::connect(&std::env::var("PRIMARY_DATABASE_URL")?).await?,
    )])),
    holds: Arc::new(Mutex::new(HashMap::new())),
};

// set by id, before the function is registered, so the hooks run for every observer
worker.on_observer_added_for("database::on-row-changed", capture.clone());
worker.on_observer_removed_for("database::on-row-changed", capture.clone());
// an observer gets the changes of its database, table and operation
worker.on_observation_match_for("database::on-row-changed", Match::keys(observer_keys, call_keys));

let row_changed = worker.register_function(
    "database::on-row-changed",
    RegisterFunction::new(|change: RowChange| Ok::<_, Error>(change))
        .args_format::<RowChangedArgs>()
        .private(),
);

// each notification of the database becomes a call, when an observer wants it
let mut listener = PgListener::connect(&std::env::var("PRIMARY_DATABASE_URL")?).await?;
listener.listen("iii_row_changed").await?;
loop {
    let change: RowChange = serde_json::from_str(listener.recv().await?.payload())?;
    let input = serde_json::to_value(&change)?;
    let call = Call {
        input,
        ..Call::default()
    };
    if !row_changed.is_observed_for(&call) {
        continue;
    }
    worker
        .call(CallRequest::new("database::on-row-changed", call.input))
        .await?;
}

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "os"
    "sync"

    "github.com/jackc/pgx/v5"
    iii "github.com/iii-hq/iii/sdk/packages/go/iii"
)

type rowChangedArgs struct {
    DB    string   `json:"db"`
    Table string   `json:"table,omitempty"`
    Ops   []string `json:"ops,omitempty"`
}

type rowChange struct {
    DB    string          `json:"db"`
    Table string          `json:"table"`
    Op    string          `json:"op"`
    Row   json.RawMessage `json:"row"`
}

var rowChangedArgsSchema = json.RawMessage(`{
    "type": "object",
    "additionalProperties": false,
    "required": ["db"],
    "properties": {
        "db": {"type": "string"},
        "table": {"type": "string"},
        "ops": {"type": "array", "items": {"type": "string", "enum": ["insert", "update", "delete"]}}
    }
}`)

func rowKey(db, table, op string) string { return db + "|" + table + "|" + op }

// the observer's table, or every table of its database
func tablesOf(ctx context.Context, pool *pgx.Conn, table string) ([]string, error) {
    if table != "" {
        return []string{table}, nil
    }
    rows, err := pool.Query(ctx, "SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
    if err != nil {
        return nil, err
    }
    return pgx.CollectRows(rows, pgx.RowTo[string])
}

func main() {
    ctx := context.Background()
    worker := iii.RegisterWorker(os.Getenv("III_URL"), iii.WithNamespace("orders"))

    primary, err := pgx.Connect(ctx, os.Getenv("PRIMARY_DATABASE_URL"))
    if err != nil {
        log.Fatal(err)
    }
    pools := map[string]*pgx.Conn{"primary": primary}

    // the tables this worker installed a trigger on, and how many observers want them
    var mu sync.Mutex
    holds := map[string]int{}

    // set by id, before the function is registered, so the hooks run for every observer
    worker.OnObserverAddedFor("database::on-row-changed", func(ctx context.Context, obs *iii.Observer) error {
        var args rowChangedArgs
        if err := obs.DecodeArgs(&args); err != nil {
            return err
        }
        pool, ok := pools[args.DB]
        if !ok {
            return iii.ObserveRejection("UNKNOWN_DB", "no database '"+args.DB+"'")
        }
        tables, err := tablesOf(ctx, pool, args.Table)
        if err != nil {
            return err
        }
        for _, name := range tables {
            key := args.DB + "|" + name
            mu.Lock()
            holds[key]++
            first := holds[key] == 1
            mu.Unlock()
            if !first {
                continue
            }
            if _, err := pool.Exec(ctx, fmt.Sprintf(
                "CREATE OR REPLACE TRIGGER iii_%[1]s_changed AFTER INSERT OR UPDATE OR DELETE ON %[1]s "+
                    "FOR EACH ROW EXECUTE FUNCTION iii_notify_row_changed()", name)); err != nil {
                return err
            }
        }
        return nil
    })

    worker.OnObserverRemovedFor("database::on-row-changed", func(ctx context.Context, obs *iii.Observer) {
        var args rowChangedArgs
        if obs.DecodeArgs(&args) != nil {
            return
        }
        pool, ok := pools[args.DB]
        if !ok {
            return
        }
        tables, err := tablesOf(ctx, pool, args.Table)
        if err != nil {
            return
        }
        for _, name := range tables {
            key := args.DB + "|" + name
            mu.Lock()
            holds[key]--
            last := holds[key] <= 0
            if last {
                delete(holds, key)
            }
            mu.Unlock()
            if !last {
                continue
            }
            pool.Exec(ctx, fmt.Sprintf("DROP TRIGGER IF EXISTS iii_%[1]s_changed ON %[1]s", name))
        }
    })

    // an observer gets the changes of its database, table and operation
    worker.OnObservationMatchFor("database::on-row-changed", iii.MatchKeys(
        func(obs *iii.Observer) []string {
            var args rowChangedArgs
            _ = obs.DecodeArgs(&args)
            table, ops := args.Table, args.Ops
            if table == "" {
                table = "*"
            }
            if len(ops) == 0 {
                ops = []string{"*"}
            }
            keys := make([]string, 0, len(ops))
            for _, op := range ops {
                keys = append(keys, rowKey(args.DB, table, op))
            }
            return keys
        },
        func(call *iii.Call) []string {
            var change rowChange
            _ = json.Unmarshal(call.Input, &change)
            keys := make([]string, 0, 4)
            for _, table := range []string{change.Table, "*"} {
                for _, op := range []string{change.Op, "*"} {
                    keys = append(keys, rowKey(change.DB, table, op))
                }
            }
            return keys
        },
    ))

    rowChanged, err := worker.RegisterFunction("database::on-row-changed", iii.Identity,
        iii.FunctionOptions{ArgsFormat: rowChangedArgsSchema, Private: true})
    if err != nil {
        log.Fatal(err)
    }

    // each notification of the database becomes a call, when an observer wants it
    if _, err := primary.Exec(ctx, "LISTEN iii_row_changed"); err != nil {
        log.Fatal(err)
    }
    for {
        notification, err := primary.WaitForNotification(ctx)
        if err != nil {
            log.Fatal(err)
        }
        var change rowChange
        if err := json.Unmarshal([]byte(notification.Payload), &change); err != nil {
            continue
        }
        payload, _ := json.Marshal(change)
        if !rowChanged.IsObservedFor(&iii.Call{Input: payload}) {
            continue
        }
        worker.Call(ctx, iii.CallRequest{FunctionID: "database::on-row-changed", Payload: payload})
    }
}