Răsfoiți Sursa

Merge pull request #1 from database-pizza/fix/unique-23505-and-composite-pk

fix: SQLSTATE 23505 for UNIQUE violations and reject composite primary keys
Danilo Fragoso 20 ore în urmă
părinte
comite
ba76c86725

+ 72 - 0
pkg/executor/composite_primary_key_test.go

@@ -0,0 +1,72 @@
+package executor
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestCreateTablePrimaryKeyConstraints(t *testing.T) {
+	tests := []struct {
+		name           string
+		sql            string
+		wantErr        string
+		wantPrimaryKey string
+	}{
+		{
+			name:    "composite table-level primary key is rejected",
+			sql:     "CREATE TABLE t (a TEXT, b TEXT, PRIMARY KEY (a, b))",
+			wantErr: "composite primary key",
+		},
+		{
+			name:    "duplicate column composite primary key is rejected",
+			sql:     "CREATE TABLE t (a TEXT, PRIMARY KEY (a, a))",
+			wantErr: "composite primary key",
+		},
+		{
+			name:           "single-column table-level primary key is accepted",
+			sql:            "CREATE TABLE t (a TEXT, PRIMARY KEY (a))",
+			wantPrimaryKey: "a",
+		},
+		{
+			name:           "inline primary key is accepted",
+			sql:            "CREATE TABLE t (a TEXT PRIMARY KEY)",
+			wantPrimaryKey: "a",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, schema, table := newTestDB(t)
+			e := newExec(schema, table)
+
+			_, err := execSQL(e, tt.sql)
+			if tt.wantErr != "" {
+				if err == nil {
+					t.Fatalf("expected %q to fail", tt.sql)
+				}
+				if !strings.Contains(err.Error(), tt.wantErr) {
+					t.Fatalf("error = %q, want it to contain %q", err, tt.wantErr)
+				}
+				if schema.TableExists("t") {
+					t.Fatal("rejected CREATE TABLE left a schema behind")
+				}
+				if _, err := execSQL(e, "SELECT * FROM t"); err == nil {
+					t.Fatal("rejected CREATE TABLE left a usable table behind")
+				}
+				return
+			}
+
+			if err != nil {
+				t.Fatalf("exec %q: %v", tt.sql, err)
+			}
+			got, err := schema.GetSchema("t")
+			if err != nil {
+				t.Fatalf("get schema: %v", err)
+			}
+			if got.PrimaryKey != tt.wantPrimaryKey {
+				t.Fatalf("PrimaryKey = %q, want %q", got.PrimaryKey, tt.wantPrimaryKey)
+			}
+			execMust(t, e, "INSERT INTO t (a) VALUES ('x')")
+		})
+	}
+}

+ 8 - 0
pkg/executor/executor.go

@@ -17,6 +17,11 @@ import (
 	"github.com/danfragoso/pizzasql-next/pkg/version"
 )
 
