Initial open source release

internal/store/postgres/message_plan_integration_test.go

@@ -0,0 +1,160 @@
+package postgres
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/jackc/pgx/v5"
+
+	"telesrv/internal/domain"
+)
+
+func TestMessagePartitionSeekPlansUseIndexes(t *testing.T) {
+	pool := testPool(t)
+	ctx := context.Background()
+	suffix := randomSuffix(t)
+
+	users := NewUserStore(pool)
+	sender, err := users.Create(ctx, domain.User{
+		AccessHash: 51,
+		Phone:      "+1999" + suffix + "01",
+		FirstName:  "PlanSender",
+	})
+	if err != nil {
+		t.Fatalf("create sender: %v", err)
+	}
+	recipient, err := users.Create(ctx, domain.User{
+		AccessHash: 52,
+		Phone:      "+1999" + suffix + "02",
+		FirstName:  "PlanRecipient",
+	})
+	if err != nil {
+		t.Fatalf("create recipient: %v", err)
+	}
+	t.Cleanup(func() {
+		_, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID})
+	})
+
+	messages := NewMessageStore(pool)
+	for i := 0; i < 3; i++ {
+		if _, err := messages.SendPrivateText(ctx, domain.SendPrivateTextRequest{
+			SenderUserID:    sender.ID,
+			RecipientUserID: recipient.ID,
+			RandomID:        int64(9000 + i),
+			Message:         "plan check",
+			Date:            1700000300 + i,
+		}); err != nil {
+			t.Fatalf("seed message %d: %v", i, err)
+		}
+	}
+	if _, err := pool.Exec(ctx, `
+		UPDATE dispatch_outbox
+		SET status = 'dispatching',
+		    updated_at = now() - interval '1 minute'
+		WHERE target_user_id = $1
+	`, recipient.ID); err != nil {
+		t.Fatalf("mark dispatch stale: %v", err)
+	}
+
+	tx, err := pool.Begin(ctx)
+	if err != nil {
+		t.Fatalf("begin explain tx: %v", err)
+	}
+	defer func() { _ = tx.Rollback(ctx) }()
+	if _, err := tx.Exec(ctx, "SET LOCAL enable_seqscan = off"); err != nil {
+		t.Fatalf("disable seqscan: %v", err)
+	}
+
+	historyPlan := explainText(t, ctx, tx, `
+SELECT box_id
+FROM message_boxes
+WHERE owner_user_id = $1
+  AND peer_type = 'user'
+  AND peer_id = $2
+  AND NOT deleted
+  AND box_id < $3
+ORDER BY box_id DESC
+LIMIT 20
+`, recipient.ID, sender.ID, 100000)
+	requirePlanUsesPartitionIndex(t, historyPlan, "message_boxes")
+	requirePlanNotContains(t, historyPlan, "Append")
+
+	updatesPlan := explainText(t, ctx, tx, `
+SELECT pts
+FROM user_update_events
+WHERE user_id = $1
+  AND pts > $2
+ORDER BY pts ASC
+LIMIT 100
+`, recipient.ID, 0)
+	requirePlanUsesPartitionIndex(t, updatesPlan, "user_update_events")
+	requirePlanNotContains(t, updatesPlan, "Append")
+
+	dispatchPlan := explainText(t, ctx, tx, `
+WITH picked AS (
+  SELECT target_user_id, id
+  FROM dispatch_outbox
+  WHERE (
+      status = 'pending'
+      AND next_attempt_at <= now()
+    )
+    OR (
+      status = 'dispatching'
+      AND updated_at < now() - interval '30 seconds'
+    )
+  ORDER BY next_attempt_at ASC, target_user_id ASC, id ASC
+  LIMIT 100
+  FOR UPDATE SKIP LOCKED
+)
+SELECT target_user_id, id
+FROM picked
+`)
+	requirePlanContains(t, dispatchPlan, "dispatch_outbox_p")
+	requirePlanContains(t, dispatchPlan, "Index")
+	requirePlanNotContains(t, dispatchPlan, "Seq Scan")
+}
+
+func explainText(t *testing.T, ctx context.Context, tx pgx.Tx, query string, args ...any) string {
+	t.Helper()
+	rows, err := tx.Query(ctx, "EXPLAIN (COSTS OFF) "+query, args...)
+	if err != nil {
+		t.Fatalf("explain query: %v", err)
+	}
+	defer rows.Close()
+
+	var b strings.Builder
+	for rows.Next() {
+		var line string
+		if err := rows.Scan(&line); err != nil {
+			t.Fatalf("scan explain row: %v", err)
+		}
+		b.WriteString(line)
+		b.WriteByte('\n')
+	}
+	if err := rows.Err(); err != nil {
+		t.Fatalf("read explain rows: %v", err)
+	}
+	return b.String()
+}
+
+func requirePlanUsesPartitionIndex(t *testing.T, plan, table string) {
+	t.Helper()
+	requirePlanContains(t, plan, table+"_p")
+	requirePlanContains(t, plan, "Index")
+	requirePlanNotContains(t, plan, "Seq Scan")
+}
+
+func requirePlanContains(t *testing.T, plan string, needle string) {
+	t.Helper()
+	if !strings.Contains(plan, needle) {
+		t.Fatalf("plan missing %q:\n%s", needle, plan)
+	}
+}
+
+func requirePlanNotContains(t *testing.T, plan string, needle string) {
+	t.Helper()
+	if strings.Contains(plan, needle) {
+		t.Fatalf("plan contains %q:\n%s", needle, plan)
+	}
+}