2
0

connection.go 44 KB

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