+// ErrUnsupportedFeature marks DDL that PizzaSQL intentionally rejects instead
+// of silently accepting with different semantics. Protocol layers map it to a
+// feature-not-supported SQLSTATE rather than a generic internal error.
+var ErrUnsupportedFeature = errors.New("unsupported feature")
+
 // Executor executes SQL statements.
 type Executor struct {
 	schema   *storage.SchemaManager
@@ -3565,6 +3570,9 @@ func (e *Executor) executeCreateTable(stmt *parser.CreateTableStmt) (*Result, er
 	// Handle table-level constraints
 	for _, constraint := range stmt.Constraints {
 		if constraint.Type == parser.ConstraintPrimaryKey && len(constraint.Columns) > 0 {
+			if len(constraint.Columns) > 1 {
+				return nil, fmt.Errorf("%w: composite primary keys are not supported; declare a single-column PRIMARY KEY", ErrUnsupportedFeature)
+			}
 			schema.PrimaryKey = constraint.Columns[0]
 			for i := range schema.Columns {
 				if strings.EqualFold(schema.Columns[i].Name, schema.PrimaryKey) {

+ 1 - 1
pkg/executor/session_state_test.go

@@ -248,7 +248,7 @@ func TestPrimaryKeyMultiRowInsertAtomic(t *testing.T) {
 func TestUpdateWithScalarSubqueryInTransaction(t *testing.T) {
 	_, schema, table := newTestDB(t)
 	e := newExec(schema, table)
-	execMust(t, e, "CREATE TABLE watch (user_id INTEGER, repo_id INTEGER, PRIMARY KEY(user_id, repo_id))")
+	execMust(t, e, "CREATE TABLE watch (user_id INTEGER, repo_id INTEGER, PRIMARY KEY(user_id))")
 	execMust(t, e, "CREATE TABLE repository (id INTEGER PRIMARY KEY, num_watches INTEGER)")
 	execMust(t, e, "INSERT INTO repository VALUES (1, 0)")
 	execMust(t, e, "BEGIN")

+ 28 - 1
pkg/pgserver/connection.go

@@ -320,13 +320,26 @@ func (c *Connection) handleQuery(msg *Message) error {
 		result, err := c.executor.Execute(stmt)
 		if err != nil {
 			code := ErrCodeInternalError
+			message := fmt.Sprintf("Execution error: %v", err)
+			var uniqueErr *storage.UniqueViolationError
 			if errors.Is(err, storage.ErrSerialization) {
 				code = ErrCodeSerializationFailure
 				c.txStatus = TxStatusIdle
+			} else if errors.As(err, &uniqueErr) {
+				code = ErrCodeUniqueViolation
+				message = fmt.Sprintf("Execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint)
+				if c.txStatus == TxStatusInBlock {
+					c.txStatus = TxStatusFailed
+				}
+			} else if errors.Is(err, executor.ErrUnsupportedFeature) {
+				code = ErrCodeFeatureNotSupported
+				if c.txStatus == TxStatusInBlock {
+					c.txStatus = TxStatusFailed
+				}
 			} else if c.txStatus == TxStatusInBlock {
 				c.txStatus = TxStatusFailed
 			}
-			c.sendError("ERROR", code, fmt.Sprintf("Execution error: %v", err))
+			c.sendError("ERROR", code, message)
 			return c.sendReadyForQuery()
 		}
 		if err := c.sendResult(result, stmt); err != nil {
@@ -489,12 +502,26 @@ func (c *Connection) handleExecute(msg *Message) error {
 	result, err := c.executor.Execute(stmt)
 	if err != nil {
 		code := ErrCodeInternalError
+		var uniqueErr *storage.UniqueViolationError
 		if errors.Is(err, storage.ErrSerialization) {
 			code = ErrCodeSerializationFailure
 			c.txStatus = TxStatusIdle
+		} else if errors.As(err, &uniqueErr) {
+			code = ErrCodeUniqueViolation
+			if c.txStatus == TxStatusInBlock {
+				c.txStatus = TxStatusFailed
+			}
+		} else if errors.Is(err, executor.ErrUnsupportedFeature) {
+			code = ErrCodeFeatureNotSupported
+			if c.txStatus == TxStatusInBlock {
+				c.txStatus = TxStatusFailed
+			}
 		} else if c.txStatus == TxStatusInBlock {
 			c.txStatus = TxStatusFailed
 		}
+		if uniqueErr != nil {
+			return c.failExtended(code, fmt.Errorf("execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint))
+		}
 		return c.failExtended(code, fmt.Errorf("execution error: %w", err))
 	}
 	return c.sendResult(result, stmt)

+ 114 - 0
pkg/pgserver/features_test.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"net"
+	"strings"
 	"testing"
 
 	"github.com/danfragoso/pizzasql-next/pkg/executor"
@@ -156,3 +157,116 @@ func TestGetCommandTagSelectUsesRowCount(t *testing.T) {
 		t.Fatalf("tag = %q, want SELECT 2", tag)
 	}
 }
+
+// errorMessage extracts the human-readable message from an ErrorResponse.
+func errorMessage(msg *Message) string {
+	if msg == nil || msg.Type != MsgErrorResponse {
+		return ""
+	}
+	data := msg.Data
+	for i := 0; i < len(data); {
+		if data[i] == 0 {
+			break
+		}
+		field := data[i]
+		i++
+		end := bytes.IndexByte(data[i:], 0)
+		if end < 0 {
+			break
+		}
+		value := string(data[i : i+end])
+		if field == ErrorFieldMessage {
+			return value
+		}
+		i += end + 1
+	}
+	return ""
+}
+
+func TestSimpleQueryUniqueViolationErrorCode(t *testing.T) {
+	c, client := newDataConnection(t)
+	runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("CREATE TABLE t (id INTEGER PRIMARY KEY, email TEXT UNIQUE)"), 0)})
+	runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("INSERT INTO t VALUES (1, 'a@x')"), 0)})
+
+	msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("INSERT INTO t VALUES (2, 'a@x')"), 0)})
+	if len(msgs) != 2 {
+		t.Fatalf("got %d messages, want ErrorResponse+ReadyForQuery", len(msgs))
+	}
+	if code := errorCode(msgs[0]); code != ErrCodeUniqueViolation {
+		t.Fatalf("error code = %q, want %q", code, ErrCodeUniqueViolation)
+	}
+	if message := errorMessage(msgs[0]); !strings.Contains(message, "uniq_t_email") {
+		t.Fatalf("error message = %q, want constraint uniq_t_email", message)
+	}
+}
+
+func TestExtendedQueryUniqueViolationErrorCode(t *testing.T) {
+	c, client := newDataConnection(t)
+	runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("CREATE TABLE t (id INTEGER PRIMARY KEY, email TEXT UNIQUE)"), 0)})
+	runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("INSERT INTO t VALUES (1, 'a@x')"), 0)})
+
+	parse := NewMessageBuilder()
+	parse.WriteString("")
+	parse.WriteString("INSERT INTO t VALUES (2, 'a@x')")
+	parse.WriteInt16(0)
+
+	bind := NewMessageBuilder()
+	bind.WriteString("")
+	bind.WriteString("")
+	bind.WriteInt16(0) // parameter format codes
+	bind.WriteInt16(0) // parameter values
+	bind.WriteInt16(0) // result format codes
+
+	execute := NewMessageBuilder()
+	execute.WriteString("")
+	execute.WriteInt32(0)
+
+	done := make(chan error, 1)
+	go func() {
+		if err := c.handleParse(&Message{Type: MsgParse, Data: parse.Bytes()}); err != nil {
+			done <- err
+			return
+		}
+		if err := c.handleBind(&Message{Type: MsgBind, Data: bind.Bytes()}); err != nil {
+			done <- err
+			return
+		}
+		done <- c.handleExecute(&Message{Type: MsgExecute, Data: execute.Bytes()})
+	}()
+
+	reader := bufio.NewReader(client)
+	var errMsg *Message
+	for errMsg == nil {
+		msg, err := ReadMessage(reader)
+		if err != nil {
+			t.Fatalf("read response: %v", err)
+		}
+		if msg.Type == MsgErrorResponse {
+			errMsg = msg
+		}
+	}
+	if err := <-done; err != nil {
+		t.Fatalf("extended query: %v", err)
+	}
+	if code := errorCode(errMsg); code != ErrCodeUniqueViolation {
+		t.Fatalf("error code = %q, want %q", code, ErrCodeUniqueViolation)
+	}
+	if message := errorMessage(errMsg); !strings.Contains(message, "uniq_t_email") {
+		t.Fatalf("error message = %q, want constraint uniq_t_email", message)
+	}
+}
+
+func TestCompositePrimaryKeyDDLErrorCode(t *testing.T) {
+	c, client := newDataConnection(t)
+
+	msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("CREATE TABLE t (a TEXT, b TEXT, PRIMARY KEY (a,b))"), 0)})
+	if len(msgs) != 2 {
+		t.Fatalf("got %d messages, want ErrorResponse+ReadyForQuery", len(msgs))
+	}
+	if code := errorCode(msgs[0]); code != ErrCodeFeatureNotSupported {
+		t.Fatalf("error code = %q, want %q", code, ErrCodeFeatureNotSupported)
+	}
+	if message := errorMessage(msgs[0]); !strings.Contains(message, "composite primary key") {
+		t.Fatalf("error message = %q, want a composite primary key error", message)
+	}
+}

