chore: import upstream snapshot with attribution

2026-07-13 12:32:25 +08:00
commit e014feafe1
2285 changed files with 1131979 additions and 0 deletions
@@ -0,0 +1,415 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package harness
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"os/exec"
+	"os/signal"
+	"path/filepath"
+	"sync"
+	"syscall"
+	"time"
+
+	"golang.org/x/sync/errgroup"
+
+	"github.com/dolthub/sqllogictest/go/logictest"
+	_ "github.com/jackc/pgx/v4/stdlib"
+)
+
+const (
+	doltgresDBDir  = "doltgresDatabases"
+	serverLogFile  = "server.log"
+	harnessLogFile = "harness.log"
+	DefaultPort    = 5432
+	defaultDbName  = "sqllogictest"
+)
+
+func noDatabaseDSN(port int) string {
+	return fmt.Sprintf("postgresql://postgres:password@127.0.0.1:%d/?sslmode=disable", port)
+}
+
+func withDatabaseDSN(port int, dbName string) string {
+	return fmt.Sprintf("postgresql://postgres:password@127.0.0.1:%d/%s?sslmode=disable", port, dbName)
+}
+
+var _ logictest.Harness = &DoltgresHarness{}
+
+// DoltgresHarness is a sqllogictest harness for doltgres databases.
+type DoltgresHarness struct {
+	db               *sql.DB
+	doltgresExec     string
+	server           *DoltgresServer
+	serverDir        string
+	timeout          int64
+	port             int
+	dbName           string
+	harnessLog       *os.File
+	stashedLogOutput io.Writer
+}
+
+// NewDoltgresHarness returns a harness that manages its own server lifecycle. The server is
+// started once at construction; Init resets the database without restarting the process.
+func NewDoltgresHarness(doltgresExec string, t int64) *DoltgresHarness {
+	cwd, err := os.Getwd()
+	if err != nil {
+		logErr(err, "getting cwd")
+	}
+	serverDir := filepath.Join(cwd, doltgresDBDir)
+	err = os.RemoveAll(serverDir)
+	if err != nil {
+		logErr(err, fmt.Sprintf("running `RemoveAll` for '%s'", serverDir))
+	}
+	err = os.MkdirAll(serverDir, os.ModePerm)
+	if err != nil {
+		logErr(err, fmt.Sprintf("running `MkdirAll` for '%s'", serverDir))
+	}
+	hl, err := os.OpenFile(harnessLogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
+	if err != nil {
+		log.Fatal(err)
+	}
+	stashLogOutput := log.Writer()
+	log.SetOutput(hl)
+	logMsg("creating a new DoltgresHarness")
+
+	h := &DoltgresHarness{
+		doltgresExec:     doltgresExec,
+		serverDir:        serverDir,
+		timeout:          t,
+		port:             DefaultPort,
+		dbName:           defaultDbName,
+		harnessLog:       hl,
+		stashedLogOutput: stashLogOutput,
+	}
+	h.server = startServerProcess(doltgresExec, serverDir)
+	return h
+}
+
+// NewDoltgresWorkerHarness returns a harness for use in a concurrent worker. It connects to an
+// already-running server on the given port using dbName as its isolated database. The caller is
+// responsible for starting and stopping the server (see StartSharedServer).
+func NewDoltgresWorkerHarness(port int, dbName string, t int64) *DoltgresHarness {
+	return &DoltgresHarness{
+		timeout: t,
+		port:    port,
+		dbName:  dbName,
+	}
+}
+
+// StartSharedServer starts a single doltgres server process for use by multiple concurrent
+// workers. The returned *DoltgresServer must be closed when testing is complete.
+func StartSharedServer(doltgresExec string) *DoltgresServer {
+	cwd, err := os.Getwd()
+	if err != nil {
+		logErr(err, "getting cwd")
+	}
+	serverDir := filepath.Join(cwd, doltgresDBDir)
+	err = os.RemoveAll(serverDir)
+	if err != nil {
+		logErr(err, fmt.Sprintf("running `RemoveAll` for '%s'", serverDir))
+	}
+	err = os.MkdirAll(serverDir, os.ModePerm)
+	if err != nil {
+		logErr(err, fmt.Sprintf("running `MkdirAll` for '%s'", serverDir))
+	}
+	return startServerProcess(doltgresExec, serverDir)
+}
+
+// startServerProcess starts a doltgres subprocess in serverDir and waits for it to be ready.
+func startServerProcess(doltgresExec string, serverDir string) *DoltgresServer {
+	withKeyCtx, cancel := context.WithCancel(context.Background())
+	gServer, serverCtx := errgroup.WithContext(withKeyCtx)
+
+	server := exec.CommandContext(serverCtx, doltgresExec, "--data-dir=.")
+	server.Dir = serverDir
+
+	l, err := os.OpenFile(serverLogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
+	if err != nil {
+		logErr(err, fmt.Sprintf("opening %s file", serverLogFile))
+	}
+	server.Stdout = l
+	server.Stderr = l
+
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		<-quit
+		defer wg.Done()
+		signal.Stop(quit)
+		cancel()
+	}()
+
+	ds := &DoltgresServer{
+		dir:       serverDir,
+		quit:      quit,
+		wg:        &wg,
+		gServer:   gServer,
+		server:    server,
+		serverLog: l,
+	}
+	ds.Start()
+	return ds
+}
+
+func (h *DoltgresHarness) EngineStr() string {
+	return "postgresql"
+}
+
+// Init resets the harness to a clean state for the next test file.
+// For worker harnesses (no server process), it keeps a persistent connection and resets the
+// public schema via DROP+CREATE — avoiding database-level connection races entirely.
+// For server-owning harnesses, it drops and recreates the database (safe because only one worker
+// uses each database in the non-concurrent path).
+func (h *DoltgresHarness) Init() error {
+	if h.server == nil {
+		return h.initWorker()
+	}
+	return h.initServer()
+}
+
+// initWorker resets a worker harness between test files.
+//
+// When useDropSchema is true (PostgreSQL), it wipes and recreates the public schema on the
+// persistent connection — no reconnect needed, no races.
+func (h *DoltgresHarness) initWorker() error {
+	ctx := context.Background()
+
+	if h.db != nil {
+		h.db.Close()
+		h.db = nil
+	}
+	noDB, err := sql.Open("pgx", noDatabaseDSN(h.port))
+	if err != nil {
+		return err
+	}
+	_, err = noDB.ExecContext(ctx, "DROP DATABASE IF EXISTS "+h.dbName)
+	if err == nil {
+		_, err = noDB.ExecContext(ctx, "CREATE DATABASE "+h.dbName)
+	}
+	noDB.Close()
+	if err != nil {
+		return fmt.Errorf("resetting database %s: %w", h.dbName, err)
+	}
+	return h.openWorkerDB(ctx)
+}
+
+// openWorkerDB establishes a connection to the worker's database, creating it if necessary.
+func (h *DoltgresHarness) openWorkerDB(ctx context.Context) error {
+	db, err := sql.Open("pgx", withDatabaseDSN(h.port, h.dbName))
+	if err != nil {
+		return err
+	}
+	if pingErr := db.Ping(); pingErr == nil {
+		h.db = db
+		return nil
+	}
+	db.Close()
+
+	// Database does not exist yet; create it.
+	noDB, err := sql.Open("pgx", noDatabaseDSN(h.port))
+	if err != nil {
+		return err
+	}
+	_, createErr := noDB.ExecContext(ctx, "CREATE DATABASE "+h.dbName)
+	noDB.Close()
+	if createErr != nil {
+		logErr(createErr, "creating database "+h.dbName)
+		return createErr
+	}
+
+	db, err = sql.Open("pgx", withDatabaseDSN(h.port, h.dbName))
+	if err != nil {
+		return err
+	}
+	if err = db.Ping(); err != nil {
+		db.Close()
+		return err
+	}
+	h.db = db
+	return nil
+}
+
+// initServer resets the server-owning harness by dropping and recreating the database.
+// This is only used in the non-concurrent (single-worker) path.
+func (h *DoltgresHarness) initServer() error {
+	ctx := context.Background()
+
+	if h.db != nil {
+		h.db.Close()
+		h.db = nil
+	}
+
+	noDB, openErr := sql.Open("pgx", noDatabaseDSN(h.port))
+	if openErr != nil {
+		logErr(openErr, "opening no-db connection")
+		return openErr
+	}
+	_, err := noDB.ExecContext(ctx, "DROP DATABASE IF EXISTS "+h.dbName)
+	if err == nil {
+		_, err = noDB.ExecContext(ctx, "CREATE DATABASE "+h.dbName)
+	}
+	noDB.Close()
+	if err != nil {
+		logErr(err, "resetting database "+h.dbName)
+		return err
+	}
+
+	db, err := sql.Open("pgx", withDatabaseDSN(h.port, h.dbName))
+	if err != nil {
+		logErr(err, "opening connection to "+h.dbName)
+		return err
+	}
+	if err = db.Ping(); err != nil {
+		db.Close()
+		return err
+	}
+	h.db = db
+	return nil
+}
+
+func (h *DoltgresHarness) Close() error {
+	if h.db != nil {
+		h.db.Close()
+		h.db = nil
+	}
+	if h.server != nil {
+		h.server.Close()
+		h.server = nil
+	}
+	if h.harnessLog != nil {
+		h.harnessLog.Close()
+		log.SetOutput(h.stashedLogOutput)
+	}
+	return os.RemoveAll(h.serverDir)
+}
+
+func (h *DoltgresHarness) ExecuteStatement(statement string) error {
+	_, err := h.db.Exec(statement)
+	return err
+}
+
+func (h *DoltgresHarness) ExecuteQuery(statement string) (schema string, results []string, err error) {
+	rows, err := h.db.Query(statement)
+	if rows != nil {
+		defer rows.Close()
+	}
+	if err != nil {
+		return "", nil, err
+	}
+	return h.getSchemaAndResults(rows)
+}
+
+func (h *DoltgresHarness) ExecuteStatementContext(ctx context.Context, statement string) error {
+	_, err := h.db.ExecContext(ctx, statement)
+	return err
+}
+
+func (h *DoltgresHarness) ExecuteQueryContext(ctx context.Context, statement string) (schema string, results []string, err error) {
+	rows, err := h.db.QueryContext(ctx, statement)
+	if rows != nil {
+		defer rows.Close()
+	}
+	if err != nil {
+		return "", nil, err
+	}
+	return h.getSchemaAndResults(rows)
+}
+
+func (h *DoltgresHarness) GetTimeout() int64 {
+	return h.timeout
+}
+
+func (h *DoltgresHarness) getSchemaAndResults(rows *sql.Rows) (schema string, results []string, err error) {
+	schema, cols, err := columns(rows)
+	if err != nil {
+		return "", nil, err
+	}
+	for rows.Next() {
+		if err = rows.Scan(cols...); err != nil {
+			return "", nil, err
+		}
+		for _, col := range cols {
+			results = append(results, stringVal(col))
+		}
+	}
+	if rows.Err() != nil {
+		return "", nil, rows.Err()
+	}
+	return schema, results, nil
+}
+
+type DoltgresServer struct {
+	dir       string
+	quit      chan os.Signal
+	wg        *sync.WaitGroup
+	gServer   *errgroup.Group
+	server    *exec.Cmd
+	serverLog *os.File
+}
+
+func (s *DoltgresServer) Start() {
+	logMsg("starting doltgres server")
+	var err error
+	s.gServer.Go(func() error {
+		err = s.server.Run()
+		return err
+	})
+	// Allow the server time to start accepting connections.
+	time.Sleep(3 * time.Second)
+	if err != nil {
+		logErr(err, "from server.Start()")
+	}
+}
+
+func (s *DoltgresServer) Stop() {
+	select {
+	case <-s.quit:
+		return
+	default:
+	}
+	s.quit <- syscall.SIGTERM
+	err := s.gServer.Wait()
+	if err != nil {
+		if err.Error() == "signal: killed" {
+			logMsg("doltgres server stopped successfully")
+		} else {
+			logErr(err, "from server.Stop()")
+		}
+	}
+	close(s.quit)
+	s.wg.Wait()
+}
+
+func (s *DoltgresServer) Close() {
+	s.Stop()
+	if err := s.serverLog.Close(); err != nil {
+		logErr(err, "closing server.log")
+	}
+}
+
+func logErr(err error, cause string) {
+	log.Printf("ERROR: %s received from %s", err.Error(), cause)
+}
+
+func logMsg(msg string) {
+	log.Println(msg)
+}