setup integration test for postgres

2026-10-14 10:08:35 +02:00 · 2025-09-02 10:51:37 -07:00
1 parent 8b4914be55
commit c9e093194d
@@ -84,19 +84,24 @@ func (s *PostgreSQLServer) handleSimpleQuery(session *PostgreSQLSession, query s
 		}
 	}

+	// Clean query by removing trailing semicolons and whitespace early
+	cleanQuery := strings.TrimSpace(query)
+	cleanQuery = strings.TrimSuffix(cleanQuery, ";")
+	cleanQuery = strings.TrimSpace(cleanQuery)
+
 	// Set database context in SQL engine if session database is different from current
 	if session.database != "" && session.database != s.sqlEngine.GetCatalog().GetCurrentDatabase() {
 		s.sqlEngine.GetCatalog().SetCurrentDatabase(session.database)
 	}

-	// Handle PostgreSQL-specific system queries
-	if postgresQuery := s.translatePostgreSQLSystemQuery(query); postgresQuery != "" {
-		query = postgresQuery
+	// Handle PostgreSQL-specific system queries directly
+	if systemResult := s.handleSystemQuery(session, cleanQuery); systemResult != nil {
+		return s.sendSystemQueryResult(session, systemResult, cleanQuery)
 	}

 	// Execute using SQL engine directly
 	ctx := context.Background()
-	result, err := s.sqlEngine.ExecuteSQL(ctx, query)
+	result, err := s.sqlEngine.ExecuteSQL(ctx, cleanQuery)
 	if err != nil {
 		return s.sendError(session, "42000", err.Error())
 	}
@@ -133,9 +138,14 @@ func (s *PostgreSQLServer) handleSimpleQuery(session *PostgreSQLSession, query s
 	return s.sendReadyForQuery(session)
 }

-// translatePostgreSQLSystemQuery translates essential PostgreSQL system queries
-// Only handles queries that PostgreSQL clients expect but SeaweedFS SQL engine doesn't natively support
-func (s *PostgreSQLServer) translatePostgreSQLSystemQuery(query string) string {
+// SystemQueryResult represents the result of a system query
+type SystemQueryResult struct {
+	Columns []string
+	Rows    [][]string
+}
+
+// handleSystemQuery handles PostgreSQL system queries directly
+func (s *PostgreSQLServer) handleSystemQuery(session *PostgreSQLSession, query string) *SystemQueryResult {
 	// Trim and normalize query
 	query = strings.TrimSpace(query)
 	query = strings.TrimSuffix(query, ";")
@@ -144,44 +154,109 @@ func (s *PostgreSQLServer) translatePostgreSQLSystemQuery(query string) string {
 	// Handle essential PostgreSQL system queries
 	switch queryLower {
 	case "select version()":
-		return "SELECT 'SeaweedFS 1.0 (PostgreSQL 14.0 compatible)' as version"
+		return &SystemQueryResult{
+			Columns: []string{"version"},
+			Rows:    [][]string{{"SeaweedFS 1.0 (PostgreSQL 14.0 compatible)"}},
+		}
 	case "select current_database()":
-		return "SELECT '" + s.config.Database + "' as current_database"
+		return &SystemQueryResult{
+			Columns: []string{"current_database"},
+			Rows:    [][]string{{s.config.Database}},
+		}
 	case "select current_user":
-		return "SELECT 'seaweedfs' as current_user"
+		return &SystemQueryResult{
+			Columns: []string{"current_user"},
+			Rows:    [][]string{{"seaweedfs"}},
+		}
 	case "select current_setting('server_version')":
-		return "SELECT '14.0' as server_version"
+		return &SystemQueryResult{
+			Columns: []string{"server_version"},
+			Rows:    [][]string{{"14.0"}},
+		}
 	case "select current_setting('server_encoding')":
-		return "SELECT 'UTF8' as server_encoding"
+		return &SystemQueryResult{
+			Columns: []string{"server_encoding"},
+			Rows:    [][]string{{"UTF8"}},
+		}
 	case "select current_setting('client_encoding')":
-		return "SELECT 'UTF8' as client_encoding"
-	}
-
-	// Handle pg_* catalog queries by mapping to equivalent SHOW commands
-	if strings.Contains(queryLower, "pg_tables") || strings.Contains(queryLower, "information_schema.tables") {
-		return "SHOW TABLES"
-	}
-	if strings.Contains(queryLower, "pg_database") || strings.Contains(queryLower, "information_schema.schemata") {
-		return "SHOW DATABASES"
+		return &SystemQueryResult{
+			Columns: []string{"client_encoding"},
+			Rows:    [][]string{{"UTF8"}},
+		}
 	}

 	// Handle transaction commands (no-op for read-only)
 	switch queryLower {
 	case "begin", "start transaction":
-		return "SELECT 'BEGIN' as status"
+		return &SystemQueryResult{
+			Columns: []string{"status"},
+			Rows:    [][]string{{"BEGIN"}},
+		}
 	case "commit":
-		return "SELECT 'COMMIT' as status"
+		return &SystemQueryResult{
+			Columns: []string{"status"},
+			Rows:    [][]string{{"COMMIT"}},
+		}
 	case "rollback":
-		return "SELECT 'ROLLBACK' as status"
+		return &SystemQueryResult{
+			Columns: []string{"status"},
+			Rows:    [][]string{{"ROLLBACK"}},
+		}
 	}

 	// If starts with SET, return a no-op
 	if strings.HasPrefix(queryLower, "set ") {
-		return "SELECT 'SET' as status"
+		return &SystemQueryResult{
+			Columns: []string{"status"},
+			Rows:    [][]string{{"SET"}},
+		}
 	}

-	// Return empty string to use original query (let SQL engine handle it)
-	return ""
+	// Return nil to use SQL engine
+	return nil
+}
+
+// sendSystemQueryResult sends the result of a system query
+func (s *PostgreSQLServer) sendSystemQueryResult(session *PostgreSQLSession, result *SystemQueryResult, query string) error {
+	// Create column descriptions for system query results
+	columns := make([]string, len(result.Columns))
+	for i, col := range result.Columns {
+		columns[i] = col
+	}
+
+	// Convert to sqltypes.Value format
+	var sqlRows [][]sqltypes.Value
+	for _, row := range result.Rows {
+		sqlRow := make([]sqltypes.Value, len(row))
+		for i, cell := range row {
+			sqlRow[i] = sqltypes.NewVarChar(cell)
+		}
+		sqlRows = append(sqlRows, sqlRow)
+	}
+
+	// Send row description
+	err := s.sendRowDescription(session, columns, sqlRows)
+	if err != nil {
+		return err
+	}
+
+	// Send data rows
+	for _, row := range sqlRows {
+		err = s.sendDataRow(session, row)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Send command complete
+	tag := s.getCommandTag(query, len(result.Rows))
+	err = s.sendCommandComplete(session, tag)
+	if err != nil {
+		return err
+	}
+
+	// Send ready for query
+	return s.sendReadyForQuery(session)
 }

 // handleParse processes a Parse message (prepared statement)
@@ -316,7 +391,7 @@ func (s *PostgreSQLServer) sendParameterStatus(session *PostgreSQLSession, name,

 // sendBackendKeyData sends backend key data
 func (s *PostgreSQLServer) sendBackendKeyData(session *PostgreSQLSession) error {
-	msg := make([]byte, 12)
+	msg := make([]byte, 13)
 	msg[0] = PG_RESP_BACKEND_KEY
 	binary.BigEndian.PutUint32(msg[1:5], 12)
 	binary.BigEndian.PutUint32(msg[5:9], session.processID)
@@ -331,7 +406,7 @@ func (s *PostgreSQLServer) sendBackendKeyData(session *PostgreSQLSession) error

 // sendReadyForQuery sends ready for query message
 func (s *PostgreSQLServer) sendReadyForQuery(session *PostgreSQLSession) error {
-	msg := make([]byte, 5)
+	msg := make([]byte, 6)
 	msg[0] = PG_RESP_READY
 	binary.BigEndian.PutUint32(msg[1:5], 5)
 	msg[5] = session.transactionState
@@ -19,6 +19,11 @@ import (

 // PostgreSQL protocol constants
 const (
+	// Protocol versions
+	PG_PROTOCOL_VERSION_3 = 196608   // PostgreSQL 3.0 protocol (0x00030000)
+	PG_SSL_REQUEST        = 80877103 // SSL request (0x04d2162f)
+	PG_GSSAPI_REQUEST     = 80877104 // GSSAPI request (0x04d21630)
+
 	// Message types from client
 	PG_MSG_STARTUP   = 0x00
 	PG_MSG_QUERY     = 'Q'
@@ -348,39 +353,68 @@ func (s *PostgreSQLServer) handleConnection(conn net.Conn) {

 // handleStartup processes the PostgreSQL startup sequence
 func (s *PostgreSQLServer) handleStartup(session *PostgreSQLSession) error {
-	// Read startup message
-	length := make([]byte, 4)
-	_, err := io.ReadFull(session.reader, length)
-	if err != nil {
-		return err
-	}
-
-	msgLength := binary.BigEndian.Uint32(length) - 4
-	msg := make([]byte, msgLength)
-	_, err = io.ReadFull(session.reader, msg)
-	if err != nil {
-		return err
-	}
-
-	// Parse startup message
-	protocolVersion := binary.BigEndian.Uint32(msg[0:4])
-	if protocolVersion != 196608 { // PostgreSQL protocol version 3.0
-		return fmt.Errorf("unsupported protocol version: %d", protocolVersion)
-	}
-
-	// Parse parameters
-	params := strings.Split(string(msg[4:]), "\x00")
-	for i := 0; i < len(params)-1; i += 2 {
-		if params[i] == "user" {
-			session.username = params[i+1]
-		} else if params[i] == "database" {
-			session.database = params[i+1]
+	for {
+		// Read startup message
+		length := make([]byte, 4)
+		_, err := io.ReadFull(session.reader, length)
+		if err != nil {
+			return err
 		}
-		session.parameters[params[i]] = params[i+1]
+
+		msgLength := binary.BigEndian.Uint32(length) - 4
+		msg := make([]byte, msgLength)
+		_, err = io.ReadFull(session.reader, msg)
+		if err != nil {
+			return err
+		}
+
+		// Parse protocol version
+		protocolVersion := binary.BigEndian.Uint32(msg[0:4])
+
+		switch protocolVersion {
+		case PG_SSL_REQUEST:
+			// Reject SSL request - send 'N' to indicate SSL not supported
+			_, err = session.conn.Write([]byte{'N'})
+			if err != nil {
+				return fmt.Errorf("failed to reject SSL request: %v", err)
+			}
+			// Continue loop to read the actual startup message
+			continue
+
+		case PG_GSSAPI_REQUEST:
+			// Reject GSSAPI request - send 'N' to indicate GSSAPI not supported
+			_, err = session.conn.Write([]byte{'N'})
+			if err != nil {
+				return fmt.Errorf("failed to reject GSSAPI request: %v", err)
+			}
+			// Continue loop to read the actual startup message
+			continue
+
+		case PG_PROTOCOL_VERSION_3:
+			// This is the actual startup message, break out of loop
+			break
+
+		default:
+			return fmt.Errorf("unsupported protocol version: %d", protocolVersion)
+		}
+
+		// Parse parameters
+		params := strings.Split(string(msg[4:]), "\x00")
+		for i := 0; i < len(params)-1; i += 2 {
+			if params[i] == "user" {
+				session.username = params[i+1]
+			} else if params[i] == "database" {
+				session.database = params[i+1]
+			}
+			session.parameters[params[i]] = params[i+1]
+		}
+
+		// Break out of the main loop - we have the startup message
+		break
 	}

 	// Handle authentication
-	err = s.handleAuthentication(session)
+	err := s.handleAuthentication(session)
 	if err != nil {
 		return err
 	}
@@ -439,7 +473,7 @@ func (s *PostgreSQLServer) handleAuthentication(session *PostgreSQLSession) erro

 // sendAuthenticationOk sends authentication OK message
 func (s *PostgreSQLServer) sendAuthenticationOk(session *PostgreSQLSession) error {
-	msg := make([]byte, 8)
+	msg := make([]byte, 9)
 	msg[0] = PG_RESP_AUTH_OK
 	binary.BigEndian.PutUint32(msg[1:5], 8)
 	binary.BigEndian.PutUint32(msg[5:9], AUTH_OK)
@@ -454,7 +488,7 @@ func (s *PostgreSQLServer) sendAuthenticationOk(session *PostgreSQLSession) erro
 // handlePasswordAuth handles clear password authentication
 func (s *PostgreSQLServer) handlePasswordAuth(session *PostgreSQLSession) error {
 	// Send password request
-	msg := make([]byte, 8)
+	msg := make([]byte, 9)
 	msg[0] = PG_RESP_AUTH_OK
 	binary.BigEndian.PutUint32(msg[1:5], 8)
 	binary.BigEndian.PutUint32(msg[5:9], AUTH_CLEAR)
@@ -511,7 +545,7 @@ func (s *PostgreSQLServer) handleMD5Auth(session *PostgreSQLSession) error {
 	}

 	// Send MD5 request
-	msg := make([]byte, 12)
+	msg := make([]byte, 13)
 	msg[0] = PG_RESP_AUTH_OK
 	binary.BigEndian.PutUint32(msg[1:5], 12)
 	binary.BigEndian.PutUint32(msg[5:9], AUTH_MD5)