+ 1 - 0
pkg/pgserver/protocol.go

@@ -92,6 +92,7 @@ const (
 	ErrCodeFeatureNotSupported  = "0A000"
 	ErrCodeTransactionAborted   = "25P02"
 	ErrCodeSerializationFailure = "40001"
+	ErrCodeUniqueViolation      = "23505"
 )
 
 // Message represents a PostgreSQL protocol message

+ 14 - 3
pkg/storage/table.go

@@ -14,6 +14,17 @@ import (
 // Row represents a database row.
 type Row map[string]interface{}
 
+// UniqueViolationError reports a duplicate value on a UNIQUE index. The
+// constraint name identifies the violated index so protocol layers can surface a
+// machine-detectable error.
+type UniqueViolationError struct {
+	Constraint string
+}
+
+func (e *UniqueViolationError) Error() string {
+	return fmt.Sprintf("UNIQUE constraint failed: %s", e.Constraint)
+}
+
 // TableManager manages table data operations.
 type TableManager struct {
 	pool     *KVPool
@@ -1804,7 +1815,7 @@ func (m *TableManager) validateUniqueRows(table string, pending []Row, excludedK
 				seen[idx.Name] = perIndex
 			}
 			if perIndex[encoded] {
-				return fmt.Errorf("UNIQUE constraint failed: %s", idx.Name)
+				return &UniqueViolationError{Constraint: idx.Name}
 			}
 			perIndex[encoded] = true
 		}
@@ -1842,7 +1853,7 @@ func (m *TableManager) validateUniqueRows(table string, pending []Row, excludedK
 						continue
 					}
 					if perIndex := seen[idx.Name]; perIndex != nil && perIndex[encoded] {
-						return fmt.Errorf("UNIQUE constraint failed: %s", idx.Name)
+						return &UniqueViolationError{Constraint: idx.Name}
 					}
 				}
 			}
