connection.go 52 KB

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