package pgserver import ( "bufio" "bytes" "encoding/binary" "fmt" "io" "log" "net" "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 } // 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), txStatus: TxStatusIdle, quiet: quiet, } } // Handle processes the connection func (c *Connection) Handle() error { defer 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, MsgBind, MsgDescribe, MsgExecute, MsgSync, MsgClose: // Extended query protocol - not implemented yet c.sendError("ERROR", ErrCodeFeatureNotSupported, "Extended query protocol not yet supported") 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 { // Parse query string (null-terminated) 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)) // lib/pq and other drivers query these for connection validation if strings.Contains(sqlUpper, "SELECT VERSION()") { // Return a fake PostgreSQL version return c.handleVersionQuery() } if strings.Contains(sqlUpper, "SELECT CURRENT_USER") { // Return the current user return c.handleCurrentUserQuery() } if 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) } // Execute query l := lexer.New(sql) p := parser.New(l) stmt, err := p.Parse() if err != nil { c.sendError("ERROR", ErrCodeSyntaxError, fmt.Sprintf("Syntax error: %v", err)) return c.sendReadyForQuery() } result, err := c.executor.Execute(stmt) if err != nil { c.sendError("ERROR", ErrCodeInternalError, fmt.Sprintf("Execution error: %v", err)) return c.sendReadyForQuery() } // Send result based on statement type if err := c.sendResult(result, stmt); err != nil { return err } return c.sendReadyForQuery() } // 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 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.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: c.txStatus = TxStatusIdle return "ROLLBACK" 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.WriteByte(c.txStatus) return c.writeMessage(MsgReadyForQuery, mb.Bytes()) } // 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.WriteByte(ErrorFieldSeverity) mb.WriteString(severity) mb.WriteByte(ErrorFieldCode) mb.WriteString(code) mb.WriteByte(ErrorFieldMessage) mb.WriteString(message) mb.WriteByte(0) // Terminator return c.writeMessage(MsgErrorResponse, mb.Bytes()) } // writeMessage writes a message to the connection func (c *Connection) writeMessage(msgType byte, data []byte) error { if !c.quiet { log.Printf("Sending message type=%c length=%d", msgType, len(data)+4) } if err := WriteMessage(c.writer, msgType, data); err != nil { return err } return c.writer.Flush() } // 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() }