pg/cmd/pg-pubsub-demo/main.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"pg/src/with"

	_ "github.com/lib/pq"
)

func main() {
	if err := with.Context(run); err != nil {
		panic(err)
	}
}

func run(ctx context.Context) error {
	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	c := fs.String("c", "postgresql://pulsegres:pulsegres@localhost:15432", "conn string")
	if err := fs.Parse(os.Args[1:]); err != nil {
		panic(err)
	}

	return with.PSQL(ctx, *c, func(pg *sql.DB) error {
		log.Println("setup...")
		if _, err := pg.ExecContext(ctx, `
			DROP SCHEMA "pubsub" CASCADE;

			CREATE SCHEMA IF NOT EXISTS "pubsub";

			CREATE TABLE IF NOT EXISTS "pubsub".topic (
				id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
				created_at TIMESTAMP DEFAULT NOW(),
				updated_at TIMESTAMP DEFAULT NOW(),
				partitions INTEGER NOT NULL DEFAULT 10,
				name TEXT UNIQUE
			);

			CREATE TABLE IF NOT EXISTS "pubsub".group (
				id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
				created_at TIMESTAMP DEFAULT NOW(),
				updated_at TIMESTAMP DEFAULT NOW(),
				name TEXT UNIQUE
			);

			CREATE TABLE IF NOT EXISTS "pubsub".group_topic_partition_offset (
				group_name TEXT REFERENCES "pubsub".group(name) ON DELETE CASCADE,
				topic_name TEXT REFERENCES "pubsub".topic(name) ON DELETE CASCADE,
				partition INTEGER NOT NULL DEFAULT 0,
				partition_offset INTEGER NOT NULL,
				PRIMARY KEY (group_name, topic_name, partition)
			);

			CREATE TABLE IF NOT EXISTS "pubsub".data (
				topic_name TEXT REFERENCES "pubsub".topic(name) ON DELETE CASCADE,
				partition INTEGER NOT NULL DEFAULT 0,
				partition_offset INTEGER NOT NULL DEFAULT 0,
				payload BYTEA NOT NULL,
				PRIMARY KEY (topic_name, partition, partition_offset)
			) PARTITION BY HASH (topic_name, partition);

			CREATE TABLE IF NOT EXISTS "pubsub".data_0 PARTITION OF "pubsub".data FOR VALUES WITH (modulus 3, remainder 0);
			CREATE TABLE IF NOT EXISTS "pubsub".data_1 PARTITION OF "pubsub".data FOR VALUES WITH (modulus 3, remainder 1);
			CREATE TABLE IF NOT EXISTS "pubsub".data_2 PARTITION OF "pubsub".data FOR VALUES WITH (modulus 3, remainder 2);
		`); err != nil {
			return err
		}

		pub := func(topic string, a any) error {
			payload, _ := json.Marshal(a)
			if _, err := pg.ExecContext(ctx, `
				INSERT INTO "pubsub".topic (name)
					VALUES ($1)
					ON CONFLICT DO NOTHING;
			`, topic); err != nil {
				return err
			}
			_, err := pg.ExecContext(ctx, `
				WITH topic AS (
					SELECT
						name,
						FLOOR(100 * RANDOM())::INTEGER % partitions AS partition
					FROM "pubsub".topic
					WHERE name=$1
				)
				, partition_offset AS (
					SELECT
						topic.partition AS partition,
						COALESCE(MAX(data.partition_offset), -1)+1 AS next_offset
					FROM topic
					LEFT JOIN "pubsub".data data ON data.topic_name=topic.name AND data.partition=topic.partition
					GROUP BY topic.partition
				)
				INSERT INTO "pubsub".data
					(topic_name, partition, partition_offset, payload)
				SELECT
					$1, partition_offset.partition, partition_offset.next_offset, $2
				FROM partition_offset
			`, topic, payload)
			return err
		}
		sub := func(topic, group string) (int, any, error) {
			return 0, nil, io.EOF
		}
		commit := func(topic, group string, offset int) error {
			return io.EOF
		}

		go with.Every(ctx, time.Second, func() {
			for i := 0; i < 2; i++ {
				topic := fmt.Sprintf("topic_%d", i)
				if err := pub(topic, 1); err != nil {
					log.Printf("failed pub: %v", err)
				} else {
					log.Printf("pubbed")
				}
			}
		})

		with.Every(ctx, time.Second, func() {
			if err := func() error {
				for i := 0; i < 2; i++ {
					topic := fmt.Sprintf("topic_%d", i)
					for j := 0; j < 2; j++ {
						group := fmt.Sprintf("group_%d", i)
						if offset, v, err := sub(topic, group); err != nil {
							return err
						} else if err := commit(topic, group, offset); err != nil {
							return err
						} else {
							log.Printf("subbed/committed offset %d: %v", offset, v)
						}
					}
				}
				return nil
			}(); err != nil {
				log.Printf("failed subs: %v", err)
			}
		})

		return ctx.Err()
	})
}