connection.go 48 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593
  1. package pgserver
  2. import (
  3. "bufio"
  4. "bytes"
  5. "encoding/binary"
  6. "encoding/hex"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "log"
  11. "net"
  12. "os"
  13. "regexp"
  14. "sort"
  15. "strconv"
  16. "strings"
  17. "sync/atomic"
  18. "github.com/danfragoso/pizzasql-next/pkg/executor"
  19. "github.com/danfragoso/pizzasql-next/pkg/lexer"
  20. "github.com/danfragoso/pizzasql-next/pkg/parser"
  21. "github.com/danfragoso/pizzasql-next/pkg/pgerror"
  22. "github.com/danfragoso/pizzasql-next/pkg/storage"
  23. "github.com/danfragoso/pizzasql-next/pkg/version"
  24. )
  25. // Connection represents a client connection
  26. type Connection struct {
  27. conn net.Conn
  28. reader *bufio.Reader
  29. writer *bufio.Writer
  30. executor *executor.Executor
  31. schema *storage.SchemaManager
  32. dbManager *storage.DatabaseManager
  33. database string
  34. params map[string]string
  35. txStatus byte
  36. quiet bool // Disable query logging
  37. statements map[string]*preparedStatement
  38. portals map[string]*portal
  39. extendedFailed bool
  40. backendPID int32
  41. // extendedCycle is true between Parse/Bind/Execute/Describe/Close and the
  42. // next Sync. Notifications are not interleaved into an extended-protocol
  43. // sequence; they are held until Sync.
  44. extendedCycle bool
  45. }
  46. // backendPIDSeq gives each connection a distinct synthetic backend PID. It is
  47. // seeded with the process ID so NotificationResponse frames look plausible and
  48. // do not collide across server processes.
  49. var backendPIDSeq atomic.Int32
  50. func nextBackendPID() int32 {
  51. return int32(os.Getpid()) + backendPIDSeq.Add(1)
  52. }
  53. type preparedStatement struct {
  54. query string
  55. paramOIDs []int32
  56. }
  57. type portal struct {
  58. query string
  59. executed bool
  60. }
  61. // NewConnection creates a new connection handler
  62. func NewConnection(conn net.Conn, dbManager *storage.DatabaseManager, quiet bool) *Connection {
  63. return &Connection{
  64. conn: conn,
  65. reader: bufio.NewReader(conn),
  66. writer: bufio.NewWriter(conn),
  67. dbManager: dbManager,
  68. params: make(map[string]string),
  69. statements: make(map[string]*preparedStatement),
  70. portals: make(map[string]*portal),
  71. txStatus: TxStatusIdle,
  72. quiet: quiet,
  73. backendPID: nextBackendPID(),
  74. }
  75. }
  76. // Handle processes the connection
  77. func (c *Connection) Handle() error {
  78. defer func() {
  79. if c.executor != nil {
  80. if err := c.executor.Close(); err != nil && !c.quiet {
  81. log.Printf("failed to close session: %v", err)
  82. }
  83. }
  84. c.conn.Close()
  85. }()
  86. // First, check for SSL request (sent before startup message)
  87. // SSL request is 8 bytes: length(4) + code(4) where code = 80877103
  88. firstBytes := make([]byte, 8)
  89. n, err := io.ReadFull(c.reader, firstBytes)
  90. if err != nil {
  91. return fmt.Errorf("failed to read initial bytes: %w", err)
  92. }
  93. // Check if it's an SSL request (code 80877103 = 0x04D2162F)
  94. if n == 8 {
  95. length := binary.BigEndian.Uint32(firstBytes[0:4])
  96. code := binary.BigEndian.Uint32(firstBytes[4:8])
  97. if length == 8 && code == 80877103 {
  98. // SSL request - we don't support SSL, send 'N'
  99. if !c.quiet {
  100. log.Printf("Client requested SSL, sending rejection")
  101. }
  102. if _, err := c.conn.Write([]byte{'N'}); err != nil {
  103. return fmt.Errorf("failed to send SSL rejection: %w", err)
  104. }
  105. // Now read the actual startup message
  106. } else {
  107. // Not SSL request, this is part of startup message
  108. // We need to prepend these bytes back for ReadStartupMessage
  109. // Create a multi-reader that first reads our buffered bytes, then continues with the reader
  110. c.reader = bufio.NewReader(io.MultiReader(bytes.NewReader(firstBytes), c.reader))
  111. }
  112. }
  113. // Read startup message
  114. if !c.quiet {
  115. log.Printf("Reading startup message...")
  116. }
  117. params, err := ReadStartupMessage(c.reader)
  118. if err != nil {
  119. return fmt.Errorf("failed to read startup message: %w", err)
  120. }
  121. c.params = params
  122. if !c.quiet {
  123. log.Printf("Startup params: %+v", params)
  124. }
  125. // Get database name from params (default to "pizzasql")
  126. dbName := params["database"]
  127. if dbName == "" {
  128. dbName = "pizzasql"
  129. }
  130. c.database = dbName
  131. if !c.quiet {
  132. log.Printf("New connection: user=%s database=%s", params["user"], dbName)
  133. }
  134. // Initialize database
  135. if err := c.initDatabase(dbName); err != nil {
  136. c.sendError("FATAL", ErrCodeConnectionFailure, fmt.Sprintf("Failed to initialize database: %v", err))
  137. return err
  138. }
  139. // Send authentication OK (no auth for now)
  140. if err := c.sendAuthenticationOk(); err != nil {
  141. return err
  142. }
  143. // Send parameter status messages
  144. if err := c.sendParameterStatus("server_version", version.PostgreSQLCompatVersion+" (PizzaSQL)"); err != nil {
  145. return err
  146. }
  147. // Build identity is a separate, custom ParameterStatus so the PostgreSQL
  148. // compatibility string above stays stable for version-sniffing drivers.
  149. if err := c.sendParameterStatus("pizzasql.version", version.String()); err != nil {
  150. return err
  151. }
  152. if err := c.sendParameterStatus("pizzasql.commit", version.Commit); err != nil {
  153. return err
  154. }
  155. if err := c.sendParameterStatus("pizzasql.proxy_errors", "1"); err != nil {
  156. return err
  157. }
  158. if err := c.sendParameterStatus("server_encoding", "UTF8"); err != nil {
  159. return err
  160. }
  161. if err := c.sendParameterStatus("client_encoding", "UTF8"); err != nil {
  162. return err
  163. }
  164. if err := c.sendParameterStatus("DateStyle", "ISO, MDY"); err != nil {
  165. return err
  166. }
  167. if err := c.sendParameterStatus("TimeZone", "UTC"); err != nil {
  168. return err
  169. }
  170. // Send backend key data (for cancellation - we don't implement this yet)
  171. if err := c.sendBackendKeyData(c.backendPID, 67890); err != nil {
  172. return err
  173. }
  174. // Send ready for query
  175. if err := c.sendReadyForQuery(); err != nil {
  176. return err
  177. }
  178. // Read messages on a separate goroutine so an idle listener can receive
  179. // asynchronous NotificationResponse frames while blocked waiting for input.
  180. // Only this goroutine writes to c.writer, which keeps frame order intact.
  181. type readResult struct {
  182. msg *Message
  183. err error
  184. }
  185. inbox := make(chan readResult, 16)
  186. done := make(chan struct{})
  187. defer close(done)
  188. go func() {
  189. for {
  190. msg, err := ReadMessage(c.reader)
  191. select {
  192. case inbox <- readResult{msg: msg, err: err}:
  193. case <-done:
  194. return
  195. }
  196. if err != nil {
  197. return
  198. }
  199. }
  200. }()
  201. notifications := c.executor.NotificationChannel()
  202. overflow := c.executor.NotificationOverflow()
  203. for {
  204. // Do not interleave notifications into an extended-protocol cycle. The
  205. // channel is disabled until Sync closes the cycle.
  206. var notifyCh <-chan storage.Notification
  207. if !c.extendedCycle {
  208. notifyCh = notifications
  209. }
  210. select {
  211. case res := <-inbox:
  212. if res.err != nil {
  213. if res.err == io.EOF {
  214. log.Printf("Connection closed by client")
  215. return nil
  216. }
  217. return fmt.Errorf("failed to read message: %w", res.err)
  218. }
  219. switch res.msg.Type {
  220. case MsgParse, MsgBind, MsgDescribe, MsgExecute, MsgClose:
  221. c.extendedCycle = true
  222. }
  223. if err := c.handleMessage(res.msg); err != nil {
  224. if err == io.EOF {
  225. log.Printf("Connection closed normally")
  226. return nil
  227. }
  228. log.Printf("Error handling message: %v", err)
  229. return err
  230. }
  231. case n := <-notifyCh:
  232. if err := c.sendNotification(n); err != nil {
  233. return err
  234. }
  235. if err := c.writer.Flush(); err != nil {
  236. return err
  237. }
  238. case <-overflow:
  239. return fmt.Errorf("notification queue overflow: slow consumer disconnected")
  240. }
  241. }
  242. }
  243. // initDatabase initializes the database connection
  244. func (c *Connection) initDatabase(dbName string) error {
  245. db, err := c.dbManager.GetDatabase(dbName)
  246. if err != nil {
  247. return err
  248. }
  249. c.schema = db.Schema
  250. c.executor = executor.New(db.Schema, db.Table)
  251. c.executor.SetBackendPID(c.backendPID)
  252. c.executor.SetNotificationHub(db.Hub)
  253. c.executor.SyncCatalog()
  254. return nil
  255. }
  256. // handleMessage processes a client message
  257. func (c *Connection) handleMessage(msg *Message) error {
  258. if c.extendedFailed && msg.Type != MsgSync && msg.Type != MsgFlush && msg.Type != MsgTerminate {
  259. return nil
  260. }
  261. switch msg.Type {
  262. case MsgProxyError:
  263. return c.handleProxyError(msg)
  264. case MsgQuery:
  265. return c.handleQuery(msg)
  266. case MsgTerminate:
  267. log.Printf("Client requested termination")
  268. return io.EOF
  269. case MsgParse:
  270. return c.handleParse(msg)
  271. case MsgBind:
  272. return c.handleBind(msg)
  273. case MsgDescribe:
  274. return c.handleDescribe(msg)
  275. case MsgExecute:
  276. return c.handleExecute(msg)
  277. case MsgClose:
  278. return c.handleClose(msg)
  279. case MsgFlush:
  280. return c.writer.Flush()
  281. case MsgSync:
  282. c.extendedFailed = false
  283. c.extendedCycle = false
  284. if err := c.sendReadyForQuery(); err != nil {
  285. return err
  286. }
  287. return c.drainNotifications()
  288. default:
  289. log.Printf("Unknown message type: %c (%d)", msg.Type, msg.Type)
  290. c.sendError("ERROR", ErrCodeProtocolViolation, fmt.Sprintf("Unknown message type: %c", msg.Type))
  291. return c.sendReadyForQuery()
  292. }
  293. }
  294. // handleQuery processes a simple query
  295. func (c *Connection) handleQuery(msg *Message) error {
  296. // A Query message carries a single NUL-terminated query string.
  297. if len(msg.Data) == 0 || msg.Data[len(msg.Data)-1] != 0 {
  298. c.sendError("ERROR", ErrCodeProtocolViolation, "invalid Query message: missing null terminator")
  299. return c.sendReadyForQuery()
  300. }
  301. sql := string(msg.Data[:len(msg.Data)-1])
  302. if !c.quiet {
  303. log.Printf("Query: %s", sql)
  304. }
  305. // Handle empty query
  306. sqlTrimmed := strings.TrimSpace(sql)
  307. if sqlTrimmed == "" || sqlTrimmed == ";" {
  308. if err := c.sendEmptyQueryResponse(); err != nil {
  309. return err
  310. }
  311. return c.sendReadyForQuery()
  312. }
  313. // Handle special PostgreSQL system queries that drivers send
  314. singleStatement := isSingleStatementSQL(sql)
  315. // In a failed transaction only ROLLBACK (or ROLLBACK TO SAVEPOINT) is
  316. // accepted. Gate before honoring the special driver queries below so they
  317. // return 25P02 instead of a result. Multi-statement batches are gated
  318. // per statement in the execution loop below.
  319. if c.txStatus == TxStatusFailed && singleStatement && !isRollbackStatement(sql) {
  320. c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block")
  321. return c.sendReadyForQuery()
  322. }
  323. // Recognize only complete initialization-query token sequences.
  324. if singleStatement {
  325. switch kind := systemQueryKind(sql); kind {
  326. case "version":
  327. return c.handleVersionQuery()
  328. case "current_user":
  329. return c.handleCurrentUserQuery()
  330. case "server_version", "server_encoding", "client_encoding":
  331. return c.handleShowCommand(strings.ToUpper(kind))
  332. }
  333. }
  334. if result, handled, err := c.catalogResult(sql); handled {
  335. if err != nil {
  336. c.sendError("ERROR", ErrCodeInternalError, err.Error())
  337. return c.sendReadyForQuery()
  338. }
  339. if err := c.sendTabularResult(result, "SELECT"); err != nil {
  340. return err
  341. }
  342. return c.sendReadyForQuery()
  343. }
  344. // Parse the complete batch before executing its first statement. This keeps
  345. // unsupported trailing clauses from turning into committed writes.
  346. l := lexer.New(sql)
  347. p := parser.New(l)
  348. stmts, err := p.ParseMultiple()
  349. if err != nil {
  350. code := ErrCodeSyntaxError
  351. if errors.Is(err, pgerror.ErrFeatureNotSupported) {
  352. code = ErrCodeFeatureNotSupported
  353. }
  354. if c.txStatus == TxStatusInBlock {
  355. c.txStatus = TxStatusFailed
  356. }
  357. c.sendError("ERROR", code, fmt.Sprintf("Syntax error: %v", err))
  358. return c.sendReadyForQuery()
  359. }
  360. for _, stmt := range stmts {
  361. if c.txStatus == TxStatusFailed {
  362. if _, rollback := stmt.(*parser.RollbackStmt); !rollback {
  363. c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block")
  364. return c.sendReadyForQuery()
  365. }
  366. }
  367. result, err := c.executor.Execute(stmt)
  368. if err != nil && isIdleRollback(stmt, c.txStatus, err) {
  369. result, err = executor.NewResult("ROLLBACK"), nil
  370. }
  371. if err != nil {
  372. code := ErrCodeInternalError
  373. message := fmt.Sprintf("Execution error: %v", err)
  374. var primaryKeyErr *storage.PrimaryKeyViolationError
  375. var uniqueErr *storage.UniqueViolationError
  376. var notNullErr *storage.NotNullViolationError
  377. var checkErr *storage.CheckViolationError
  378. var foreignKeyErr *storage.ForeignKeyViolationError
  379. if errors.Is(err, storage.ErrSerialization) {
  380. code = ErrCodeSerializationFailure
  381. c.txStatus = TxStatusIdle
  382. } else if errors.As(err, &notNullErr) {
  383. code = "23502"
  384. if _, commit := stmt.(*parser.CommitStmt); commit {
  385. c.txStatus = TxStatusIdle
  386. } else if c.txStatus == TxStatusInBlock {
  387. c.txStatus = TxStatusFailed
  388. }
  389. } else if errors.As(err, &checkErr) {
  390. code = ErrCodeCheckViolation
  391. if _, commit := stmt.(*parser.CommitStmt); commit {
  392. c.txStatus = TxStatusIdle
  393. } else if c.txStatus == TxStatusInBlock {
  394. c.txStatus = TxStatusFailed
  395. }
  396. } else if errors.As(err, &foreignKeyErr) {
  397. code = ErrCodeForeignKeyViolation
  398. if _, commit := stmt.(*parser.CommitStmt); commit {
  399. c.txStatus = TxStatusIdle
  400. } else if c.txStatus == TxStatusInBlock {
  401. c.txStatus = TxStatusFailed
  402. }
  403. } else if errors.As(err, &uniqueErr) || errors.As(err, &primaryKeyErr) {
  404. code = ErrCodeUniqueViolation
  405. if uniqueErr != nil {
  406. message = fmt.Sprintf("Execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint)
  407. }
  408. if _, commit := stmt.(*parser.CommitStmt); commit {
  409. c.txStatus = TxStatusIdle
  410. } else if c.txStatus == TxStatusInBlock {
  411. c.txStatus = TxStatusFailed
  412. }
  413. } else if errors.Is(err, pgerror.ErrFeatureNotSupported) || errors.Is(err, executor.ErrUnsupportedFeature) {
  414. code = ErrCodeFeatureNotSupported
  415. if c.txStatus == TxStatusInBlock {
  416. c.txStatus = TxStatusFailed
  417. }
  418. } else if _, commit := stmt.(*parser.CommitStmt); commit {
  419. // A failed COMMIT has already rolled the transaction back in
  420. // storage; leaving the session in the failed state would make
  421. // every following statement (including ROLLBACK) fail.
  422. c.txStatus = TxStatusIdle
  423. } else if c.txStatus == TxStatusInBlock {
  424. c.txStatus = TxStatusFailed
  425. }
  426. c.sendError("ERROR", code, message)
  427. return c.sendReadyForQuery()
  428. }
  429. if err := c.sendResult(result, stmt); err != nil {
  430. return err
  431. }
  432. if result.LSN > 0 {
  433. if err := c.sendLSN(result.LSN); err != nil {
  434. return err
  435. }
  436. }
  437. }
  438. return c.sendReadyForQuery()
  439. }
  440. func (c *Connection) handleParse(msg *Message) error {
  441. if c.extendedFailed {
  442. return nil
  443. }
  444. name, pos, err := readCString(msg.Data, 0)
  445. if err != nil {
  446. return c.failExtended(ErrCodeProtocolViolation, err)
  447. }
  448. query, pos, err := readCString(msg.Data, pos)
  449. if err != nil || pos+2 > len(msg.Data) {
  450. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse message"))
  451. }
  452. count := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2]))
  453. pos += 2
  454. if pos+count*4 != len(msg.Data) {
  455. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse parameter list"))
  456. }
  457. oids := make([]int32, count)
  458. for i := range oids {
  459. oids[i] = int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
  460. pos += 4
  461. }
  462. if !isSingleStatementSQL(query) {
  463. return c.failExtended(ErrCodeSyntaxError, fmt.Errorf("prepared statements must contain one command"))
  464. }
  465. c.statements[name] = &preparedStatement{query: query, paramOIDs: oids}
  466. return c.writeMessage(MsgParseComplete, nil)
  467. }
  468. func (c *Connection) handleBind(msg *Message) error {
  469. if c.extendedFailed {
  470. return nil
  471. }
  472. portalName, pos, err := readCString(msg.Data, 0)
  473. if err != nil {
  474. return c.failExtended(ErrCodeProtocolViolation, err)
  475. }
  476. statementName, pos, err := readCString(msg.Data, pos)
  477. if err != nil {
  478. return c.failExtended(ErrCodeProtocolViolation, err)
  479. }
  480. statement, ok := c.statements[statementName]
  481. if !ok {
  482. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", statementName))
  483. }
  484. formats, pos, err := readInt16List(msg.Data, pos)
  485. if err != nil || pos+2 > len(msg.Data) {
  486. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind format list"))
  487. }
  488. paramCount := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2]))
  489. pos += 2
  490. if len(formats) > 1 && len(formats) != paramCount {
  491. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("Bind parameter format count does not match parameters"))
  492. }
  493. for _, format := range formats {
  494. if format != 0 && format != 1 {
  495. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter format code"))
  496. }
  497. }
  498. params := make([]boundParameter, paramCount)
  499. for i := range params {
  500. if pos+4 > len(msg.Data) {
  501. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter"))
  502. }
  503. length := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
  504. pos += 4
  505. params[i].oid = parameterOID(statement.paramOIDs, i)
  506. params[i].format = parameterFormat(formats, i)
  507. if length == -1 {
  508. params[i].null = true
  509. continue
  510. }
  511. if length < 0 || pos+int(length) > len(msg.Data) {
  512. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter length"))
  513. }
  514. params[i].value = append([]byte(nil), msg.Data[pos:pos+int(length)]...)
  515. pos += int(length)
  516. }
  517. resultFormats, pos, err := readInt16List(msg.Data, pos)
  518. if err != nil || pos != len(msg.Data) {
  519. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format list"))
  520. }
  521. for _, format := range resultFormats {
  522. if format != 0 && format != 1 {
  523. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format code"))
  524. }
  525. if format == 1 {
  526. return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("binary result format is not supported; request text results"))
  527. }
  528. }
  529. if len(resultFormats) > 1 {
  530. return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("per-column result format lists are not supported; request all-text results"))
  531. }
  532. query, err := bindQuery(statement.query, params)
  533. if err != nil {
  534. code := ErrCodeInvalidParameter
  535. if errors.Is(err, pgerror.ErrFeatureNotSupported) {
  536. code = ErrCodeFeatureNotSupported
  537. }
  538. return c.failExtended(code, err)
  539. }
  540. c.portals[portalName] = &portal{query: query}
  541. return c.writeMessage(MsgBindComplete, nil)
  542. }
  543. func (c *Connection) handleDescribe(msg *Message) error {
  544. if c.extendedFailed {
  545. return nil
  546. }
  547. if len(msg.Data) < 2 {
  548. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe message"))
  549. }
  550. name, pos, err := readCString(msg.Data, 1)
  551. if err != nil || pos != len(msg.Data) {
  552. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe name"))
  553. }
  554. switch msg.Data[0] {
  555. case 'S':
  556. statement, ok := c.statements[name]
  557. if !ok {
  558. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", name))
  559. }
  560. mb := NewMessageBuilder()
  561. mb.WriteInt16(int16(len(statement.paramOIDs)))
  562. for _, oid := range statement.paramOIDs {
  563. mb.WriteInt32(oid)
  564. }
  565. if err := c.writeMessage(MsgParameterDescription, mb.Bytes()); err != nil {
  566. return err
  567. }
  568. case 'P':
  569. if _, ok := c.portals[name]; !ok {
  570. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
  571. }
  572. default:
  573. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe target"))
  574. }
  575. // Row metadata requires analysis without side effects. Return NoData for
  576. // compatibility; Execute still emits the real RowDescription. This is the
  577. // pre-0.4.0 behavior used by common extended-protocol clients.
  578. return c.writeMessage(MsgNoData, nil)
  579. }
  580. func (c *Connection) handleExecute(msg *Message) error {
  581. if c.extendedFailed {
  582. return nil
  583. }
  584. name, pos, err := readCString(msg.Data, 0)
  585. if err != nil || pos+4 != len(msg.Data) {
  586. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute message"))
  587. }
  588. portal, ok := c.portals[name]
  589. if !ok {
  590. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
  591. }
  592. if portal.executed {
  593. return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("portal can only be executed once"))
  594. }
  595. maxRows := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
  596. if maxRows < 0 {
  597. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute row limit"))
  598. }
  599. // Check the transaction state before catalog emulation. Catalog queries do
  600. // not all parse as regular PizzaSQL statements, but must still return 25P02
  601. // while the transaction is aborted.
  602. if c.txStatus == TxStatusFailed && !isRollbackStatement(portal.query) {
  603. return c.failExtended(ErrCodeTransactionAborted, fmt.Errorf("current transaction is aborted, commands ignored until end of transaction block"))
  604. }
  605. if maxRows != 0 {
  606. return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("bounded Execute is not supported; use maxRows=0"))
  607. }
  608. portal.executed = true
  609. if result, handled, catalogErr := c.catalogResult(portal.query); handled {
  610. if catalogErr != nil {
  611. return c.failExtended(ErrCodeInternalError, catalogErr)
  612. }
  613. return c.sendTabularResult(result, "SELECT")
  614. }
  615. l := lexer.New(portal.query)
  616. stmt, err := parser.New(l).Parse()
  617. if err != nil {
  618. code := ErrCodeSyntaxError
  619. if errors.Is(err, pgerror.ErrFeatureNotSupported) {
  620. code = ErrCodeFeatureNotSupported
  621. if c.txStatus == TxStatusInBlock {
  622. c.txStatus = TxStatusFailed
  623. }
  624. }
  625. return c.failExtended(code, fmt.Errorf("syntax error: %w", err))
  626. }
  627. result, err := c.executor.Execute(stmt)
  628. if err != nil && isIdleRollback(stmt, c.txStatus, err) {
  629. result, err = executor.NewResult("ROLLBACK"), nil
  630. }
  631. if err != nil {
  632. code := ErrCodeInternalError
  633. var primaryKeyErr *storage.PrimaryKeyViolationError
  634. var uniqueErr *storage.UniqueViolationError
  635. var notNullErr *storage.NotNullViolationError
  636. var checkErr *storage.CheckViolationError
  637. var foreignKeyErr *storage.ForeignKeyViolationError
  638. if errors.Is(err, storage.ErrSerialization) {
  639. code = ErrCodeSerializationFailure
  640. c.txStatus = TxStatusIdle
  641. } else if errors.As(err, &notNullErr) {
  642. code = "23502"
  643. if _, commit := stmt.(*parser.CommitStmt); commit {
  644. c.txStatus = TxStatusIdle
  645. } else if c.txStatus == TxStatusInBlock {
  646. c.txStatus = TxStatusFailed
  647. }
  648. } else if errors.As(err, &checkErr) {
  649. code = ErrCodeCheckViolation
  650. if _, commit := stmt.(*parser.CommitStmt); commit {
  651. c.txStatus = TxStatusIdle
  652. } else if c.txStatus == TxStatusInBlock {
  653. c.txStatus = TxStatusFailed
  654. }
  655. } else if errors.As(err, &foreignKeyErr) {
  656. code = ErrCodeForeignKeyViolation
  657. if _, commit := stmt.(*parser.CommitStmt); commit {
  658. c.txStatus = TxStatusIdle
  659. } else if c.txStatus == TxStatusInBlock {
  660. c.txStatus = TxStatusFailed
  661. }
  662. } else if errors.As(err, &uniqueErr) || errors.As(err, &primaryKeyErr) {
  663. code = ErrCodeUniqueViolation
  664. if _, commit := stmt.(*parser.CommitStmt); commit {
  665. c.txStatus = TxStatusIdle
  666. } else if c.txStatus == TxStatusInBlock {
  667. c.txStatus = TxStatusFailed
  668. }
  669. } else if errors.Is(err, pgerror.ErrFeatureNotSupported) || errors.Is(err, executor.ErrUnsupportedFeature) {
  670. code = ErrCodeFeatureNotSupported
  671. if c.txStatus == TxStatusInBlock {
  672. c.txStatus = TxStatusFailed
  673. }
  674. } else if _, commit := stmt.(*parser.CommitStmt); commit {
  675. c.txStatus = TxStatusIdle
  676. } else if c.txStatus == TxStatusInBlock {
  677. c.txStatus = TxStatusFailed
  678. }
  679. if uniqueErr != nil {
  680. return c.failExtended(code, fmt.Errorf("execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint))
  681. }
  682. return c.failExtended(code, fmt.Errorf("execution error: %w", err))
  683. }
  684. if err := c.sendResult(result, stmt); err != nil {
  685. return err
  686. }
  687. if result.LSN > 0 {
  688. return c.sendLSN(result.LSN)
  689. }
  690. return nil
  691. }
  692. var catalogFilterPattern = regexp.MustCompile(`(?i)\b(table_name|tablename|table_schema|schemaname|constraint_name|indexname)\s*=\s*'((?:''|[^'])*)'`)
  693. func (c *Connection) catalogResult(sql string) (*executor.Result, bool, error) {
  694. if !isSingleStatementSQL(sql) {
  695. return nil, false, nil
  696. }
  697. source := catalogSource(sql)
  698. if source == "" {
  699. return nil, false, nil
  700. }
  701. if c.txStatus == TxStatusIdle {
  702. c.schema.LockStatement()
  703. defer c.schema.UnlockStatement()
  704. }
  705. tables, err := c.schema.ListTables()
  706. if err != nil {
  707. return nil, true, err
  708. }
  709. rows := make([]map[string]interface{}, 0)
  710. switch source {
  711. case "INFORMATION_SCHEMA.TABLES":
  712. for _, table := range tables {
  713. rows = append(rows, map[string]interface{}{
  714. "table_catalog": c.database, "table_schema": "public", "table_name": table, "table_type": "BASE TABLE",
  715. })
  716. }
  717. case "PG_TABLES":
  718. for _, table := range tables {
  719. indexes, _ := c.schema.ListTableIndexes(table)
  720. rows = append(rows, map[string]interface{}{
  721. "schemaname": "public", "tablename": table, "tableowner": c.params["user"], "tablespace": nil,
  722. "hasindexes": len(indexes) > 0, "hasrules": false, "hastriggers": false, "rowsecurity": false,
  723. })
  724. }
  725. case "INFORMATION_SCHEMA.COLUMNS":
  726. for _, table := range tables {
  727. schema, schemaErr := c.schema.GetSchema(table)
  728. if schemaErr != nil {
  729. continue
  730. }
  731. for i, column := range schema.Columns {
  732. rows = append(rows, map[string]interface{}{
  733. "table_catalog": c.database, "table_schema": "public", "table_name": table,
  734. "column_name": column.Name, "ordinal_position": int64(i + 1), "column_default": column.DefaultLiteral(),
  735. "is_nullable": yesNo(column.Nullable), "data_type": strings.ToLower(column.Type),
  736. })
  737. }
  738. }
  739. case "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE":
  740. for _, table := range tables {
  741. schema, schemaErr := c.schema.GetSchema(table)
  742. if schemaErr != nil || schema.UsesSyntheticRowID() {
  743. continue
  744. }
  745. columns := schema.PrimaryKeyList()
  746. if source == "INFORMATION_SCHEMA.TABLE_CONSTRAINTS" {
  747. rows = append(rows, map[string]interface{}{
  748. "constraint_catalog": c.database, "constraint_schema": "public", "constraint_name": table + "_pkey",
  749. "table_catalog": c.database, "table_schema": "public", "table_name": table,
  750. "constraint_type": "PRIMARY KEY", "is_deferrable": "NO", "initially_deferred": "NO",
  751. })
  752. continue
  753. }
  754. // KEY_COLUMN_USAGE reports one row per key column, in key order.
  755. for i, column := range columns {
  756. rows = append(rows, map[string]interface{}{
  757. "constraint_catalog": c.database, "constraint_schema": "public", "constraint_name": table + "_pkey",
  758. "table_catalog": c.database, "table_schema": "public", "table_name": table,
  759. "column_name": column, "ordinal_position": int64(i + 1),
  760. })
  761. }
  762. }
  763. case "PG_INDEXES":
  764. for _, table := range tables {
  765. indexes, _ := c.schema.ListTableIndexes(table)
  766. for _, index := range indexes {
  767. columns := make([]string, len(index.Columns))
  768. for i, column := range index.Columns {
  769. columns[i] = column.Name
  770. }
  771. unique := ""
  772. if index.Unique {
  773. unique = "UNIQUE "
  774. }
  775. rows = append(rows, map[string]interface{}{
  776. "schemaname": "public", "tablename": table, "indexname": index.Name, "tablespace": nil,
  777. "indexdef": fmt.Sprintf("CREATE %sINDEX %s ON %s (%s)", unique, index.Name, table, strings.Join(columns, ", ")),
  778. })
  779. }
  780. }
  781. }
  782. for _, match := range catalogFilterPattern.FindAllStringSubmatch(sql, -1) {
  783. column, expected := strings.ToLower(match[1]), strings.ReplaceAll(match[2], "''", "'")
  784. filtered := rows[:0]
  785. for _, row := range rows {
  786. if value, ok := row[column]; ok && strings.EqualFold(fmt.Sprintf("%v", value), expected) {
  787. filtered = append(filtered, row)
  788. }
  789. }
  790. rows = filtered
  791. }
  792. columns := catalogProjection(sql, rows)
  793. result := executor.NewResult("SELECT")
  794. if len(columns) == 1 && columns[0] == "count(*)" {
  795. result.AddColumnWithType("count", "INTEGER")
  796. result.AddRow(int64(len(rows)))
  797. return result, true, nil
  798. }
  799. for _, column := range columns {
  800. columnType := "TEXT"
  801. if column == "ordinal_position" {
  802. columnType = "INTEGER"
  803. } else if strings.HasPrefix(column, "has") || column == "rowsecurity" {
  804. columnType = "BOOLEAN"
  805. }
  806. result.AddColumnWithType(column, columnType)
  807. }
  808. for _, row := range rows {
  809. values := make([]interface{}, len(columns))
  810. for i, column := range columns {
  811. values[i] = row[column]
  812. }
  813. result.AddRow(values...)
  814. }
  815. return result, true, nil
  816. }
  817. // isRollbackStatement reports whether sql parses as a single ROLLBACK
  818. // statement, including ROLLBACK TO SAVEPOINT.
  819. func isRollbackStatement(sql string) bool {
  820. stmt, err := parser.New(lexer.New(sql)).Parse()
  821. if err != nil {
  822. return false
  823. }
  824. _, ok := stmt.(*parser.RollbackStmt)
  825. return ok
  826. }
  827. func catalogProjection(sql string, rows []map[string]interface{}) []string {
  828. upper := strings.ToUpper(sql)
  829. selectPos, fromPos := strings.Index(upper, "SELECT"), strings.Index(upper, " FROM ")
  830. if selectPos < 0 || fromPos < 0 || fromPos <= selectPos+6 {
  831. return nil
  832. }
  833. projection := strings.TrimSpace(sql[selectPos+6 : fromPos])
  834. projection = strings.TrimSpace(strings.TrimPrefix(strings.ToUpper(projection), "DISTINCT "))
  835. if projection == "*" && len(rows) > 0 {
  836. columns := make([]string, 0, len(rows[0]))
  837. for column := range rows[0] {
  838. columns = append(columns, column)
  839. }
  840. sort.Strings(columns)
  841. return columns
  842. }
  843. parts := strings.Split(projection, ",")
  844. columns := make([]string, 0, len(parts))
  845. for _, part := range parts {
  846. column := strings.TrimSpace(part)
  847. if index := strings.Index(strings.ToUpper(column), " AS "); index >= 0 {
  848. column = strings.TrimSpace(column[:index])
  849. }
  850. if index := strings.LastIndex(column, "."); index >= 0 {
  851. column = column[index+1:]
  852. }
  853. columns = append(columns, strings.ToLower(strings.Trim(column, `"`)))
  854. }
  855. return columns
  856. }
  857. func yesNo(value bool) string {
  858. if value {
  859. return "YES"
  860. }
  861. return "NO"
  862. }
  863. func (c *Connection) sendTabularResult(result *executor.Result, tag string) error {
  864. if err := c.sendRowDescription(result.Columns, result.ColumnTypes); err != nil {
  865. return err
  866. }
  867. for _, row := range result.Rows {
  868. if err := c.sendDataRow(row, result.Columns); err != nil {
  869. return err
  870. }
  871. }
  872. return c.sendCommandComplete(fmt.Sprintf("%s %d", tag, len(result.Rows)))
  873. }
  874. func (c *Connection) handleClose(msg *Message) error {
  875. if c.extendedFailed {
  876. return nil
  877. }
  878. if len(msg.Data) < 2 {
  879. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close message"))
  880. }
  881. name, pos, err := readCString(msg.Data, 1)
  882. if err != nil || pos != len(msg.Data) {
  883. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close name"))
  884. }
  885. if msg.Data[0] == 'S' {
  886. delete(c.statements, name)
  887. } else if msg.Data[0] == 'P' {
  888. delete(c.portals, name)
  889. } else {
  890. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close target"))
  891. }
  892. return c.writeMessage(MsgCloseComplete, nil)
  893. }
  894. func (c *Connection) failExtended(code string, err error) error {
  895. c.extendedFailed = true
  896. if c.txStatus == TxStatusInBlock {
  897. c.txStatus = TxStatusFailed
  898. }
  899. return c.sendError("ERROR", code, err.Error())
  900. }
  901. func (c *Connection) handleProxyError(msg *Message) error {
  902. if len(msg.Data) != 1 || (msg.Data[0] != MsgQuery && msg.Data[0] != MsgParse) {
  903. _ = c.sendError("FATAL", ErrCodeProtocolViolation, "invalid proxy error relay")
  904. return io.EOF
  905. }
  906. code, message := "42501", "permission denied by API key scopes"
  907. if c.txStatus == TxStatusFailed {
  908. code, message = ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block"
  909. }
  910. if msg.Data[0] == MsgParse {
  911. c.extendedCycle = true
  912. return c.failExtended(code, errors.New(message))
  913. }
  914. if c.txStatus == TxStatusInBlock {
  915. c.txStatus = TxStatusFailed
  916. }
  917. if err := c.sendError("ERROR", code, message); err != nil {
  918. return err
  919. }
  920. return c.sendReadyForQuery()
  921. }
  922. type boundParameter struct {
  923. value []byte
  924. oid int32
  925. format int16
  926. null bool
  927. }
  928. func readCString(data []byte, pos int) (string, int, error) {
  929. if pos < 0 || pos >= len(data) {
  930. return "", pos, fmt.Errorf("missing null-terminated string")
  931. }
  932. end := bytes.IndexByte(data[pos:], 0)
  933. if end < 0 {
  934. return "", pos, fmt.Errorf("unterminated string")
  935. }
  936. return string(data[pos : pos+end]), pos + end + 1, nil
  937. }
  938. func readInt16List(data []byte, pos int) ([]int16, int, error) {
  939. if pos+2 > len(data) {
  940. return nil, pos, fmt.Errorf("missing list length")
  941. }
  942. count := int(binary.BigEndian.Uint16(data[pos : pos+2]))
  943. pos += 2
  944. if pos+count*2 > len(data) {
  945. return nil, pos, fmt.Errorf("truncated list")
  946. }
  947. result := make([]int16, count)
  948. for i := range result {
  949. result[i] = int16(binary.BigEndian.Uint16(data[pos : pos+2]))
  950. pos += 2
  951. }
  952. return result, pos, nil
  953. }
  954. func parameterOID(oids []int32, index int) int32 {
  955. if index < len(oids) {
  956. return oids[index]
  957. }
  958. return 0
  959. }
  960. func parameterFormat(formats []int16, index int) int16 {
  961. if len(formats) == 1 {
  962. return formats[0]
  963. }
  964. if index < len(formats) {
  965. return formats[index]
  966. }
  967. return 0
  968. }
  969. func bindQuery(query string, params []boundParameter) (string, error) {
  970. var result strings.Builder
  971. var quote byte
  972. for i := 0; i < len(query); {
  973. ch := query[i]
  974. if quote != 0 {
  975. result.WriteByte(ch)
  976. i++
  977. if ch == quote {
  978. if i < len(query) && query[i] == quote {
  979. result.WriteByte(query[i])
  980. i++
  981. } else {
  982. quote = 0
  983. }
  984. }
  985. continue
  986. }
  987. if i+1 < len(query) && ch == '-' && query[i+1] == '-' {
  988. end := strings.IndexByte(query[i:], '\n')
  989. if end < 0 {
  990. result.WriteString(query[i:])
  991. break
  992. }
  993. result.WriteString(query[i : i+end+1])
  994. i += end + 1
  995. continue
  996. }
  997. if i+1 < len(query) && ch == '/' && query[i+1] == '*' {
  998. end := strings.Index(query[i+2:], "*/")
  999. if end < 0 {
  1000. result.WriteString(query[i:])
  1001. break
  1002. }
  1003. result.WriteString(query[i : i+end+4])
  1004. i += end + 4
  1005. continue
  1006. }
  1007. if ch == '\'' || ch == '"' || ch == '`' || ch == '[' {
  1008. quote = ch
  1009. if ch == '[' {
  1010. quote = ']'
  1011. }
  1012. result.WriteByte(ch)
  1013. i++
  1014. continue
  1015. }
  1016. if ch == '$' && i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' {
  1017. end := i + 1
  1018. for end < len(query) && query[end] >= '0' && query[end] <= '9' {
  1019. end++
  1020. }
  1021. n, err := strconv.Atoi(query[i+1 : end])
  1022. if err != nil || n < 1 || n > len(params) {
  1023. return "", fmt.Errorf("parameter %s was not provided", query[i:end])
  1024. }
  1025. literal, err := parameterLiteral(params[n-1])
  1026. if err != nil {
  1027. return "", fmt.Errorf("parameter $%d: %w", n, err)
  1028. }
  1029. result.WriteString(literal)
  1030. i = end
  1031. continue
  1032. }
  1033. result.WriteByte(ch)
  1034. i++
  1035. }
  1036. return result.String(), nil
  1037. }
  1038. func parameterLiteral(param boundParameter) (string, error) {
  1039. if param.null {
  1040. return "NULL", nil
  1041. }
  1042. if param.format == 1 {
  1043. switch param.oid {
  1044. case 16:
  1045. if len(param.value) != 1 {
  1046. return "", fmt.Errorf("invalid binary boolean")
  1047. }
  1048. if param.value[0] == 0 {
  1049. return "FALSE", nil
  1050. }
  1051. return "TRUE", nil
  1052. case 21:
  1053. if len(param.value) != 2 {
  1054. return "", fmt.Errorf("invalid binary int2")
  1055. }
  1056. return strconv.FormatInt(int64(int16(binary.BigEndian.Uint16(param.value))), 10), nil
  1057. case 23:
  1058. if len(param.value) != 4 {
  1059. return "", fmt.Errorf("invalid binary int4")
  1060. }
  1061. return strconv.FormatInt(int64(int32(binary.BigEndian.Uint32(param.value))), 10), nil
  1062. case 20:
  1063. if len(param.value) != 8 {
  1064. return "", fmt.Errorf("invalid binary int8")
  1065. }
  1066. return strconv.FormatInt(int64(binary.BigEndian.Uint64(param.value)), 10), nil
  1067. default:
  1068. return "", fmt.Errorf("%w: binary parameter format for OID %d", pgerror.ErrFeatureNotSupported, param.oid)
  1069. }
  1070. }
  1071. value := string(param.value)
  1072. switch param.oid {
  1073. case 0:
  1074. if strings.EqualFold(value, "true") || strings.EqualFold(value, "false") {
  1075. return strings.ToUpper(value), nil
  1076. }
  1077. if _, err := strconv.ParseFloat(value, 64); err == nil && value != "" {
  1078. return value, nil
  1079. }
  1080. return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil
  1081. case 16:
  1082. if strings.EqualFold(value, "true") || value == "1" || value == "t" {
  1083. return "TRUE", nil
  1084. }
  1085. return "FALSE", nil
  1086. case 20, 21, 23, 26, 700, 701, 1700:
  1087. if _, err := strconv.ParseFloat(value, 64); err != nil {
  1088. return "", fmt.Errorf("invalid numeric value")
  1089. }
  1090. return value, nil
  1091. default:
  1092. return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil
  1093. }
  1094. }
  1095. // sendResult sends query results
  1096. func (c *Connection) sendResult(result *executor.Result, stmt parser.Statement) error {
  1097. // Any statement that produced a row set (SELECT, or INSERT/UPDATE/DELETE
  1098. // with RETURNING) sends a row description and its data rows. Other
  1099. // statements only send the command completion tag.
  1100. if len(result.Columns) > 0 {
  1101. if err := c.sendRowDescription(result.Columns, wireColumnTypes(result)); err != nil {
  1102. return err
  1103. }
  1104. for _, row := range result.Rows {
  1105. if err := c.sendDataRow(row, result.Columns); err != nil {
  1106. return err
  1107. }
  1108. }
  1109. return c.sendCommandComplete(c.getCommandTag(stmt, result))
  1110. }
  1111. // For other statements, just send command complete
  1112. tag := c.getCommandTag(stmt, result)
  1113. return c.sendCommandComplete(tag)
  1114. }
  1115. // wireColumnTypes corrects declared-affinity metadata when SQLite has stored a
  1116. // value of another type. In particular, binding []byte to a VARCHAR column is
  1117. // still a BLOB in SQLite; advertising TEXT while encoding PostgreSQL bytea text
  1118. // would make clients receive the literal "\\x..." instead of the original
  1119. // bytes.
  1120. func wireColumnTypes(result *executor.Result) []string {
  1121. types := append([]string(nil), result.ColumnTypes...)
  1122. if len(types) < len(result.Columns) {
  1123. types = append(types, make([]string, len(result.Columns)-len(types))...)
  1124. }
  1125. for _, row := range result.Rows {
  1126. for i, value := range row {
  1127. if i >= len(types) {
  1128. break
  1129. }
  1130. if _, ok := value.([]byte); ok {
  1131. types[i] = "BLOB"
  1132. }
  1133. }
  1134. }
  1135. return types
  1136. }
  1137. // getCommandTag returns the command completion tag
  1138. func (c *Connection) getCommandTag(stmt parser.Statement, result *executor.Result) string {
  1139. switch s := stmt.(type) {
  1140. case *parser.SelectStmt:
  1141. return fmt.Sprintf("SELECT %d", len(result.Rows))
  1142. case *parser.CreateTableStmt:
  1143. return "CREATE TABLE"
  1144. case *parser.DropTableStmt:
  1145. return "DROP TABLE"
  1146. case *parser.CreateIndexStmt:
  1147. return "CREATE INDEX"
  1148. case *parser.DropIndexStmt:
  1149. return "DROP INDEX"
  1150. case *parser.AlterTableStmt:
  1151. return "ALTER TABLE"
  1152. case *parser.InsertStmt:
  1153. return fmt.Sprintf("INSERT 0 %d", result.RowsAffected)
  1154. case *parser.UpdateStmt:
  1155. return fmt.Sprintf("UPDATE %d", result.RowsAffected)
  1156. case *parser.DeleteStmt:
  1157. return fmt.Sprintf("DELETE %d", result.RowsAffected)
  1158. case *parser.BeginStmt:
  1159. c.txStatus = TxStatusInBlock
  1160. return "BEGIN"
  1161. case *parser.CommitStmt:
  1162. c.txStatus = TxStatusIdle
  1163. return "COMMIT"
  1164. case *parser.RollbackStmt:
  1165. if s.Savepoint != "" {
  1166. c.txStatus = TxStatusInBlock
  1167. return "ROLLBACK"
  1168. }
  1169. c.txStatus = TxStatusIdle
  1170. return "ROLLBACK"
  1171. case *parser.SavepointStmt:
  1172. c.txStatus = TxStatusInBlock
  1173. return "SAVEPOINT"
  1174. case *parser.ReleaseStmt:
  1175. c.txStatus = TxStatusInBlock
  1176. return "RELEASE"
  1177. case *parser.ListenStmt:
  1178. return "LISTEN"
  1179. case *parser.UnlistenStmt:
  1180. return "UNLISTEN"
  1181. case *parser.NotifyStmt:
  1182. return "NOTIFY"
  1183. default:
  1184. return "OK"
  1185. }
  1186. }
  1187. // sendAuthenticationOk sends authentication OK message
  1188. func (c *Connection) sendAuthenticationOk() error {
  1189. mb := NewMessageBuilder()
  1190. mb.WriteInt32(0) // Auth OK
  1191. return c.writeMessage(MsgAuthenticationOk, mb.Bytes())
  1192. }
  1193. // drainNotifications writes every notification currently buffered for this
  1194. // session. It is called at protocol-safe points (after Sync/ReadyForQuery).
  1195. func (c *Connection) drainNotifications() error {
  1196. ch := c.executor.NotificationChannel()
  1197. for {
  1198. select {
  1199. case n := <-ch:
  1200. if err := c.sendNotification(n); err != nil {
  1201. return err
  1202. }
  1203. default:
  1204. return c.writer.Flush()
  1205. }
  1206. }
  1207. }
  1208. // sendNotification writes a NotificationResponse frame. The frame is only
  1209. // written by the connection's main goroutine, never from a publisher.
  1210. func (c *Connection) sendNotification(n storage.Notification) error {
  1211. mb := NewMessageBuilder()
  1212. mb.WriteInt32(n.PID)
  1213. mb.WriteString(n.Channel)
  1214. mb.WriteString(n.Payload)
  1215. return c.writeMessage(MsgNotificationResponse, mb.Bytes())
  1216. }
  1217. // sendLSN exposes the last commit LSN as a ParameterStatus. Drivers that do
  1218. // not know the key ignore it; subscribers can use it as an ordering cursor.
  1219. func (c *Connection) sendLSN(lsn uint64) error {
  1220. return c.sendParameterStatus("pizzasql.lsn", strconv.FormatUint(lsn, 10))
  1221. }
  1222. // sendParameterStatus sends a parameter status message
  1223. func (c *Connection) sendParameterStatus(name, value string) error {
  1224. mb := NewMessageBuilder()
  1225. mb.WriteString(name)
  1226. mb.WriteString(value)
  1227. return c.writeMessage(MsgParameterStatus, mb.Bytes())
  1228. }
  1229. // sendBackendKeyData sends backend key data
  1230. func (c *Connection) sendBackendKeyData(processID, secretKey int32) error {
  1231. mb := NewMessageBuilder()
  1232. mb.WriteInt32(processID)
  1233. mb.WriteInt32(secretKey)
  1234. return c.writeMessage(MsgBackendKeyData, mb.Bytes())
  1235. }
  1236. // isIdleRollback reports whether a bare ROLLBACK found no transaction to roll
  1237. // back. A COMMIT rejected by storage has already ended the transaction, so the
  1238. // client's cleanup ROLLBACK must complete (PostgreSQL reports a warning and
  1239. // succeeds) instead of surfacing an internal error on an idle session.
  1240. func isIdleRollback(stmt parser.Statement, status byte, err error) bool {
  1241. rollback, ok := stmt.(*parser.RollbackStmt)
  1242. if !ok || rollback.Savepoint != "" {
  1243. return false
  1244. }
  1245. return status == TxStatusIdle && errors.Is(err, executor.ErrNoTransaction)
  1246. }
  1247. // sendReadyForQuery sends ready for query message
  1248. func (c *Connection) sendReadyForQuery() error {
  1249. mb := NewMessageBuilder()
  1250. mb.AppendByte(c.txStatus)
  1251. if err := c.writeMessage(MsgReadyForQuery, mb.Bytes()); err != nil {
  1252. return err
  1253. }
  1254. // ReadyForQuery closes out a response cycle, so flush everything buffered
  1255. // so far. This is what makes simple-query results and Sync responses
  1256. // visible to the client.
  1257. return c.writer.Flush()
  1258. }
  1259. // sendEmptyQueryResponse sends empty query response
  1260. func (c *Connection) sendEmptyQueryResponse() error {
  1261. return c.writeMessage(MsgEmptyQueryResponse, []byte{})
  1262. }
  1263. // sendRowDescription sends row description (column metadata)
  1264. func (c *Connection) sendRowDescription(columns []string, columnTypes []string) error {
  1265. mb := NewMessageBuilder()
  1266. mb.WriteInt16(int16(len(columns)))
  1267. for i, col := range columns {
  1268. colType := ""
  1269. if i < len(columnTypes) {
  1270. colType = columnTypes[i]
  1271. }
  1272. mb.WriteString(col)
  1273. mb.WriteInt32(0) // table OID
  1274. mb.WriteInt16(0) // column attribute number
  1275. mb.WriteInt32(c.getOIDForType(colType)) // type OID
  1276. mb.WriteInt16(c.getTypeSizeForType(colType)) // type size
  1277. mb.WriteInt32(-1) // type modifier
  1278. mb.WriteInt16(0) // format code (text)
  1279. }
  1280. return c.writeMessage(MsgRowDescription, mb.Bytes())
  1281. }
  1282. // sendDataRow sends a data row
  1283. func (c *Connection) sendDataRow(row []interface{}, columns []string) error {
  1284. mb := NewMessageBuilder()
  1285. mb.WriteInt16(int16(len(row)))
  1286. for _, value := range row {
  1287. if value == nil {
  1288. mb.WriteInt32(-1) // NULL indicator
  1289. continue
  1290. }
  1291. // Convert value to string
  1292. strValue := c.valueToString(value)
  1293. mb.WriteInt32(int32(len(strValue)))
  1294. mb.WriteBytes([]byte(strValue))
  1295. }
  1296. return c.writeMessage(MsgDataRow, mb.Bytes())
  1297. }
  1298. // sendCommandComplete sends command complete message
  1299. func (c *Connection) sendCommandComplete(tag string) error {
  1300. mb := NewMessageBuilder()
  1301. mb.WriteString(tag)
  1302. return c.writeMessage(MsgCommandComplete, mb.Bytes())
  1303. }
  1304. // sendError sends an error response
  1305. func (c *Connection) sendError(severity, code, message string) error {
  1306. if severity == "ERROR" && c.txStatus == TxStatusInBlock {
  1307. c.txStatus = TxStatusFailed
  1308. }
  1309. mb := NewMessageBuilder()
  1310. mb.AppendByte(ErrorFieldSeverity)
  1311. mb.WriteString(severity)
  1312. mb.AppendByte(ErrorFieldCode)
  1313. mb.WriteString(code)
  1314. mb.AppendByte(ErrorFieldMessage)
  1315. mb.WriteString(message)
  1316. mb.AppendByte(0) // Terminator
  1317. if err := c.writeMessage(MsgErrorResponse, mb.Bytes()); err != nil {
  1318. return err
  1319. }
  1320. // Errors must be visible promptly, including FATAL startup failures that
  1321. // are not followed by a ReadyForQuery before the connection closes.
  1322. return c.writer.Flush()
  1323. }
  1324. // writeMessage buffers a message for the connection. It does not flush, so
  1325. // callers that need to make a response visible to the client must flush at the
  1326. // appropriate protocol boundary (ReadyForQuery, an explicit Flush message, or
  1327. // an error response). Buffering amortizes the per-message syscalls that a
  1328. // result set would otherwise incur; the underlying bufio.Writer bounds memory
  1329. // use so large result sets cannot grow the buffer without limit.
  1330. func (c *Connection) writeMessage(msgType byte, data []byte) error {
  1331. if !c.quiet {
  1332. log.Printf("Sending message type=%c length=%d", msgType, len(data)+4)
  1333. }
  1334. return WriteMessage(c.writer, msgType, data)
  1335. }
  1336. // getOIDForType returns PostgreSQL OID for type.
  1337. //
  1338. // Types narrower than the value the engine carries are reported as the carrier
  1339. // type, because the engine does not range-check stored or cast values: a value
  1340. // that does not fit int2/int4 or float4 would fail client-side decoding rather
  1341. // than having been rejected on write. INTEGER and REAL keep their declared
  1342. // mappings, which predate this rule and are what existing clients expect.
  1343. func (c *Connection) getOIDForType(typeName string) int32 {
  1344. switch strings.ToUpper(typeName) {
  1345. case "SMALLINT", "INT2", "TINYINT", "MEDIUMINT":
  1346. // Carried as int64 and not range-checked, so int8 is the honest type.
  1347. return 20 // INT8OID
  1348. case "INTEGER", "INT", "INT4":
  1349. return 23 // INT4OID
  1350. case "BIGINT", "INT8":
  1351. return 20 // INT8OID
  1352. case "NUMERIC", "DECIMAL":
  1353. // Values round-trip through JSON as float64; there is no exact decimal
  1354. // carrier, so NUMERIC is not advertised.
  1355. return 701 // FLOAT8OID
  1356. case "TEXT", "VARCHAR", "CHAR", "UUID":
  1357. // UUID is stored as text (SQLite-style), so it reports TEXTOID rather
  1358. // than the native PG UUID OID.
  1359. return 25 // TEXTOID
  1360. case "REAL", "FLOAT":
  1361. return 700 // FLOAT4OID
  1362. case "DOUBLE":
  1363. return 701 // FLOAT8OID
  1364. case "BOOLEAN", "BOOL":
  1365. return 16 // BOOLOID
  1366. case "BLOB":
  1367. return 17 // BYTEAOID
  1368. case "DATETIME", "TIMESTAMP":
  1369. // Datetime values are exchanged as UTC RFC3339 (ISO8601 with a timezone
  1370. // offset), which timestamptz decodes directly.
  1371. return 1184 // TIMESTAMPTZOID
  1372. case "DATE":
  1373. return 1082 // DATEOID
  1374. default:
  1375. return 25 // Default to TEXT
  1376. }
  1377. }
  1378. // getTypeSizeForType returns type size in bytes, or -1 for variable length.
  1379. func (c *Connection) getTypeSizeForType(typeName string) int16 {
  1380. switch strings.ToUpper(typeName) {
  1381. case "SMALLINT", "INT2", "TINYINT", "MEDIUMINT":
  1382. return 8 // Reported as int8; see getOIDForType.
  1383. case "INTEGER", "INT", "INT4":
  1384. return 4
  1385. case "BIGINT", "INT8":
  1386. return 8
  1387. case "NUMERIC", "DECIMAL":
  1388. return 8 // Reported as float8; see getOIDForType.
  1389. case "REAL", "FLOAT":
  1390. return 4
  1391. case "DOUBLE":
  1392. return 8
  1393. case "BOOLEAN", "BOOL":
  1394. return 1
  1395. case "DATETIME", "TIMESTAMP":
  1396. return 8
  1397. case "DATE":
  1398. return 4
  1399. default:
  1400. return -1 // Variable length
  1401. }
  1402. }
  1403. // valueToString converts a value to string
  1404. func (c *Connection) valueToString(value interface{}) string {
  1405. if value == nil {
  1406. return ""
  1407. }
  1408. // PostgreSQL's text format for bytea is \x followed by hex. Emitting that
  1409. // lets libpq/pgx decode a BLOB column (OID 17) losslessly.
  1410. if b, ok := value.([]byte); ok {
  1411. return `\x` + hex.EncodeToString(b)
  1412. }
  1413. return fmt.Sprintf("%v", value)
  1414. }
  1415. // handleVersionQuery handles SELECT version()
  1416. func (c *Connection) handleVersionQuery() error {
  1417. columns := []string{"version"}
  1418. columnTypes := []string{"TEXT"}
  1419. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1420. return err
  1421. }
  1422. row := []interface{}{"PostgreSQL " + version.PostgreSQLCompatVersion + " (PizzaSQL)"}
  1423. if err := c.sendDataRow(row, columns); err != nil {
  1424. return err
  1425. }
  1426. if err := c.sendCommandComplete("SELECT 1"); err != nil {
  1427. return err
  1428. }
  1429. return c.sendReadyForQuery()
  1430. }
  1431. // handleCurrentUserQuery handles SELECT current_user
  1432. func (c *Connection) handleCurrentUserQuery() error {
  1433. columns := []string{"current_user"}
  1434. columnTypes := []string{"TEXT"}
  1435. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1436. return err
  1437. }
  1438. user := c.params["user"]
  1439. if user == "" {
  1440. user = "pizzasql"
  1441. }
  1442. row := []interface{}{user}
  1443. if err := c.sendDataRow(row, columns); err != nil {
  1444. return err
  1445. }
  1446. if err := c.sendCommandComplete("SELECT 1"); err != nil {
  1447. return err
  1448. }
  1449. return c.sendReadyForQuery()
  1450. }
  1451. // handleShowCommand handles SHOW commands
  1452. func (c *Connection) handleShowCommand(sqlUpper string) error {
  1453. var value string
  1454. var name string
  1455. if sqlUpper == "SERVER_VERSION" {
  1456. name = "server_version"
  1457. value = version.PostgreSQLCompatVersion
  1458. } else if sqlUpper == "SERVER_ENCODING" {
  1459. name = "server_encoding"
  1460. value = "UTF8"
  1461. } else if sqlUpper == "CLIENT_ENCODING" {
  1462. name = "client_encoding"
  1463. value = "UTF8"
  1464. } else {
  1465. // Unknown SHOW command
  1466. c.sendError("ERROR", ErrCodeFeatureNotSupported, "SHOW command not supported")
  1467. return c.sendReadyForQuery()
  1468. }
  1469. columns := []string{name}
  1470. columnTypes := []string{"TEXT"}
  1471. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1472. return err
  1473. }
  1474. row := []interface{}{value}
  1475. if err := c.sendDataRow(row, columns); err != nil {
  1476. return err
  1477. }
  1478. if err := c.sendCommandComplete("SHOW"); err != nil {
  1479. return err
  1480. }
  1481. return c.sendReadyForQuery()
  1482. }