restub i do it myself

cmd/pg-pubsub-demo/main.go

@@ -3,13 +3,11 @@ package main
 import (
 	"context"
 	"database/sql"
-	"encoding/json"
 	"flag"
 	"fmt"
+	"io"
 	"log"
-	"math/rand"
 	"os"
-	"strconv"
 	"time"
 
 	"pg/src/with"
@@ -31,109 +29,25 @@ func run(ctx context.Context) error {
 	}
 
 	return with.PSQL(ctx, *c, func(pg *sql.DB) error {
-		log.Println("creating tables...")
+		if _, err := pg.ExecContext(ctx, `
-		for _, sql := range []string{
+		`); err != nil {
-			`CREATE SCHEMA IF NOT EXISTS "pubsub"`,
-
-			`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a" (
-				partition INT PRIMARY KEY,
-				partition_offset BIGINT NOT NULL DEFAULT 0
-			)`,
-			`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a_partition_0" (
-				id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
-				partition_offset BIGINT UNIQUE NOT NULL,
-				payload BYTEA NOT NULL,
-				created_at TIMESTAMP NOT NULL DEFAULT NOW()
-			)`,
-			`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a_partition_1" (
-				id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
-				partition_offset BIGINT UNIQUE NOT NULL,
-				payload BYTEA NOT NULL,
-				created_at TIMESTAMP NOT NULL DEFAULT NOW()
-			)`,
-
-			`INSERT INTO "pubsub"."topic_a" (partition, partition_offset) VALUES (0, 0) ON CONFLICT DO NOTHING`,
-			`INSERT INTO "pubsub"."topic_a" (partition, partition_offset) VALUES (1, 0) ON CONFLICT DO NOTHING`,
-
-			`CREATE TABLE IF NOT EXISTS "pubsub"."consumer_partition_offset" (
-				consumer_group TEXT NOT NULL,
-				partition INT REFERENCES "pubsub"."topic_a"("partition"),
-				partition_offset BIGINT NOT NULL DEFAULT 0,
-				PRIMARY KEY (consumer_group, partition)
-			)`,
-		} {
-			if _, err := pg.ExecContext(ctx, sql); err != nil {
-				return fmt.Errorf("failed setup sql %s: %w", sql, err)
-			}
-		}
-
-		pub := func(a any) error {
-			b, _ := json.Marshal(a)
-			partition := rand.Int() % 2
-			partitionS := strconv.Itoa(partition)
-			_, err := pg.ExecContext(ctx, `
-				WITH reserve AS (
-					UPDATE "pubsub"."topic_a"
-					SET partition_offset=partition_offset+1
-					WHERE partition=`+partitionS+`
-					RETURNING (partition_offset-1) AS offset
-				)
-				INSERT INTO "pubsub"."topic_a_partition_`+partitionS+`" (partition_offset, payload)
-				SELECT reserve.offset, $1 FROM reserve
-			`, b)
 			return err
 		}
-		sub := func(group string) (any, error) {
-			partition := rand.Int() % 2
-			partitionS := strconv.Itoa(partition)
 
-			if _, err := pg.ExecContext(ctx, `
+		pub := func(topic string, a any) error {
-				INSERT INTO
+			return io.EOF
-				"pubsub"."consumer_partition_offset"
+		}
-				(consumer_group, partition)
+		sub := func(topic, group string) (int, any, error) {
-				VALUES
+			return 0, nil, io.EOF
-				($1, $2)
+		}
-				ON CONFLICT DO NOTHING
+		commit := func(topic, group string, offset int) error {
-			`, group, partition); err != nil {
+			return io.EOF
-				return nil, err
-			}
-
-			row := pg.QueryRowContext(ctx, `
-				WITH tip AS (
-					SELECT partition_offset-1
-					FROM "pubsub"."topic_a"
-					WHERE partition=`+partitionS+`
-				)
-				, to_claim AS (
-					SELECT partition_offset
-					FROM "pubsub"."consumer_partition_offset"
-					WHERE consumer_group=$1 AND partition=$2
-					FOR UPDATE SKIP LOCKED
-				)
-				, claimed AS (
-					UPDATE "pubsub"."consumer_partition_offset" consumer_partition_offset
-					SET consumer_partition_offset.partition_offset=to_claim.partition_offset+1
-					FROM to_claim
-					WHERE consumer_partition_offset.consumer_group=$1 AND consumer_partition_offset.partition=$2
-					RETURNING consumer_partition_offset.partition_offset
-				)
-				SELECT partition_offset FROM claimed
-			`, group, partition)
-			if err := row.Err(); err != nil {
-				return nil, err
-			}
-
-			var result any
-			if err := row.Scan(&result); err != nil {
-				return nil, err
-			}
-
-			return result, row.Err()
 		}
 
 		go with.Every(ctx, time.Second, func() {
 			for i := 0; i < 2; i++ {
-				if err := pub(1); err != nil {
+				topic := fmt.Sprintf("topic_%d", i)
+				if err := pub(topic, 1); err != nil {
 					log.Printf("failed pub: %v", err)
 				} else {
 					log.Printf("pubbed")
@@ -144,11 +58,16 @@ func run(ctx context.Context) error {
 		with.Every(ctx, time.Second, func() {
 			if err := func() error {
 				for i := 0; i < 2; i++ {
-					group := fmt.Sprintf("group_%d", i)
+					topic := fmt.Sprintf("topic_%d", i)
-					if v, err := sub(group); err != nil {
+					for j := 0; j < 2; j++ {
-						return err
+						group := fmt.Sprintf("group_%d", i)
-					} else {
+						if offset, v, err := sub(topic, group); err != nil {
-						log.Printf("%s => %v", group, v)
+							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