@@ -1915,7 +1926,7 @@ func (m *TableManager) ValidateUniqueIndex(index *Index) error {
 			return false, nil
 		}
 		if seen[encoded] {
-			return true, fmt.Errorf("UNIQUE constraint failed: %s", index.Name)
+			return true, &UniqueViolationError{Constraint: index.Name}
 		}
 		seen[encoded] = true
 		return false, nil

+ 25 - 0
pkg/storage/unique_test.go

@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 	"sync"
@@ -77,6 +78,30 @@ func TestUniqueIndexEnforcesOnInsert(t *testing.T) {
 	}
 }
 
+func TestUniqueViolationErrorCarriesConstraint(t *testing.T) {
+	schemas, tables := uniqueTable(t,
+		Column{Name: "id", Type: "INTEGER", PrimaryKey: true},
+		Column{Name: "email", Type: "TEXT"},
+	)
+	if err := schemas.CreateIndex(&Index{Name: "uq_email", Table: "t", Unique: true, Columns: []IndexColumn{{Name: "email"}}}); err != nil {
+		t.Fatal(err)
+	}
+	if err := tables.Insert("t", Row{"id": int64(1), "email": "a@x"}); err != nil {
+		t.Fatal(err)
+	}
+	err := tables.Insert("t", Row{"id": int64(2), "email": "a@x"})
+	var violation *UniqueViolationError
+	if !errors.As(err, &violation) {
+		t.Fatalf("error %v is not a *UniqueViolationError", err)
+	}
+	if violation.Constraint != "uq_email" {
+		t.Fatalf("constraint = %q, want uq_email", violation.Constraint)
+	}
+	if err.Error() != "UNIQUE constraint failed: uq_email" {
+		t.Fatalf("message = %q, want UNIQUE constraint failed: uq_email", err.Error())
+	}
+}
+
 func TestUniqueIndexAllowsMultipleNulls(t *testing.T) {
 	schemas, tables := uniqueTable(t,
 		Column{Name: "id", Type: "INTEGER", PrimaryKey: true},