瀏覽代碼

fix(executor): reject composite primary keys with SQLSTATE 0A000

A table-level PRIMARY KEY with more than one column used to silently keep
only the first column, merging rows that differ in the remaining key
columns. CREATE TABLE now rejects the definition before writing any
schema, and the executor marks the error as an unsupported feature so the
PostgreSQL wire server returns SQLSTATE 0A000 on both protocol paths.

Single-column table-level and inline primary keys are unchanged.
Composite keys are still not implemented.
Danilo Fragoso 17 小時之前
父節點
當前提交
8693f5b50c

+ 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")

+ 10 - 0
pkg/pgserver/connection.go

@@ -331,6 +331,11 @@ func (c *Connection) handleQuery(msg *Message) error {
 				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
 			}
@@ -506,6 +511,11 @@ func (c *Connection) handleExecute(msg *Message) error {
 			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
 		}

+ 15 - 0
pkg/pgserver/features_test.go

@@ -255,3 +255,18 @@ func TestExtendedQueryUniqueViolationErrorCode(t *testing.T) {
 		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)
+	}
+}