2
0

connection.go 35 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220
  1. package pgserver
  2. import (
  3. "bufio"
  4. "bytes"
  5. "encoding/binary"
  6. "fmt"
  7. "io"
  8. "log"
  9. "net"
  10. "regexp"
  11. "sort"
  12. "strconv"
  13. "strings"
  14. "github.com/danfragoso/pizzasql-next/pkg/executor"
  15. "github.com/danfragoso/pizzasql-next/pkg/lexer"
  16. "github.com/danfragoso/pizzasql-next/pkg/parser"
  17. "github.com/danfragoso/pizzasql-next/pkg/storage"
  18. )
  19. // Connection represents a client connection
  20. type Connection struct {
  21. conn net.Conn
  22. reader *bufio.Reader
  23. writer *bufio.Writer
  24. executor *executor.Executor
  25. schema *storage.SchemaManager
  26. dbManager *storage.DatabaseManager
  27. database string
  28. params map[string]string
  29. txStatus byte
  30. quiet bool // Disable query logging
  31. statements map[string]*preparedStatement
  32. portals map[string]*portal
  33. extendedFailed bool
  34. }
  35. type preparedStatement struct {
  36. query string
  37. paramOIDs []int32
  38. }
  39. type portal struct {
  40. query string
  41. executed bool
  42. }
  43. // NewConnection creates a new connection handler
  44. func NewConnection(conn net.Conn, dbManager *storage.DatabaseManager, quiet bool) *Connection {
  45. return &Connection{
  46. conn: conn,
  47. reader: bufio.NewReader(conn),
  48. writer: bufio.NewWriter(conn),
  49. dbManager: dbManager,
  50. params: make(map[string]string),
  51. statements: make(map[string]*preparedStatement),
  52. portals: make(map[string]*portal),
  53. txStatus: TxStatusIdle,
  54. quiet: quiet,
  55. }
  56. }
  57. // Handle processes the connection
  58. func (c *Connection) Handle() error {
  59. defer func() {
  60. if c.executor != nil {
  61. if err := c.executor.RollbackActive(); err != nil && !c.quiet {
  62. log.Printf("failed to roll back disconnected transaction: %v", err)
  63. }
  64. }
  65. c.conn.Close()
  66. }()
  67. // First, check for SSL request (sent before startup message)
  68. // SSL request is 8 bytes: length(4) + code(4) where code = 80877103
  69. firstBytes := make([]byte, 8)
  70. n, err := io.ReadFull(c.reader, firstBytes)
  71. if err != nil {
  72. return fmt.Errorf("failed to read initial bytes: %w", err)
  73. }
  74. // Check if it's an SSL request (code 80877103 = 0x04D2162F)
  75. if n == 8 {
  76. length := binary.BigEndian.Uint32(firstBytes[0:4])
  77. code := binary.BigEndian.Uint32(firstBytes[4:8])
  78. if length == 8 && code == 80877103 {
  79. // SSL request - we don't support SSL, send 'N'
  80. if !c.quiet {
  81. log.Printf("Client requested SSL, sending rejection")
  82. }
  83. if _, err := c.conn.Write([]byte{'N'}); err != nil {
  84. return fmt.Errorf("failed to send SSL rejection: %w", err)
  85. }
  86. // Now read the actual startup message
  87. } else {
  88. // Not SSL request, this is part of startup message
  89. // We need to prepend these bytes back for ReadStartupMessage
  90. // Create a multi-reader that first reads our buffered bytes, then continues with the reader
  91. c.reader = bufio.NewReader(io.MultiReader(bytes.NewReader(firstBytes), c.reader))
  92. }
  93. }
  94. // Read startup message
  95. if !c.quiet {
  96. log.Printf("Reading startup message...")
  97. }
  98. params, err := ReadStartupMessage(c.reader)
  99. if err != nil {
  100. return fmt.Errorf("failed to read startup message: %w", err)
  101. }
  102. c.params = params
  103. if !c.quiet {
  104. log.Printf("Startup params: %+v", params)
  105. }
  106. // Get database name from params (default to "pizzasql")
  107. dbName := params["database"]
  108. if dbName == "" {
  109. dbName = "pizzasql"
  110. }
  111. c.database = dbName
  112. if !c.quiet {
  113. log.Printf("New connection: user=%s database=%s", params["user"], dbName)
  114. }
  115. // Initialize database
  116. if err := c.initDatabase(dbName); err != nil {
  117. c.sendError("FATAL", ErrCodeConnectionFailure, fmt.Sprintf("Failed to initialize database: %v", err))
  118. return err
  119. }
  120. // Send authentication OK (no auth for now)
  121. if err := c.sendAuthenticationOk(); err != nil {
  122. return err
  123. }
  124. // Send parameter status messages
  125. if err := c.sendParameterStatus("server_version", "14.0 (PizzaSQL)"); err != nil {
  126. return err
  127. }
  128. if err := c.sendParameterStatus("server_encoding", "UTF8"); err != nil {
  129. return err
  130. }
  131. if err := c.sendParameterStatus("client_encoding", "UTF8"); err != nil {
  132. return err
  133. }
  134. if err := c.sendParameterStatus("DateStyle", "ISO, MDY"); err != nil {
  135. return err
  136. }
  137. if err := c.sendParameterStatus("TimeZone", "UTC"); err != nil {
  138. return err
  139. }
  140. // Send backend key data (for cancellation - we don't implement this yet)
  141. if err := c.sendBackendKeyData(12345, 67890); err != nil {
  142. return err
  143. }
  144. // Send ready for query
  145. if err := c.sendReadyForQuery(); err != nil {
  146. return err
  147. }
  148. // Message loop
  149. for {
  150. msg, err := ReadMessage(c.reader)
  151. if err != nil {
  152. if err == io.EOF {
  153. log.Printf("Connection closed by client")
  154. return nil
  155. }
  156. return fmt.Errorf("failed to read message: %w", err)
  157. }
  158. if err := c.handleMessage(msg); err != nil {
  159. if err == io.EOF {
  160. // Normal termination
  161. log.Printf("Connection closed normally")
  162. return nil
  163. }
  164. log.Printf("Error handling message: %v", err)
  165. return err
  166. }
  167. }
  168. }
  169. // initDatabase initializes the database connection
  170. func (c *Connection) initDatabase(dbName string) error {
  171. db, err := c.dbManager.GetDatabase(dbName)
  172. if err != nil {
  173. return err
  174. }
  175. c.schema = db.Schema
  176. c.executor = executor.New(db.Schema, db.Table)
  177. c.executor.SyncCatalog()
  178. return nil
  179. }
  180. // handleMessage processes a client message
  181. func (c *Connection) handleMessage(msg *Message) error {
  182. switch msg.Type {
  183. case MsgQuery:
  184. return c.handleQuery(msg)
  185. case MsgTerminate:
  186. log.Printf("Client requested termination")
  187. return io.EOF
  188. case MsgParse:
  189. return c.handleParse(msg)
  190. case MsgBind:
  191. return c.handleBind(msg)
  192. case MsgDescribe:
  193. return c.handleDescribe(msg)
  194. case MsgExecute:
  195. return c.handleExecute(msg)
  196. case MsgClose:
  197. return c.handleClose(msg)
  198. case MsgFlush:
  199. return c.writer.Flush()
  200. case MsgSync:
  201. c.extendedFailed = false
  202. return c.sendReadyForQuery()
  203. default:
  204. log.Printf("Unknown message type: %c (%d)", msg.Type, msg.Type)
  205. c.sendError("ERROR", ErrCodeProtocolViolation, fmt.Sprintf("Unknown message type: %c", msg.Type))
  206. return c.sendReadyForQuery()
  207. }
  208. }
  209. // handleQuery processes a simple query
  210. func (c *Connection) handleQuery(msg *Message) error {
  211. // A Query message carries a single NUL-terminated query string.
  212. if len(msg.Data) == 0 || msg.Data[len(msg.Data)-1] != 0 {
  213. c.sendError("ERROR", ErrCodeProtocolViolation, "invalid Query message: missing null terminator")
  214. return c.sendReadyForQuery()
  215. }
  216. sql := string(msg.Data[:len(msg.Data)-1])
  217. if !c.quiet {
  218. log.Printf("Query: %s", sql)
  219. }
  220. // Handle empty query
  221. sqlTrimmed := strings.TrimSpace(sql)
  222. if sqlTrimmed == "" || sqlTrimmed == ";" {
  223. if err := c.sendEmptyQueryResponse(); err != nil {
  224. return err
  225. }
  226. return c.sendReadyForQuery()
  227. }
  228. // Handle special PostgreSQL system queries that drivers send
  229. sqlUpper := strings.ToUpper(strings.TrimSpace(sql))
  230. singleStatement := isSingleStatementSQL(sql)
  231. // In a failed transaction only ROLLBACK (or ROLLBACK TO SAVEPOINT) is
  232. // accepted. Gate before honoring the special driver queries below so they
  233. // return 25P02 instead of a result. Multi-statement batches are gated
  234. // per statement in the execution loop below.
  235. if c.txStatus == TxStatusFailed && singleStatement && !isRollbackStatement(sql) {
  236. c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block")
  237. return c.sendReadyForQuery()
  238. }
  239. // lib/pq and other drivers query these for connection validation
  240. if singleStatement && strings.Contains(sqlUpper, "SELECT VERSION()") {
  241. // Return a fake PostgreSQL version
  242. return c.handleVersionQuery()
  243. }
  244. if singleStatement && strings.Contains(sqlUpper, "SELECT CURRENT_USER") {
  245. // Return the current user
  246. return c.handleCurrentUserQuery()
  247. }
  248. if singleStatement && strings.Contains(sqlUpper, "SHOW") && (strings.Contains(sqlUpper, "SERVER_VERSION") ||
  249. strings.Contains(sqlUpper, "SERVER_ENCODING") ||
  250. strings.Contains(sqlUpper, "CLIENT_ENCODING")) {
  251. // Handle SHOW commands
  252. return c.handleShowCommand(sqlUpper)
  253. }
  254. if result, handled, err := c.catalogResult(sql); handled {
  255. if err != nil {
  256. c.sendError("ERROR", ErrCodeInternalError, err.Error())
  257. return c.sendReadyForQuery()
  258. }
  259. if err := c.sendTabularResult(result, "SELECT"); err != nil {
  260. return err
  261. }
  262. return c.sendReadyForQuery()
  263. }
  264. // Parse the complete batch before executing its first statement. This keeps
  265. // unsupported trailing clauses from turning into committed writes.
  266. l := lexer.New(sql)
  267. p := parser.New(l)
  268. stmts, err := p.ParseMultiple()
  269. if err != nil {
  270. c.sendError("ERROR", ErrCodeSyntaxError, fmt.Sprintf("Syntax error: %v", err))
  271. return c.sendReadyForQuery()
  272. }
  273. for _, stmt := range stmts {
  274. if c.txStatus == TxStatusFailed {
  275. if _, rollback := stmt.(*parser.RollbackStmt); !rollback {
  276. c.sendError("ERROR", ErrCodeTransactionAborted, "current transaction is aborted, commands ignored until end of transaction block")
  277. return c.sendReadyForQuery()
  278. }
  279. }
  280. result, err := c.executor.Execute(stmt)
  281. if err != nil {
  282. if c.txStatus == TxStatusInBlock {
  283. c.txStatus = TxStatusFailed
  284. }
  285. c.sendError("ERROR", ErrCodeInternalError, fmt.Sprintf("Execution error: %v", err))
  286. return c.sendReadyForQuery()
  287. }
  288. if err := c.sendResult(result, stmt); err != nil {
  289. return err
  290. }
  291. }
  292. return c.sendReadyForQuery()
  293. }
  294. func (c *Connection) handleParse(msg *Message) error {
  295. if c.extendedFailed {
  296. return nil
  297. }
  298. name, pos, err := readCString(msg.Data, 0)
  299. if err != nil {
  300. return c.failExtended(ErrCodeProtocolViolation, err)
  301. }
  302. query, pos, err := readCString(msg.Data, pos)
  303. if err != nil || pos+2 > len(msg.Data) {
  304. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse message"))
  305. }
  306. count := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2]))
  307. pos += 2
  308. if pos+count*4 != len(msg.Data) {
  309. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Parse parameter list"))
  310. }
  311. oids := make([]int32, count)
  312. for i := range oids {
  313. oids[i] = int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
  314. pos += 4
  315. }
  316. c.statements[name] = &preparedStatement{query: query, paramOIDs: oids}
  317. return c.writeMessage(MsgParseComplete, nil)
  318. }
  319. func (c *Connection) handleBind(msg *Message) error {
  320. if c.extendedFailed {
  321. return nil
  322. }
  323. portalName, pos, err := readCString(msg.Data, 0)
  324. if err != nil {
  325. return c.failExtended(ErrCodeProtocolViolation, err)
  326. }
  327. statementName, pos, err := readCString(msg.Data, pos)
  328. if err != nil {
  329. return c.failExtended(ErrCodeProtocolViolation, err)
  330. }
  331. statement, ok := c.statements[statementName]
  332. if !ok {
  333. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", statementName))
  334. }
  335. formats, pos, err := readInt16List(msg.Data, pos)
  336. if err != nil || pos+2 > len(msg.Data) {
  337. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind format list"))
  338. }
  339. paramCount := int(binary.BigEndian.Uint16(msg.Data[pos : pos+2]))
  340. pos += 2
  341. params := make([]boundParameter, paramCount)
  342. for i := range params {
  343. if pos+4 > len(msg.Data) {
  344. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter"))
  345. }
  346. length := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
  347. pos += 4
  348. params[i].oid = parameterOID(statement.paramOIDs, i)
  349. params[i].format = parameterFormat(formats, i)
  350. if length == -1 {
  351. params[i].null = true
  352. continue
  353. }
  354. if length < 0 || pos+int(length) > len(msg.Data) {
  355. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind parameter length"))
  356. }
  357. params[i].value = append([]byte(nil), msg.Data[pos:pos+int(length)]...)
  358. pos += int(length)
  359. }
  360. _, pos, err = readInt16List(msg.Data, pos) // result formats; text output is currently used
  361. if err != nil || pos != len(msg.Data) {
  362. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind result format list"))
  363. }
  364. query, err := bindQuery(statement.query, params)
  365. if err != nil {
  366. return c.failExtended(ErrCodeInvalidParameter, err)
  367. }
  368. c.portals[portalName] = &portal{query: query}
  369. return c.writeMessage(MsgBindComplete, nil)
  370. }
  371. func (c *Connection) handleDescribe(msg *Message) error {
  372. if c.extendedFailed {
  373. return nil
  374. }
  375. if len(msg.Data) < 2 {
  376. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe message"))
  377. }
  378. name, pos, err := readCString(msg.Data, 1)
  379. if err != nil || pos != len(msg.Data) {
  380. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe name"))
  381. }
  382. switch msg.Data[0] {
  383. case 'S':
  384. statement, ok := c.statements[name]
  385. if !ok {
  386. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", name))
  387. }
  388. mb := NewMessageBuilder()
  389. mb.WriteInt16(int16(len(statement.paramOIDs)))
  390. for _, oid := range statement.paramOIDs {
  391. mb.WriteInt32(oid)
  392. }
  393. if err := c.writeMessage(MsgParameterDescription, mb.Bytes()); err != nil {
  394. return err
  395. }
  396. case 'P':
  397. if _, ok := c.portals[name]; !ok {
  398. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
  399. }
  400. default:
  401. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Describe target"))
  402. }
  403. // Execute sends the row description once the bound statement has been
  404. // analyzed, avoiding side effects during Describe.
  405. return c.writeMessage(MsgNoData, nil)
  406. }
  407. func (c *Connection) handleExecute(msg *Message) error {
  408. if c.extendedFailed {
  409. return nil
  410. }
  411. name, pos, err := readCString(msg.Data, 0)
  412. if err != nil || pos+4 != len(msg.Data) {
  413. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute message"))
  414. }
  415. portal, ok := c.portals[name]
  416. if !ok {
  417. return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
  418. }
  419. if portal.executed {
  420. return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("portal can only be executed once"))
  421. }
  422. portal.executed = true
  423. // Check the transaction state before catalog emulation. Catalog queries do
  424. // not all parse as regular PizzaSQL statements, but must still return 25P02
  425. // while the transaction is aborted.
  426. if c.txStatus == TxStatusFailed && !isRollbackStatement(portal.query) {
  427. return c.failExtended(ErrCodeTransactionAborted, fmt.Errorf("current transaction is aborted, commands ignored until end of transaction block"))
  428. }
  429. if result, handled, catalogErr := c.catalogResult(portal.query); handled {
  430. if catalogErr != nil {
  431. return c.failExtended(ErrCodeInternalError, catalogErr)
  432. }
  433. return c.sendTabularResult(result, "SELECT")
  434. }
  435. l := lexer.New(portal.query)
  436. stmt, err := parser.New(l).Parse()
  437. if err != nil {
  438. return c.failExtended(ErrCodeSyntaxError, fmt.Errorf("syntax error: %w", err))
  439. }
  440. result, err := c.executor.Execute(stmt)
  441. if err != nil {
  442. if c.txStatus == TxStatusInBlock {
  443. c.txStatus = TxStatusFailed
  444. }
  445. return c.failExtended(ErrCodeInternalError, fmt.Errorf("execution error: %w", err))
  446. }
  447. return c.sendResult(result, stmt)
  448. }
  449. var catalogFilterPattern = regexp.MustCompile(`(?i)\b(table_name|tablename|table_schema|schemaname|constraint_name|indexname)\s*=\s*'((?:''|[^'])*)'`)
  450. func (c *Connection) catalogResult(sql string) (*executor.Result, bool, error) {
  451. if !isSingleStatementSQL(sql) {
  452. return nil, false, nil
  453. }
  454. upper := strings.ToUpper(sql)
  455. var source string
  456. for _, candidate := range []string{
  457. "INFORMATION_SCHEMA.TABLES", "INFORMATION_SCHEMA.COLUMNS",
  458. "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE",
  459. "PG_TABLES", "PG_INDEXES",
  460. } {
  461. if strings.Contains(upper, "FROM "+candidate) {
  462. source = candidate
  463. break
  464. }
  465. }
  466. if source == "" {
  467. return nil, false, nil
  468. }
  469. if c.txStatus == TxStatusIdle {
  470. c.schema.LockStatement()
  471. defer c.schema.UnlockStatement()
  472. }
  473. tables, err := c.schema.ListTables()
  474. if err != nil {
  475. return nil, true, err
  476. }
  477. rows := make([]map[string]interface{}, 0)
  478. switch source {
  479. case "INFORMATION_SCHEMA.TABLES":
  480. for _, table := range tables {
  481. rows = append(rows, map[string]interface{}{
  482. "table_catalog": c.database, "table_schema": "public", "table_name": table, "table_type": "BASE TABLE",
  483. })
  484. }
  485. case "PG_TABLES":
  486. for _, table := range tables {
  487. indexes, _ := c.schema.ListTableIndexes(table)
  488. rows = append(rows, map[string]interface{}{
  489. "schemaname": "public", "tablename": table, "tableowner": c.params["user"], "tablespace": nil,
  490. "hasindexes": len(indexes) > 0, "hasrules": false, "hastriggers": false, "rowsecurity": false,
  491. })
  492. }
  493. case "INFORMATION_SCHEMA.COLUMNS":
  494. for _, table := range tables {
  495. schema, schemaErr := c.schema.GetSchema(table)
  496. if schemaErr != nil {
  497. continue
  498. }
  499. for i, column := range schema.Columns {
  500. rows = append(rows, map[string]interface{}{
  501. "table_catalog": c.database, "table_schema": "public", "table_name": table,
  502. "column_name": column.Name, "ordinal_position": int64(i + 1), "column_default": column.Default,
  503. "is_nullable": yesNo(column.Nullable), "data_type": strings.ToLower(column.Type),
  504. })
  505. }
  506. }
  507. case "INFORMATION_SCHEMA.TABLE_CONSTRAINTS", "INFORMATION_SCHEMA.KEY_COLUMN_USAGE":
  508. for _, table := range tables {
  509. schema, schemaErr := c.schema.GetSchema(table)
  510. if schemaErr != nil || schema.PrimaryKey == "" || schema.PrimaryKey == "_rowid_" {
  511. continue
  512. }
  513. row := map[string]interface{}{
  514. "constraint_catalog": c.database, "constraint_schema": "public", "constraint_name": table + "_pkey",
  515. "table_catalog": c.database, "table_schema": "public", "table_name": table,
  516. }
  517. if source == "INFORMATION_SCHEMA.TABLE_CONSTRAINTS" {
  518. row["constraint_type"] = "PRIMARY KEY"
  519. row["is_deferrable"] = "NO"
  520. row["initially_deferred"] = "NO"
  521. } else {
  522. row["column_name"] = schema.PrimaryKey
  523. row["ordinal_position"] = int64(1)
  524. }
  525. rows = append(rows, row)
  526. }
  527. case "PG_INDEXES":
  528. for _, table := range tables {
  529. indexes, _ := c.schema.ListTableIndexes(table)
  530. for _, index := range indexes {
  531. columns := make([]string, len(index.Columns))
  532. for i, column := range index.Columns {
  533. columns[i] = column.Name
  534. }
  535. unique := ""
  536. if index.Unique {
  537. unique = "UNIQUE "
  538. }
  539. rows = append(rows, map[string]interface{}{
  540. "schemaname": "public", "tablename": table, "indexname": index.Name, "tablespace": nil,
  541. "indexdef": fmt.Sprintf("CREATE %sINDEX %s ON %s (%s)", unique, index.Name, table, strings.Join(columns, ", ")),
  542. })
  543. }
  544. }
  545. }
  546. for _, match := range catalogFilterPattern.FindAllStringSubmatch(sql, -1) {
  547. column, expected := strings.ToLower(match[1]), strings.ReplaceAll(match[2], "''", "'")
  548. filtered := rows[:0]
  549. for _, row := range rows {
  550. if value, ok := row[column]; ok && strings.EqualFold(fmt.Sprintf("%v", value), expected) {
  551. filtered = append(filtered, row)
  552. }
  553. }
  554. rows = filtered
  555. }
  556. columns := catalogProjection(sql, rows)
  557. result := executor.NewResult("SELECT")
  558. if len(columns) == 1 && columns[0] == "count(*)" {
  559. result.AddColumnWithType("count", "INTEGER")
  560. result.AddRow(int64(len(rows)))
  561. return result, true, nil
  562. }
  563. for _, column := range columns {
  564. columnType := "TEXT"
  565. if column == "ordinal_position" {
  566. columnType = "INTEGER"
  567. } else if strings.HasPrefix(column, "has") || column == "rowsecurity" {
  568. columnType = "BOOLEAN"
  569. }
  570. result.AddColumnWithType(column, columnType)
  571. }
  572. for _, row := range rows {
  573. values := make([]interface{}, len(columns))
  574. for i, column := range columns {
  575. values[i] = row[column]
  576. }
  577. result.AddRow(values...)
  578. }
  579. return result, true, nil
  580. }
  581. func isSingleStatementSQL(sql string) bool {
  582. trimmed := strings.TrimSpace(sql)
  583. if strings.HasSuffix(trimmed, ";") {
  584. trimmed = strings.TrimSpace(strings.TrimSuffix(trimmed, ";"))
  585. }
  586. inString, inIdent := false, false
  587. for i := 0; i < len(trimmed); i++ {
  588. switch trimmed[i] {
  589. case '\'':
  590. if !inIdent {
  591. if inString && i+1 < len(trimmed) && trimmed[i+1] == '\'' {
  592. i++
  593. continue
  594. }
  595. inString = !inString
  596. }
  597. case '"':
  598. if !inString {
  599. inIdent = !inIdent
  600. }
  601. case ';':
  602. if !inString && !inIdent {
  603. return false
  604. }
  605. }
  606. }
  607. return true
  608. }
  609. // isRollbackStatement reports whether sql parses as a single ROLLBACK
  610. // statement, including ROLLBACK TO SAVEPOINT.
  611. func isRollbackStatement(sql string) bool {
  612. stmt, err := parser.New(lexer.New(sql)).Parse()
  613. if err != nil {
  614. return false
  615. }
  616. _, ok := stmt.(*parser.RollbackStmt)
  617. return ok
  618. }
  619. func catalogProjection(sql string, rows []map[string]interface{}) []string {
  620. upper := strings.ToUpper(sql)
  621. selectPos, fromPos := strings.Index(upper, "SELECT"), strings.Index(upper, " FROM ")
  622. if selectPos < 0 || fromPos < 0 || fromPos <= selectPos+6 {
  623. return nil
  624. }
  625. projection := strings.TrimSpace(sql[selectPos+6 : fromPos])
  626. projection = strings.TrimSpace(strings.TrimPrefix(strings.ToUpper(projection), "DISTINCT "))
  627. if projection == "*" && len(rows) > 0 {
  628. columns := make([]string, 0, len(rows[0]))
  629. for column := range rows[0] {
  630. columns = append(columns, column)
  631. }
  632. sort.Strings(columns)
  633. return columns
  634. }
  635. parts := strings.Split(projection, ",")
  636. columns := make([]string, 0, len(parts))
  637. for _, part := range parts {
  638. column := strings.TrimSpace(part)
  639. if index := strings.Index(strings.ToUpper(column), " AS "); index >= 0 {
  640. column = strings.TrimSpace(column[:index])
  641. }
  642. if index := strings.LastIndex(column, "."); index >= 0 {
  643. column = column[index+1:]
  644. }
  645. columns = append(columns, strings.ToLower(strings.Trim(column, `"`)))
  646. }
  647. return columns
  648. }
  649. func yesNo(value bool) string {
  650. if value {
  651. return "YES"
  652. }
  653. return "NO"
  654. }
  655. func (c *Connection) sendTabularResult(result *executor.Result, tag string) error {
  656. if err := c.sendRowDescription(result.Columns, result.ColumnTypes); err != nil {
  657. return err
  658. }
  659. for _, row := range result.Rows {
  660. if err := c.sendDataRow(row, result.Columns); err != nil {
  661. return err
  662. }
  663. }
  664. return c.sendCommandComplete(fmt.Sprintf("%s %d", tag, len(result.Rows)))
  665. }
  666. func (c *Connection) handleClose(msg *Message) error {
  667. if c.extendedFailed {
  668. return nil
  669. }
  670. if len(msg.Data) < 2 {
  671. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close message"))
  672. }
  673. name, pos, err := readCString(msg.Data, 1)
  674. if err != nil || pos != len(msg.Data) {
  675. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close name"))
  676. }
  677. if msg.Data[0] == 'S' {
  678. delete(c.statements, name)
  679. } else if msg.Data[0] == 'P' {
  680. delete(c.portals, name)
  681. } else {
  682. return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Close target"))
  683. }
  684. return c.writeMessage(MsgCloseComplete, nil)
  685. }
  686. func (c *Connection) failExtended(code string, err error) error {
  687. c.extendedFailed = true
  688. return c.sendError("ERROR", code, err.Error())
  689. }
  690. type boundParameter struct {
  691. value []byte
  692. oid int32
  693. format int16
  694. null bool
  695. }
  696. func readCString(data []byte, pos int) (string, int, error) {
  697. if pos < 0 || pos >= len(data) {
  698. return "", pos, fmt.Errorf("missing null-terminated string")
  699. }
  700. end := bytes.IndexByte(data[pos:], 0)
  701. if end < 0 {
  702. return "", pos, fmt.Errorf("unterminated string")
  703. }
  704. return string(data[pos : pos+end]), pos + end + 1, nil
  705. }
  706. func readInt16List(data []byte, pos int) ([]int16, int, error) {
  707. if pos+2 > len(data) {
  708. return nil, pos, fmt.Errorf("missing list length")
  709. }
  710. count := int(binary.BigEndian.Uint16(data[pos : pos+2]))
  711. pos += 2
  712. if pos+count*2 > len(data) {
  713. return nil, pos, fmt.Errorf("truncated list")
  714. }
  715. result := make([]int16, count)
  716. for i := range result {
  717. result[i] = int16(binary.BigEndian.Uint16(data[pos : pos+2]))
  718. pos += 2
  719. }
  720. return result, pos, nil
  721. }
  722. func parameterOID(oids []int32, index int) int32 {
  723. if index < len(oids) {
  724. return oids[index]
  725. }
  726. return 0
  727. }
  728. func parameterFormat(formats []int16, index int) int16 {
  729. if len(formats) == 1 {
  730. return formats[0]
  731. }
  732. if index < len(formats) {
  733. return formats[index]
  734. }
  735. return 0
  736. }
  737. func bindQuery(query string, params []boundParameter) (string, error) {
  738. var result strings.Builder
  739. inString, inIdent := false, false
  740. for i := 0; i < len(query); {
  741. ch := query[i]
  742. if ch == '\'' && !inIdent {
  743. result.WriteByte(ch)
  744. if inString && i+1 < len(query) && query[i+1] == '\'' {
  745. result.WriteByte(query[i+1])
  746. i += 2
  747. continue
  748. }
  749. inString = !inString
  750. i++
  751. continue
  752. }
  753. if ch == '"' && !inString {
  754. inIdent = !inIdent
  755. result.WriteByte(ch)
  756. i++
  757. continue
  758. }
  759. if ch == '$' && !inString && !inIdent && i+1 < len(query) && query[i+1] >= '0' && query[i+1] <= '9' {
  760. end := i + 1
  761. for end < len(query) && query[end] >= '0' && query[end] <= '9' {
  762. end++
  763. }
  764. n, _ := strconv.Atoi(query[i+1 : end])
  765. if n < 1 || n > len(params) {
  766. return "", fmt.Errorf("parameter $%d was not provided", n)
  767. }
  768. literal, err := parameterLiteral(params[n-1])
  769. if err != nil {
  770. return "", fmt.Errorf("parameter $%d: %w", n, err)
  771. }
  772. result.WriteString(literal)
  773. i = end
  774. continue
  775. }
  776. result.WriteByte(ch)
  777. i++
  778. }
  779. return result.String(), nil
  780. }
  781. func parameterLiteral(param boundParameter) (string, error) {
  782. if param.null {
  783. return "NULL", nil
  784. }
  785. if param.format == 1 {
  786. switch param.oid {
  787. case 16:
  788. if len(param.value) != 1 {
  789. return "", fmt.Errorf("invalid binary boolean")
  790. }
  791. if param.value[0] == 0 {
  792. return "FALSE", nil
  793. }
  794. return "TRUE", nil
  795. case 21:
  796. if len(param.value) != 2 {
  797. return "", fmt.Errorf("invalid binary int2")
  798. }
  799. return strconv.FormatInt(int64(int16(binary.BigEndian.Uint16(param.value))), 10), nil
  800. case 23:
  801. if len(param.value) != 4 {
  802. return "", fmt.Errorf("invalid binary int4")
  803. }
  804. return strconv.FormatInt(int64(int32(binary.BigEndian.Uint32(param.value))), 10), nil
  805. case 20:
  806. if len(param.value) != 8 {
  807. return "", fmt.Errorf("invalid binary int8")
  808. }
  809. return strconv.FormatInt(int64(binary.BigEndian.Uint64(param.value)), 10), nil
  810. default:
  811. return "", fmt.Errorf("binary format is unsupported for OID %d", param.oid)
  812. }
  813. }
  814. value := string(param.value)
  815. switch param.oid {
  816. case 0:
  817. if strings.EqualFold(value, "true") || strings.EqualFold(value, "false") {
  818. return strings.ToUpper(value), nil
  819. }
  820. if _, err := strconv.ParseFloat(value, 64); err == nil && value != "" {
  821. return value, nil
  822. }
  823. return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil
  824. case 16:
  825. if strings.EqualFold(value, "true") || value == "1" || value == "t" {
  826. return "TRUE", nil
  827. }
  828. return "FALSE", nil
  829. case 20, 21, 23, 26, 700, 701, 1700:
  830. if _, err := strconv.ParseFloat(value, 64); err != nil {
  831. return "", fmt.Errorf("invalid numeric value")
  832. }
  833. return value, nil
  834. default:
  835. return "'" + strings.ReplaceAll(value, "'", "''") + "'", nil
  836. }
  837. }
  838. // sendResult sends query results
  839. func (c *Connection) sendResult(result *executor.Result, stmt parser.Statement) error {
  840. // For SELECT statements, send row description and data rows
  841. if _, isSelect := stmt.(*parser.SelectStmt); isSelect && len(result.Columns) > 0 {
  842. // Send row description
  843. if err := c.sendRowDescription(result.Columns, result.ColumnTypes); err != nil {
  844. return err
  845. }
  846. // Send data rows
  847. for _, row := range result.Rows {
  848. if err := c.sendDataRow(row, result.Columns); err != nil {
  849. return err
  850. }
  851. }
  852. // Send command complete
  853. tag := fmt.Sprintf("SELECT %d", len(result.Rows))
  854. return c.sendCommandComplete(tag)
  855. }
  856. // For other statements, just send command complete
  857. tag := c.getCommandTag(stmt, result)
  858. return c.sendCommandComplete(tag)
  859. }
  860. // getCommandTag returns the command completion tag
  861. func (c *Connection) getCommandTag(stmt parser.Statement, result *executor.Result) string {
  862. switch s := stmt.(type) {
  863. case *parser.CreateTableStmt:
  864. return "CREATE TABLE"
  865. case *parser.DropTableStmt:
  866. return "DROP TABLE"
  867. case *parser.CreateIndexStmt:
  868. return "CREATE INDEX"
  869. case *parser.DropIndexStmt:
  870. return "DROP INDEX"
  871. case *parser.AlterTableStmt:
  872. return "ALTER TABLE"
  873. case *parser.InsertStmt:
  874. return fmt.Sprintf("INSERT 0 %d", result.RowsAffected)
  875. case *parser.UpdateStmt:
  876. return fmt.Sprintf("UPDATE %d", result.RowsAffected)
  877. case *parser.DeleteStmt:
  878. return fmt.Sprintf("DELETE %d", result.RowsAffected)
  879. case *parser.BeginStmt:
  880. c.txStatus = TxStatusInBlock
  881. return "BEGIN"
  882. case *parser.CommitStmt:
  883. c.txStatus = TxStatusIdle
  884. return "COMMIT"
  885. case *parser.RollbackStmt:
  886. if s.Savepoint != "" {
  887. c.txStatus = TxStatusInBlock
  888. return "ROLLBACK"
  889. }
  890. c.txStatus = TxStatusIdle
  891. return "ROLLBACK"
  892. case *parser.SavepointStmt:
  893. c.txStatus = TxStatusInBlock
  894. return "SAVEPOINT"
  895. case *parser.ReleaseStmt:
  896. c.txStatus = TxStatusInBlock
  897. return "RELEASE"
  898. default:
  899. return "OK"
  900. }
  901. }
  902. // sendAuthenticationOk sends authentication OK message
  903. func (c *Connection) sendAuthenticationOk() error {
  904. mb := NewMessageBuilder()
  905. mb.WriteInt32(0) // Auth OK
  906. return c.writeMessage(MsgAuthenticationOk, mb.Bytes())
  907. }
  908. // sendParameterStatus sends a parameter status message
  909. func (c *Connection) sendParameterStatus(name, value string) error {
  910. mb := NewMessageBuilder()
  911. mb.WriteString(name)
  912. mb.WriteString(value)
  913. return c.writeMessage(MsgParameterStatus, mb.Bytes())
  914. }
  915. // sendBackendKeyData sends backend key data
  916. func (c *Connection) sendBackendKeyData(processID, secretKey int32) error {
  917. mb := NewMessageBuilder()
  918. mb.WriteInt32(processID)
  919. mb.WriteInt32(secretKey)
  920. return c.writeMessage(MsgBackendKeyData, mb.Bytes())
  921. }
  922. // sendReadyForQuery sends ready for query message
  923. func (c *Connection) sendReadyForQuery() error {
  924. mb := NewMessageBuilder()
  925. mb.AppendByte(c.txStatus)
  926. if err := c.writeMessage(MsgReadyForQuery, mb.Bytes()); err != nil {
  927. return err
  928. }
  929. // ReadyForQuery closes out a response cycle, so flush everything buffered
  930. // so far. This is what makes simple-query results and Sync responses
  931. // visible to the client.
  932. return c.writer.Flush()
  933. }
  934. // sendEmptyQueryResponse sends empty query response
  935. func (c *Connection) sendEmptyQueryResponse() error {
  936. return c.writeMessage(MsgEmptyQueryResponse, []byte{})
  937. }
  938. // sendRowDescription sends row description (column metadata)
  939. func (c *Connection) sendRowDescription(columns []string, columnTypes []string) error {
  940. mb := NewMessageBuilder()
  941. mb.WriteInt16(int16(len(columns)))
  942. for i, col := range columns {
  943. colType := ""
  944. if i < len(columnTypes) {
  945. colType = columnTypes[i]
  946. }
  947. mb.WriteString(col)
  948. mb.WriteInt32(0) // table OID
  949. mb.WriteInt16(0) // column attribute number
  950. mb.WriteInt32(c.getOIDForType(colType)) // type OID
  951. mb.WriteInt16(c.getTypeSizeForType(colType)) // type size
  952. mb.WriteInt32(-1) // type modifier
  953. mb.WriteInt16(0) // format code (text)
  954. }
  955. return c.writeMessage(MsgRowDescription, mb.Bytes())
  956. }
  957. // sendDataRow sends a data row
  958. func (c *Connection) sendDataRow(row []interface{}, columns []string) error {
  959. mb := NewMessageBuilder()
  960. mb.WriteInt16(int16(len(row)))
  961. for _, value := range row {
  962. if value == nil {
  963. mb.WriteInt32(-1) // NULL indicator
  964. continue
  965. }
  966. // Convert value to string
  967. strValue := c.valueToString(value)
  968. mb.WriteInt32(int32(len(strValue)))
  969. mb.WriteBytes([]byte(strValue))
  970. }
  971. return c.writeMessage(MsgDataRow, mb.Bytes())
  972. }
  973. // sendCommandComplete sends command complete message
  974. func (c *Connection) sendCommandComplete(tag string) error {
  975. mb := NewMessageBuilder()
  976. mb.WriteString(tag)
  977. return c.writeMessage(MsgCommandComplete, mb.Bytes())
  978. }
  979. // sendError sends an error response
  980. func (c *Connection) sendError(severity, code, message string) error {
  981. mb := NewMessageBuilder()
  982. mb.AppendByte(ErrorFieldSeverity)
  983. mb.WriteString(severity)
  984. mb.AppendByte(ErrorFieldCode)
  985. mb.WriteString(code)
  986. mb.AppendByte(ErrorFieldMessage)
  987. mb.WriteString(message)
  988. mb.AppendByte(0) // Terminator
  989. if err := c.writeMessage(MsgErrorResponse, mb.Bytes()); err != nil {
  990. return err
  991. }
  992. // Errors must be visible promptly, including FATAL startup failures that
  993. // are not followed by a ReadyForQuery before the connection closes.
  994. return c.writer.Flush()
  995. }
  996. // writeMessage buffers a message for the connection. It does not flush, so
  997. // callers that need to make a response visible to the client must flush at the
  998. // appropriate protocol boundary (ReadyForQuery, an explicit Flush message, or
  999. // an error response). Buffering amortizes the per-message syscalls that a
  1000. // result set would otherwise incur; the underlying bufio.Writer bounds memory
  1001. // use so large result sets cannot grow the buffer without limit.
  1002. func (c *Connection) writeMessage(msgType byte, data []byte) error {
  1003. if !c.quiet {
  1004. log.Printf("Sending message type=%c length=%d", msgType, len(data)+4)
  1005. }
  1006. return WriteMessage(c.writer, msgType, data)
  1007. }
  1008. // getOIDForType returns PostgreSQL OID for type
  1009. func (c *Connection) getOIDForType(typeName string) int32 {
  1010. switch strings.ToUpper(typeName) {
  1011. case "INTEGER", "INT":
  1012. return 23 // INT4OID
  1013. case "TEXT", "VARCHAR", "CHAR":
  1014. return 25 // TEXTOID
  1015. case "REAL", "FLOAT":
  1016. return 700 // FLOAT4OID
  1017. case "DOUBLE":
  1018. return 701 // FLOAT8OID
  1019. case "BOOLEAN", "BOOL":
  1020. return 16 // BOOLOID
  1021. case "BLOB":
  1022. return 17 // BYTEAOID
  1023. default:
  1024. return 25 // Default to TEXT
  1025. }
  1026. }
  1027. // getTypeSizeForType returns type size
  1028. func (c *Connection) getTypeSizeForType(typeName string) int16 {
  1029. switch strings.ToUpper(typeName) {
  1030. case "INTEGER", "INT":
  1031. return 4
  1032. case "REAL", "FLOAT":
  1033. return 4
  1034. case "DOUBLE":
  1035. return 8
  1036. case "BOOLEAN", "BOOL":
  1037. return 1
  1038. default:
  1039. return -1 // Variable length
  1040. }
  1041. }
  1042. // valueToString converts a value to string
  1043. func (c *Connection) valueToString(value interface{}) string {
  1044. if value == nil {
  1045. return ""
  1046. }
  1047. return fmt.Sprintf("%v", value)
  1048. }
  1049. // handleVersionQuery handles SELECT version()
  1050. func (c *Connection) handleVersionQuery() error {
  1051. columns := []string{"version"}
  1052. columnTypes := []string{"TEXT"}
  1053. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1054. return err
  1055. }
  1056. row := []interface{}{"PostgreSQL 14.0 (PizzaSQL)"}
  1057. if err := c.sendDataRow(row, columns); err != nil {
  1058. return err
  1059. }
  1060. if err := c.sendCommandComplete("SELECT 1"); err != nil {
  1061. return err
  1062. }
  1063. return c.sendReadyForQuery()
  1064. }
  1065. // handleCurrentUserQuery handles SELECT current_user
  1066. func (c *Connection) handleCurrentUserQuery() error {
  1067. columns := []string{"current_user"}
  1068. columnTypes := []string{"TEXT"}
  1069. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1070. return err
  1071. }
  1072. user := c.params["user"]
  1073. if user == "" {
  1074. user = "pizzasql"
  1075. }
  1076. row := []interface{}{user}
  1077. if err := c.sendDataRow(row, columns); err != nil {
  1078. return err
  1079. }
  1080. if err := c.sendCommandComplete("SELECT 1"); err != nil {
  1081. return err
  1082. }
  1083. return c.sendReadyForQuery()
  1084. }
  1085. // handleShowCommand handles SHOW commands
  1086. func (c *Connection) handleShowCommand(sqlUpper string) error {
  1087. var value string
  1088. var name string
  1089. if strings.Contains(sqlUpper, "SERVER_VERSION") {
  1090. name = "server_version"
  1091. value = "14.0"
  1092. } else if strings.Contains(sqlUpper, "SERVER_ENCODING") {
  1093. name = "server_encoding"
  1094. value = "UTF8"
  1095. } else if strings.Contains(sqlUpper, "CLIENT_ENCODING") {
  1096. name = "client_encoding"
  1097. value = "UTF8"
  1098. } else {
  1099. // Unknown SHOW command
  1100. c.sendError("ERROR", ErrCodeFeatureNotSupported, "SHOW command not supported")
  1101. return c.sendReadyForQuery()
  1102. }
  1103. columns := []string{name}
  1104. columnTypes := []string{"TEXT"}
  1105. if err := c.sendRowDescription(columns, columnTypes); err != nil {
  1106. return err
  1107. }
  1108. row := []interface{}{value}
  1109. if err := c.sendDataRow(row, columns); err != nil {
  1110. return err
  1111. }
  1112. if err := c.sendCommandComplete("SHOW"); err != nil {
  1113. return err
  1114. }
  1115. return c.sendReadyForQuery()
  1116. }