| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220 |
- 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()
- }
|