package pgserver import ( "bufio" "bytes" "encoding/binary" "fmt" "io" "log" "net" "regexp" "sort" "strconv" "strings" "github.com/danfragoso/pizzasql-next/pkg/executor" "github.com/danfragoso/pizzasql-next/pkg/lexer" "github.com/danfragoso/pizzasql-next/pkg/parser" "github.com/danfragoso/pizzasql-next/pkg/storage" ) // Connection represents a client connection type Connection struct { conn net.Conn reader *bufio.Reader writer *bufio.Writer executor *executor.Executor schema *storage.SchemaManager dbManager *storage.DatabaseManager database string params map[string]string txStatus byte quiet bool // Disable query logging statements map[string]*preparedStatement portals map[string]*portal extendedFailed bool } type preparedStatement struct { query string paramOIDs []int32 } type portal struct { query string executed bool } // NewConnection creates a new connection handler func NewConnection(conn net.Conn, dbManager *storage.DatabaseManager, quiet bool) *Connection { return &Connection{ conn: conn, reader: bufio.NewReader(conn), writer: bufio.NewWriter(conn), dbManager: dbManager, params: make(map[string]string), statements: make(map[string]*preparedStatement), portals: make(map[string]*portal), txStatus: TxStatusIdle, quiet: quiet, } } // Handle processes the connection func (c *Connection) Handle() error { defer func() { if c.executor != nil { if err := c.executor.RollbackActive(); err != nil && !c.quiet { log.Printf("failed to roll back disconnected transaction: %v", err) } } c.conn.Close() }() // First, check for SSL request (sent before startup message) // SSL request is 8 bytes: length(4) + code(4) where code = 80877103 firstBytes := make([]byte, 8) n, err := io.ReadFull(c.reader, firstBytes) if err != nil { return fmt.Errorf("failed to read initial bytes: %w", err) } // Check if it's an SSL request (code 80877103 = 0x04D2162F) if n == 8 { length := binary.BigEndian.Uint32(firstBytes[0:4]) code := binary.BigEndian.Uint32(firstBytes[4:8]) if length == 8 && code == 80877103 { // SSL request - we don't support SSL, send 'N' if !c.quiet { log.Printf("Client requested SSL, sending rejection") } if _, err := c.conn.Write([]byte{'N'}); err != nil { return fmt.Errorf("failed to send SSL rejection: %w", err) } // Now read the actual startup message } else { // Not SSL request, this is part of startup message // We need to prepend these bytes back for ReadStartupMessage // Create a multi-reader that first reads our buffered bytes, then continues with the reader c.reader = bufio.NewReader(io.MultiReader(bytes.NewReader(firstBytes), c.reader)) } } // Read startup message if !c.quiet { log.Printf("Reading startup message...") } params, err := ReadStartupMessage(c.reader) if err != nil { return fmt.Errorf("failed to read startup message: %w", err) } c.params = params if !c.quiet { log.Printf("Startup params: %+v", params) } // Get database name from params (default to "pizzasql") dbName := params["database"] if dbName == "" { dbName = "pizzasql" } c.database = dbName if !c.quiet { log.Printf("New connection: user=%s database=%s", params["user"], dbName) } // Initialize database if err := c.initDatabase(dbName); err != nil { c.sendError("FATAL", ErrCodeConnectionFailure, fmt.Sprintf("Failed to initialize database: %v", err)) return err } // Send authentication OK (no auth for now) if err := c.sendAuthenticationOk(); err != nil { return err } // Send parameter status messages if err := c.sendParameterStatus("server_version", "14.0 (PizzaSQL)"); err != nil { return err } if err := c.sendParameterStatus("server_encoding", "UTF8"); err != nil { return err } if err := c.sendParameterStatus("client_encoding", "UTF8"); err != nil { return err } if err := c.sendParameterStatus("DateStyle", "ISO, MDY"); err != nil { return err } if err := c.sendParameterStatus("TimeZone", "UTC"); err != nil { return err } // Send backend key data (for cancellation - we don't implement this yet) if err := c.sendBackendKeyData(12345, 67890); err != nil { return err } // Send ready for query if err := c.sendReadyForQuery(); err != nil { return err } // Message loop for { msg, err := ReadMessage(c.reader) if err != nil { if err == io.EOF { log.Printf("Connection closed by client") return nil } return fmt.Errorf("failed to read message: %w", err) } if err := c.handleMessage(msg); err != nil { if err == io.EOF { // Normal termination log.Printf("Connection closed normally") return nil } log.Printf("Error handling message: %v", err) return err } } } // initDatabase initializes the database connection func (c *Connection) initDatabase(dbName string) error { db, err := c.dbManager.GetDatabase(dbName) if err != nil { return err } c.schema = db.Schema c.executor = executor.New(db.Schema, db.Table) c.executor.SyncCatalog() return nil } // handleMessage processes a client message func (c *Connection) handleMessage(msg *Message) error { switch msg.Type { case MsgQuery: return c.handleQuery(msg) case MsgTerminate: log.Printf("Client requested termination") return io.EOF case MsgParse: return c.handleParse(msg) case MsgBind: return c.handleBind(msg) case MsgDescribe: return c.handleDescribe(msg) case MsgExecute: return c.handleExecute(msg) case MsgClose: return c.handleClose(msg) case MsgFlush: return c.writer.Flush() case MsgSync: c.extendedFailed = false return c.sendReadyForQuery() default: log.Printf("Unknown message type: %c (%d)", msg.Type, msg.Type) c.sendError("ERROR", ErrCodeProtocolViolation, fmt.Sprintf("Unknown message type: %c", msg.Type)) return c.sendReadyForQuery() } } // handleQuery processes a simple query func (c *Connection) handleQuery(msg *Message) error { // A Query message carries a single NUL-terminated query string. if len(msg.Data) == 0 || msg.Data[len(msg.Data)-1] != 0 { c.sendError("ERROR", ErrCodeProtocolViolation, "invalid Query message: missing null terminator") return c.sendReadyForQuery() } sql := string(msg.Data[:len(msg.Data)-1]) if !c.quiet { log.Printf("Query: %s", sql) } // Handle empty query sqlTrimmed := strings.TrimSpace(sql) if sqlTrimmed == "" || sqlTrimmed == ";" { if err := c.sendEmptyQueryResponse(); err != nil { return err } return c.sendReadyForQuery() } // Handle special PostgreSQL system queries that drivers send sqlUpper := strings.ToUpper(strings.TrimSpace(sql)) singleStatement := isSingleStatementSQL(sql) // In a failed transaction only ROLLBACK (or ROLLBACK TO SAVEPOINT) is // accepted. Gate before honoring the special driver queries below so they // return 25P02 instead of a result. Multi-statement batches are gated // per statement in the execution loop below. if c.txStatus == TxStatusFailed && singleStatement && !isRollbackStatement(sql) { c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block") return c.sendReadyForQuery() } // lib/pq and other drivers query these for connection validation if singleStatement && strings.Contains(sqlUpper, "SELECT VERSION()") { // Return a fake PostgreSQL version return c.handleVersionQuery() } if singleStatement && strings.Contains(sqlUpper, "SELECT CURRENT_USER") { // Return the current user return c.handleCurrentUserQuery() } if singleStatement && strings.Contains(sqlUpper, "SHOW") && (strings.Contains(sqlUpper, "SERVER_VERSION") || strings.Contains(sqlUpper, "SERVER_ENCODING") || strings.Contains(sqlUpper, "CLIENT_ENCODING")) { // Handle SHOW commands return c.handleShowCommand(sqlUpper) } if result, handled, err := c.catalogResult(sql); handled { if err != nil { c.sendError("ERROR", ErrCodeInternalError, err.Error()) return c.sendReadyForQuery() } if err := c.sendTabularResult(result, "SELECT"); err != nil { return err } return c.sendReadyForQuery() } // Parse the complete batch before executing its first statement. This keeps // unsupported trailing clauses from turning into committed writes. l := lexer.New(sql) p := parser.New(l) stmts, err := p.ParseMultiple() if err != nil { c.sendError("ERROR", ErrCodeSyntaxError, fmt.Sprintf("Syntax error: %v", err)) return c.sendReadyForQuery() } for _, stmt := range stmts { if c.txStatus == TxStatusFailed { if _, rollback := stmt.(*parser.RollbackStmt); !rollback { c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block") return c.sendReadyForQuery() } } result, err := c.executor.Execute(stmt) if err != nil { if c.txStatus == TxStatusInBlock { c.txStatus = TxStatusFailed } c.sendError("ERROR", ErrCodeInternalError, fmt.Sprintf("Execution error: %v", err)) return c.sendReadyForQuery() } if err := c.sendResult(result, stmt); err != nil { return err } } return c.sendReadyForQuery() } func (c *Connection) handleParse(msg *Message) error { if c.extendedFailed { return nil } name, pos, err := readCString(msg.Data, 0) if err != nil { return c.failExtended(ErrCodeProtocolViolation, err) } query, pos, err := readCString(msg.Data, pos) if err != nil || pos+2 > len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse message")) } count := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2])) pos += 2 if pos+count*4 != len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse parameter list")) } oids := make([]int32, count) for i := range oids { oids[i] = int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4])) pos += 4 } c.statements[name] = &preparedStatement{query: query, paramOIDs: oids} return c.writeMessage(MsgParseComplete, nil) } func (c *Connection) handleBind(msg *Message) error { if c.extendedFailed { return nil } portalName, pos, err := readCString(msg.Data, 0) if err != nil { return c.failExtended(ErrCodeProtocolViolation, err) } statementName, pos, err := readCString(msg.Data, pos) if err != nil { return c.failExtended(ErrCodeProtocolViolation, err) } statement, ok := c.statements[statementName] if !ok { return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", statementName)) } formats, pos, err := readInt16List(msg.Data, pos) if err != nil || pos+2 > len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind format list")) } paramCount := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2])) pos += 2 params := make([]boundParameter, paramCount) for i := range params { if pos+4 > len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter")) } length := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4])) pos += 4 params[i].oid = parameterOID(statement.paramOIDs, i) params[i].format = parameterFormat(formats, i) if length == -1 { params[i].null = true continue } if length < 0 || pos+int(length) > len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter length")) } params[i].value = append([]byte(nil), msg.Data[pos:pos+int(length)]...) pos += int(length) } _, pos, err = readInt16List(msg.Data, pos) // result formats; text output is currently used if err != nil || pos != len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format list")) } query, err := bindQuery(statement.query, params) if err != nil { return c.failExtended(ErrCodeInvalidParameter, err) } c.portals[portalName] = &portal{query: query} return c.writeMessage(MsgBindComplete, nil) } func (c *Connection) handleDescribe(msg *Message) error { if c.extendedFailed { return nil } if len(msg.Data) < 2 { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe message")) } name, pos, err := readCString(msg.Data, 1) if err != nil || pos != len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe name")) } switch msg.Data[0] { case 'S': statement, ok := c.statements[name] if !ok { return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", name)) } mb := NewMessageBuilder() mb.WriteInt16(int16(len(statement.paramOIDs))) for _, oid := range statement.paramOIDs { mb.WriteInt32(oid) } if err := c.writeMessage(MsgParameterDescription, mb.Bytes()); err != nil { return err } case 'P': if _, ok := c.portals[name]; !ok { return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name)) } default: return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe target")) } // Execute sends the row description once the bound statement has been // analyzed, avoiding side effects during Describe. return c.writeMessage(MsgNoData, nil) } func (c *Connection) handleExecute(msg *Message) error { if c.extendedFailed { return nil } name, pos, err := readCString(msg.Data, 0) if err != nil || pos+4 != len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute message")) } portal, ok := c.portals[name] if !ok { return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name)) } if portal.executed { return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("portal can only be executed once")) } portal.executed = true // Check the transaction state before catalog emulation. Catalog queries do // not all parse as regular PizzaSQL statements, but must still return 25P02 // while the transaction is aborted. if c.txStatus == TxStatusFailed && !isRollbackStatement(portal.query) { return c.failExtended(ErrCodeTransactionAborted, fmt.Errorf("current transaction is aborted, commands ignored until end of transaction block")) } if result, handled, catalogErr := c.catalogResult(portal.query); handled { if catalogErr != nil { return c.failExtended(ErrCodeInternalError, catalogErr) } return c.sendTabularResult(result, "SELECT") } l := lexer.New(portal.query) stmt, err := parser.New(l).Parse() if err != nil { return c.failExtended(ErrCodeSyntaxError, fmt.Errorf("syntax error: %w", err)) } result, err := c.executor.Execute(stmt) if err != nil { if c.txStatus == TxStatusInBlock { c.txStatus = TxStatusFailed } return c.failExtended(ErrCodeInternalError, fmt.Errorf("execution error: %w", err)) } return c.sendResult(result, stmt) } var catalogFilterPattern = regexp.MustCompile(`(?i)\b(table_name|tablename|table_schema|schemaname|constraint_name|indexname)\s*=\s*'((?:''|[^'])*)'`) func (c *Connection) catalogResult(sql string) (*executor.Result, bool, error) { if !isSingleStatementSQL(sql) { return nil, false, nil } upper := strings.ToUpper(sql) var source string for _, candidate := range []string{ "INFORMATION_SCHEMA.TABLES", "INFORMATION_SCHEMA.COLUMNS", "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE", "PG_TABLES", "PG_INDEXES", } { if strings.Contains(upper, "FROM "+candidate) { source = candidate break } } if source == "" { return nil, false, nil } if c.txStatus == TxStatusIdle { c.schema.LockStatement() defer c.schema.UnlockStatement() } tables, err := c.schema.ListTables() if err != nil { return nil, true, err } rows := make([]map[string]interface{}, 0) switch source { case "INFORMATION_SCHEMA.TABLES": for _, table := range tables { rows = append(rows, map[string]interface{}{ "table_catalog": c.database, "table_schema": "public", "table_name": table, "table_type": "BASE TABLE", }) } case "PG_TABLES": for _, table := range tables { indexes, _ := c.schema.ListTableIndexes(table) rows = append(rows, map[string]interface{}{ "schemaname": "public", "tablename": table, "tableowner": c.params["user"], "tablespace": nil, "hasindexes": len(indexes) > 0, "hasrules": false, "hastriggers": false, "rowsecurity": false, }) } case "INFORMATION_SCHEMA.COLUMNS": for _, table := range tables { schema, schemaErr := c.schema.GetSchema(table) if schemaErr != nil { continue } for i, column := range schema.Columns { rows = append(rows, map[string]interface{}{ "table_catalog": c.database, "table_schema": "public", "table_name": table, "column_name": column.Name, "ordinal_position": int64(i + 1), "column_default": column.Default, "is_nullable": yesNo(column.Nullable), "data_type": strings.ToLower(column.Type), }) } } case "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE": for _, table := range tables { schema, schemaErr := c.schema.GetSchema(table) if schemaErr != nil || schema.PrimaryKey == "" || schema.PrimaryKey == "_rowid_" { continue } row := map[string]interface{}{ "constraint_catalog": c.database, "constraint_schema": "public", "constraint_name": table + "_pkey", "table_catalog": c.database, "table_schema": "public", "table_name": table, } if source == "INFORMATION_SCHEMA.TABLE_CONSTRAINTS" { row["constraint_type"] = "PRIMARY KEY" row["is_deferrable"] = "NO" row["initially_deferred"] = "NO" } else { row["column_name"] = schema.PrimaryKey row["ordinal_position"] = int64(1) } rows = append(rows, row) } case "PG_INDEXES": for _, table := range tables { indexes, _ := c.schema.ListTableIndexes(table) for _, index := range indexes { columns := make([]string, len(index.Columns)) for i, column := range index.Columns { columns[i] = column.Name } unique := "" if index.Unique { unique = "UNIQUE " } rows = append(rows, map[string]interface{}{ "schemaname": "public", "tablename": table, "indexname": index.Name, "tablespace": nil, "indexdef": fmt.Sprintf("CREATE %sINDEX %s ON %s (%s)", unique, index.Name, table, strings.Join(columns, ", ")), }) } } } for _, match := range catalogFilterPattern.FindAllStringSubmatch(sql, -1) { column, expected := strings.ToLower(match[1]), strings.ReplaceAll(match[2], "''", "'") filtered := rows[:0] for _, row := range rows { if value, ok := row[column]; ok && strings.EqualFold(fmt.Sprintf("%v", value), expected) { filtered = append(filtered, row) } } rows = filtered } columns := catalogProjection(sql, rows) result := executor.NewResult("SELECT") if len(columns) == 1 && columns[0] == "count(*)" { result.AddColumnWithType("count", "INTEGER") result.AddRow(int64(len(rows))) return result, true, nil } for _, column := range columns { columnType := "TEXT" if column == "ordinal_position" { columnType = "INTEGER" } else if strings.HasPrefix(column, "has") || column == "rowsecurity" { columnType = "BOOLEAN" } result.AddColumnWithType(column, columnType) } for _, row := range rows { values := make([]interface{}, len(columns)) for i, column := range columns { values[i] = row[column] } result.AddRow(values...) } return result, true, nil } func isSingleStatementSQL(sql string) bool { trimmed := strings.TrimSpace(sql) if strings.HasSuffix(trimmed, ";") { trimmed = strings.TrimSpace(strings.TrimSuffix(trimmed, ";")) } inString, inIdent := false, false for i := 0; i < len(trimmed); i++ { switch trimmed[i] { case '\'': if !inIdent { if inString && i+1 < len(trimmed) && trimmed[i+1] == '\'' { i++ continue } inString = !inString } case '"': if !inString { inIdent = !inIdent } case ';': if !inString && !inIdent { return false } } } return true } // isRollbackStatement reports whether sql parses as a single ROLLBACK // statement, including ROLLBACK TO SAVEPOINT. func isRollbackStatement(sql string) bool { stmt, err := parser.New(lexer.New(sql)).Parse() if err != nil { return false } _, ok := stmt.(*parser.RollbackStmt) return ok } func catalogProjection(sql string, rows []map[string]interface{}) []string { upper := strings.ToUpper(sql) selectPos, fromPos := strings.Index(upper, "SELECT"), strings.Index(upper, " FROM ") if selectPos < 0 || fromPos < 0 || fromPos <= selectPos+6 { return nil } projection := strings.TrimSpace(sql[selectPos+6 : fromPos]) projection = strings.TrimSpace(strings.TrimPrefix(strings.ToUpper(projection), "DISTINCT ")) if projection == "*" && len(rows) > 0 { columns := make([]string, 0, len(rows[0])) for column := range rows[0] { columns = append(columns, column) } sort.Strings(columns) return columns } parts := strings.Split(projection, ",") columns := make([]string, 0, len(parts)) for _, part := range parts { column := strings.TrimSpace(part) if index := strings.Index(strings.ToUpper(column), " AS "); index >= 0 { column = strings.TrimSpace(column[:index]) } if index := strings.LastIndex(column, "."); index >= 0 { column = column[index+1:] } columns = append(columns, strings.ToLower(strings.Trim(column, `"`))) } return columns } func yesNo(value bool) string { if value { return "YES" } return "NO" } func (c *Connection) sendTabularResult(result *executor.Result, tag string) error { if err := c.sendRowDescription(result.Columns, result.ColumnTypes); err != nil { return err } for _, row := range result.Rows { if err := c.sendDataRow(row, result.Columns); err != nil { return err } } return c.sendCommandComplete(fmt.Sprintf("%s %d", tag, len(result.Rows))) } func (c *Connection) handleClose(msg *Message) error { if c.extendedFailed { return nil } if len(msg.Data) < 2 { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close message")) } name, pos, err := readCString(msg.Data, 1) if err != nil || pos != len(msg.Data) { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close name")) } if msg.Data[0] == 'S' { delete(c.statements, name) } else if msg.Data[0] == 'P' { delete(c.portals, name) } else { return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close target")) } return c.writeMessage(MsgCloseComplete, nil) } func (c *Connection) failExtended(code string, err error) error { c.extendedFailed = true return c.sendError("ERROR", code, err.Error()) } type boundParameter struct { value []byte oid int32 format int16 null bool } func readCString(data []byte, pos int) (string, int, error) { if pos < 0 || pos >= len(data) { return "", pos, fmt.Errorf("missing null-terminated string") } end := bytes.IndexByte(data[pos:], 0) if end < 0 { return "", pos, fmt.Errorf("unterminated string") } return string(data[pos : pos+end]), pos + end + 1, nil } func readInt16List(data []byte, pos int) ([]int16, int, error) { if pos+2 > len(data) { return nil, pos, fmt.Errorf("missing list length") } count := int(binary.BigEndian.Uint16(data[pos : pos+2])) pos += 2 if pos+count*2 > len(data) { return nil, pos, fmt.Errorf("truncated list") } result := make([]int16, count) for i := range result { result[i] = int16(binary.BigEndian.Uint16(data[pos : pos+2])) pos += 2 } return result, pos, nil } func parameterOID(oids []int32, index int) int32 { if index < len(oids) { return oids[index] } return 0 } func parameterFormat(formats []int16, index int) int16 { if len(formats) == 1 { return formats[0] } if index < len(formats) { return formats[index] } return 0 } func bindQuery(query string, params []boundParameter) (string, error) { var result strings.Builder inString, inIdent := false, false for i := 0; i < len(query); { ch := query[i] if ch == '\'' && !inIdent { result.WriteByte(ch) if inString && i+1 < len(query) && query[i+1] == '\'' { result.WriteByte(query[i+1]) i += 2 continue } inString = !inString i++ continue } if ch == '"' && !inString { inIdent = !inIdent result.WriteByte(ch) i++ continue } if ch == '$' && !inString && !inIdent && i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' { end := i + 1 for end < len(query) && query[end] >= '0' && query[end] <= '9' { end++ } n, _ := strconv.Atoi(query[i+1 : end]) if n < 1 || n > len(params) { return "", fmt.Errorf("parameter $%d was not provided", n) } literal, err := parameterLiteral(params[n-1]) if err != nil { return "", fmt.Errorf("parameter $%d: %w", n, err) } result.WriteString(literal) i = end continue } result.WriteByte(ch) i++ } return result.String(), nil } func parameterLiteral(param boundParameter) (string, error) { if param.null { return "NULL", nil } if param.format == 1 { switch param.oid { case 16: if len(param.value) != 1 { return "", fmt.Errorf("invalid binary boolean") } if param.value[0] == 0 { return "FALSE", nil } return "TRUE", nil case 21: if len(param.value) != 2 { return "", fmt.Errorf("invalid binary int2") } return strconv.FormatInt(int64(int16(binary.BigEndian.Uint16(param.value))), 10), nil case 23: if len(param.value) != 4 { return "", fmt.Errorf("invalid binary int4") } return strconv.FormatInt(int64(int32(binary.BigEndian.Uint32(param.value))), 10), nil case 20: if len(param.value) != 8 { return "", fmt.Errorf("invalid binary int8") } return strconv.FormatInt(int64(binary.BigEndian.Uint64(param.value)), 10), nil default: return "", fmt.Errorf("binary format is unsupported for OID %d", param.oid) } } value := string(param.value) switch param.oid { case 0: if strings.EqualFold(value, "true") || strings.EqualFold(value, "false") { return strings.ToUpper(value), nil } if _, err := strconv.ParseFloat(value, 64); err == nil && value != "" { return value, nil } return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil case 16: if strings.EqualFold(value, "true") || value == "1" || value == "t" { return "TRUE", nil } return "FALSE", nil case 20, 21, 23, 26, 700, 701, 1700: if _, err := strconv.ParseFloat(value, 64); err != nil { return "", fmt.Errorf("invalid numeric value") } return value, nil default: return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil } } // sendResult sends query results func (c *Connection) sendResult(result *executor.Result, stmt parser.Statement) error { // For SELECT statements, send row description and data rows if _, isSelect := stmt.(*parser.SelectStmt); isSelect && len(result.Columns) > 0 { // Send row description if err := c.sendRowDescription(result.Columns, result.ColumnTypes); err != nil { return err } // Send data rows for _, row := range result.Rows { if err := c.sendDataRow(row, result.Columns); err != nil { return err } } // Send command complete tag := fmt.Sprintf("SELECT %d", len(result.Rows)) return c.sendCommandComplete(tag) } // For other statements, just send command complete tag := c.getCommandTag(stmt, result) return c.sendCommandComplete(tag) } // getCommandTag returns the command completion tag func (c *Connection) getCommandTag(stmt parser.Statement, result *executor.Result) string { switch s := stmt.(type) { case *parser.CreateTableStmt: return "CREATE TABLE" case *parser.DropTableStmt: return "DROP TABLE" case *parser.CreateIndexStmt: return "CREATE INDEX" case *parser.DropIndexStmt: return "DROP INDEX" case *parser.AlterTableStmt: return "ALTER TABLE" case *parser.InsertStmt: return fmt.Sprintf("INSERT 0 %d", result.RowsAffected) case *parser.UpdateStmt: return fmt.Sprintf("UPDATE %d", result.RowsAffected) case *parser.DeleteStmt: return fmt.Sprintf("DELETE %d", result.RowsAffected) case *parser.BeginStmt: c.txStatus = TxStatusInBlock return "BEGIN" case *parser.CommitStmt: c.txStatus = TxStatusIdle return "COMMIT" case *parser.RollbackStmt: if s.Savepoint != "" { c.txStatus = TxStatusInBlock return "ROLLBACK" } c.txStatus = TxStatusIdle return "ROLLBACK" case *parser.SavepointStmt: c.txStatus = TxStatusInBlock return "SAVEPOINT" case *parser.ReleaseStmt: c.txStatus = TxStatusInBlock return "RELEASE" default: return "OK" } } // sendAuthenticationOk sends authentication OK message func (c *Connection) sendAuthenticationOk() error { mb := NewMessageBuilder() mb.WriteInt32(0) // Auth OK return c.writeMessage(MsgAuthenticationOk, mb.Bytes()) } // sendParameterStatus sends a parameter status message func (c *Connection) sendParameterStatus(name, value string) error { mb := NewMessageBuilder() mb.WriteString(name) mb.WriteString(value) return c.writeMessage(MsgParameterStatus, mb.Bytes()) } // sendBackendKeyData sends backend key data func (c *Connection) sendBackendKeyData(processID, secretKey int32) error { mb := NewMessageBuilder() mb.WriteInt32(processID) mb.WriteInt32(secretKey) return c.writeMessage(MsgBackendKeyData, mb.Bytes()) } // sendReadyForQuery sends ready for query message func (c *Connection) sendReadyForQuery() error { mb := NewMessageBuilder() mb.AppendByte(c.txStatus) if err := c.writeMessage(MsgReadyForQuery, mb.Bytes()); err != nil { return err } // ReadyForQuery closes out a response cycle, so flush everything buffered // so far. This is what makes simple-query results and Sync responses // visible to the client. return c.writer.Flush() } // sendEmptyQueryResponse sends empty query response func (c *Connection) sendEmptyQueryResponse() error { return c.writeMessage(MsgEmptyQueryResponse, []byte{}) } // sendRowDescription sends row description (column metadata) func (c *Connection) sendRowDescription(columns []string, columnTypes []string) error { mb := NewMessageBuilder() mb.WriteInt16(int16(len(columns))) for i, col := range columns { colType := "" if i < len(columnTypes) { colType = columnTypes[i] } mb.WriteString(col) mb.WriteInt32(0) // table OID mb.WriteInt16(0) // column attribute number mb.WriteInt32(c.getOIDForType(colType)) // type OID mb.WriteInt16(c.getTypeSizeForType(colType)) // type size mb.WriteInt32(-1) // type modifier mb.WriteInt16(0) // format code (text) } return c.writeMessage(MsgRowDescription, mb.Bytes()) } // sendDataRow sends a data row func (c *Connection) sendDataRow(row []interface{}, columns []string) error { mb := NewMessageBuilder() mb.WriteInt16(int16(len(row))) for _, value := range row { if value == nil { mb.WriteInt32(-1) // NULL indicator continue } // Convert value to string strValue := c.valueToString(value) mb.WriteInt32(int32(len(strValue))) mb.WriteBytes([]byte(strValue)) } return c.writeMessage(MsgDataRow, mb.Bytes()) } // sendCommandComplete sends command complete message func (c *Connection) sendCommandComplete(tag string) error { mb := NewMessageBuilder() mb.WriteString(tag) return c.writeMessage(MsgCommandComplete, mb.Bytes()) } // sendError sends an error response func (c *Connection) sendError(severity, code, message string) error { mb := NewMessageBuilder() mb.AppendByte(ErrorFieldSeverity) mb.WriteString(severity) mb.AppendByte(ErrorFieldCode) mb.WriteString(code) mb.AppendByte(ErrorFieldMessage) mb.WriteString(message) mb.AppendByte(0) // Terminator if err := c.writeMessage(MsgErrorResponse, mb.Bytes()); err != nil { return err } // Errors must be visible promptly, including FATAL startup failures that // are not followed by a ReadyForQuery before the connection closes. return c.writer.Flush() } // writeMessage buffers a message for the connection. It does not flush, so // callers that need to make a response visible to the client must flush at the // appropriate protocol boundary (ReadyForQuery, an explicit Flush message, or // an error response). Buffering amortizes the per-message syscalls that a // result set would otherwise incur; the underlying bufio.Writer bounds memory // use so large result sets cannot grow the buffer without limit. func (c *Connection) writeMessage(msgType byte, data []byte) error { if !c.quiet { log.Printf("Sending message type=%c length=%d", msgType, len(data)+4) } return WriteMessage(c.writer, msgType, data) } // getOIDForType returns PostgreSQL OID for type func (c *Connection) getOIDForType(typeName string) int32 { switch strings.ToUpper(typeName) { case "INTEGER", "INT": return 23 // INT4OID case "TEXT", "VARCHAR", "CHAR": return 25 // TEXTOID case "REAL", "FLOAT": return 700 // FLOAT4OID case "DOUBLE": return 701 // FLOAT8OID case "BOOLEAN", "BOOL": return 16 // BOOLOID case "BLOB": return 17 // BYTEAOID default: return 25 // Default to TEXT } } // getTypeSizeForType returns type size func (c *Connection) getTypeSizeForType(typeName string) int16 { switch strings.ToUpper(typeName) { case "INTEGER", "INT": return 4 case "REAL", "FLOAT": return 4 case "DOUBLE": return 8 case "BOOLEAN", "BOOL": return 1 default: return -1 // Variable length } } // valueToString converts a value to string func (c *Connection) valueToString(value interface{}) string { if value == nil { return "" } return fmt.Sprintf("%v", value) } // handleVersionQuery handles SELECT version() func (c *Connection) handleVersionQuery() error { columns := []string{"version"} columnTypes := []string{"TEXT"} if err := c.sendRowDescription(columns, columnTypes); err != nil { return err } row := []interface{}{"PostgreSQL 14.0 (PizzaSQL)"} if err := c.sendDataRow(row, columns); err != nil { return err } if err := c.sendCommandComplete("SELECT 1"); err != nil { return err } return c.sendReadyForQuery() } // handleCurrentUserQuery handles SELECT current_user func (c *Connection) handleCurrentUserQuery() error { columns := []string{"current_user"} columnTypes := []string{"TEXT"} if err := c.sendRowDescription(columns, columnTypes); err != nil { return err } user := c.params["user"] if user == "" { user = "pizzasql" } row := []interface{}{user} if err := c.sendDataRow(row, columns); err != nil { return err } if err := c.sendCommandComplete("SELECT 1"); err != nil { return err } return c.sendReadyForQuery() } // handleShowCommand handles SHOW commands func (c *Connection) handleShowCommand(sqlUpper string) error { var value string var name string if strings.Contains(sqlUpper, "SERVER_VERSION") { name = "server_version" value = "14.0" } else if strings.Contains(sqlUpper, "SERVER_ENCODING") { name = "server_encoding" value = "UTF8" } else if strings.Contains(sqlUpper, "CLIENT_ENCODING") { name = "client_encoding" value = "UTF8" } else { // Unknown SHOW command c.sendError("ERROR", ErrCodeFeatureNotSupported, "SHOW command not supported") return c.sendReadyForQuery() } columns := []string{name} columnTypes := []string{"TEXT"} if err := c.sendRowDescription(columns, columnTypes); err != nil { return err } row := []interface{}{value} if err := c.sendDataRow(row, columns); err != nil { return err } if err := c.sendCommandComplete("SHOW"); err != nil { return err } return c.sendReadyForQuery() }