connection.go 64 KB

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