conn.go 42 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823
  1. package pq
  2. import (
  3. "bufio"
  4. "crypto/md5"
  5. "database/sql"
  6. "database/sql/driver"
  7. "encoding/binary"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "net"
  12. "os"
  13. "os/user"
  14. "path"
  15. "path/filepath"
  16. "strconv"
  17. "strings"
  18. "time"
  19. "unicode"
  20. "github.com/lib/pq/oid"
  21. )
  22. // Common error types
  23. var (
  24. ErrNotSupported = errors.New("pq: Unsupported command")
  25. ErrInFailedTransaction = errors.New("pq: Could not complete operation in a failed transaction")
  26. ErrSSLNotSupported = errors.New("pq: SSL is not enabled on the server")
  27. ErrSSLKeyHasWorldPermissions = errors.New("pq: Private key file has group or world access. Permissions should be u=rw (0600) or less")
  28. ErrCouldNotDetectUsername = errors.New("pq: Could not detect default username. Please provide one explicitly")
  29. errUnexpectedReady = errors.New("unexpected ReadyForQuery")
  30. errNoRowsAffected = errors.New("no RowsAffected available after the empty statement")
  31. errNoLastInsertID = errors.New("no LastInsertId available after the empty statement")
  32. )
  33. type Driver struct{}
  34. func (d *Driver) Open(name string) (driver.Conn, error) {
  35. return Open(name)
  36. }
  37. func init() {
  38. sql.Register("postgres", &Driver{})
  39. }
  40. type parameterStatus struct {
  41. // server version in the same format as server_version_num, or 0 if
  42. // unavailable
  43. serverVersion int
  44. // the current location based on the TimeZone value of the session, if
  45. // available
  46. currentLocation *time.Location
  47. }
  48. type transactionStatus byte
  49. const (
  50. txnStatusIdle transactionStatus = 'I'
  51. txnStatusIdleInTransaction transactionStatus = 'T'
  52. txnStatusInFailedTransaction transactionStatus = 'E'
  53. )
  54. func (s transactionStatus) String() string {
  55. switch s {
  56. case txnStatusIdle:
  57. return "idle"
  58. case txnStatusIdleInTransaction:
  59. return "idle in transaction"
  60. case txnStatusInFailedTransaction:
  61. return "in a failed transaction"
  62. default:
  63. errorf("unknown transactionStatus %d", s)
  64. }
  65. panic("not reached")
  66. }
  67. type Dialer interface {
  68. Dial(network, address string) (net.Conn, error)
  69. DialTimeout(network, address string, timeout time.Duration) (net.Conn, error)
  70. }
  71. type defaultDialer struct{}
  72. func (d defaultDialer) Dial(ntw, addr string) (net.Conn, error) {
  73. return net.Dial(ntw, addr)
  74. }
  75. func (d defaultDialer) DialTimeout(ntw, addr string, timeout time.Duration) (net.Conn, error) {
  76. return net.DialTimeout(ntw, addr, timeout)
  77. }
  78. type conn struct {
  79. c net.Conn
  80. buf *bufio.Reader
  81. namei int
  82. scratch [512]byte
  83. txnStatus transactionStatus
  84. txnFinish func()
  85. // Save connection arguments to use during CancelRequest.
  86. dialer Dialer
  87. opts values
  88. // Cancellation key data for use with CancelRequest messages.
  89. processID int
  90. secretKey int
  91. parameterStatus parameterStatus
  92. saveMessageType byte
  93. saveMessageBuffer []byte
  94. // If true, this connection is bad and all public-facing functions should
  95. // return ErrBadConn.
  96. bad bool
  97. // If set, this connection should never use the binary format when
  98. // receiving query results from prepared statements. Only provided for
  99. // debugging.
  100. disablePreparedBinaryResult bool
  101. // Whether to always send []byte parameters over as binary. Enables single
  102. // round-trip mode for non-prepared Query calls.
  103. binaryParameters bool
  104. // If true this connection is in the middle of a COPY
  105. inCopy bool
  106. }
  107. // Handle driver-side settings in parsed connection string.
  108. func (cn *conn) handleDriverSettings(o values) (err error) {
  109. boolSetting := func(key string, val *bool) error {
  110. if value, ok := o[key]; ok {
  111. if value == "yes" {
  112. *val = true
  113. } else if value == "no" {
  114. *val = false
  115. } else {
  116. return fmt.Errorf("unrecognized value %q for %s", value, key)
  117. }
  118. }
  119. return nil
  120. }
  121. err = boolSetting("disable_prepared_binary_result", &cn.disablePreparedBinaryResult)
  122. if err != nil {
  123. return err
  124. }
  125. err = boolSetting("binary_parameters", &cn.binaryParameters)
  126. if err != nil {
  127. return err
  128. }
  129. return nil
  130. }
  131. func (cn *conn) handlePgpass(o values) {
  132. // if a password was supplied, do not process .pgpass
  133. if _, ok := o["password"]; ok {
  134. return
  135. }
  136. filename := os.Getenv("PGPASSFILE")
  137. if filename == "" {
  138. // XXX this code doesn't work on Windows where the default filename is
  139. // XXX %APPDATA%\postgresql\pgpass.conf
  140. user, err := user.Current()
  141. if err != nil {
  142. return
  143. }
  144. filename = filepath.Join(user.HomeDir, ".pgpass")
  145. }
  146. fileinfo, err := os.Stat(filename)
  147. if err != nil {
  148. return
  149. }
  150. mode := fileinfo.Mode()
  151. if mode&(0x77) != 0 {
  152. // XXX should warn about incorrect .pgpass permissions as psql does
  153. return
  154. }
  155. file, err := os.Open(filename)
  156. if err != nil {
  157. return
  158. }
  159. defer file.Close()
  160. scanner := bufio.NewScanner(io.Reader(file))
  161. hostname := o["host"]
  162. ntw, _ := network(o)
  163. port := o["port"]
  164. db := o["dbname"]
  165. username := o["user"]
  166. // From: https://github.com/tg/pgpass/blob/master/reader.go
  167. getFields := func(s string) []string {
  168. fs := make([]string, 0, 5)
  169. f := make([]rune, 0, len(s))
  170. var esc bool
  171. for _, c := range s {
  172. switch {
  173. case esc:
  174. f = append(f, c)
  175. esc = false
  176. case c == '\\':
  177. esc = true
  178. case c == ':':
  179. fs = append(fs, string(f))
  180. f = f[:0]
  181. default:
  182. f = append(f, c)
  183. }
  184. }
  185. return append(fs, string(f))
  186. }
  187. for scanner.Scan() {
  188. line := scanner.Text()
  189. if len(line) == 0 || line[0] == '#' {
  190. continue
  191. }
  192. split := getFields(line)
  193. if len(split) != 5 {
  194. continue
  195. }
  196. if (split[0] == "*" || split[0] == hostname || (split[0] == "localhost" && (hostname == "" || ntw == "unix"))) && (split[1] == "*" || split[1] == port) && (split[2] == "*" || split[2] == db) && (split[3] == "*" || split[3] == username) {
  197. o["password"] = split[4]
  198. return
  199. }
  200. }
  201. }
  202. func (cn *conn) writeBuf(b byte) *writeBuf {
  203. cn.scratch[0] = b
  204. return &writeBuf{
  205. buf: cn.scratch[:5],
  206. pos: 1,
  207. }
  208. }
  209. func Open(name string) (_ driver.Conn, err error) {
  210. return DialOpen(defaultDialer{}, name)
  211. }
  212. func DialOpen(d Dialer, name string) (_ driver.Conn, err error) {
  213. // Handle any panics during connection initialization. Note that we
  214. // specifically do *not* want to use errRecover(), as that would turn any
  215. // connection errors into ErrBadConns, hiding the real error message from
  216. // the user.
  217. defer errRecoverNoErrBadConn(&err)
  218. o := make(values)
  219. // A number of defaults are applied here, in this order:
  220. //
  221. // * Very low precedence defaults applied in every situation
  222. // * Environment variables
  223. // * Explicitly passed connection information
  224. o["host"] = "localhost"
  225. o["port"] = "5432"
  226. // N.B.: Extra float digits should be set to 3, but that breaks
  227. // Postgres 8.4 and older, where the max is 2.
  228. o["extra_float_digits"] = "2"
  229. for k, v := range parseEnviron(os.Environ()) {
  230. o[k] = v
  231. }
  232. if strings.HasPrefix(name, "postgres://") || strings.HasPrefix(name, "postgresql://") {
  233. name, err = ParseURL(name)
  234. if err != nil {
  235. return nil, err
  236. }
  237. }
  238. if err := parseOpts(name, o); err != nil {
  239. return nil, err
  240. }
  241. // Use the "fallback" application name if necessary
  242. if fallback, ok := o["fallback_application_name"]; ok {
  243. if _, ok := o["application_name"]; !ok {
  244. o["application_name"] = fallback
  245. }
  246. }
  247. // We can't work with any client_encoding other than UTF-8 currently.
  248. // However, we have historically allowed the user to set it to UTF-8
  249. // explicitly, and there's no reason to break such programs, so allow that.
  250. // Note that the "options" setting could also set client_encoding, but
  251. // parsing its value is not worth it. Instead, we always explicitly send
  252. // client_encoding as a separate run-time parameter, which should override
  253. // anything set in options.
  254. if enc, ok := o["client_encoding"]; ok && !isUTF8(enc) {
  255. return nil, errors.New("client_encoding must be absent or 'UTF8'")
  256. }
  257. o["client_encoding"] = "UTF8"
  258. // DateStyle needs a similar treatment.
  259. if datestyle, ok := o["datestyle"]; ok {
  260. if datestyle != "ISO, MDY" {
  261. panic(fmt.Sprintf("setting datestyle must be absent or %v; got %v",
  262. "ISO, MDY", datestyle))
  263. }
  264. } else {
  265. o["datestyle"] = "ISO, MDY"
  266. }
  267. // If a user is not provided by any other means, the last
  268. // resort is to use the current operating system provided user
  269. // name.
  270. if _, ok := o["user"]; !ok {
  271. u, err := userCurrent()
  272. if err != nil {
  273. return nil, err
  274. }
  275. o["user"] = u
  276. }
  277. cn := &conn{
  278. opts: o,
  279. dialer: d,
  280. }
  281. err = cn.handleDriverSettings(o)
  282. if err != nil {
  283. return nil, err
  284. }
  285. cn.handlePgpass(o)
  286. cn.c, err = dial(d, o)
  287. if err != nil {
  288. return nil, err
  289. }
  290. cn.ssl(o)
  291. cn.buf = bufio.NewReader(cn.c)
  292. cn.startup(o)
  293. // reset the deadline, in case one was set (see dial)
  294. if timeout, ok := o["connect_timeout"]; ok && timeout != "0" {
  295. err = cn.c.SetDeadline(time.Time{})
  296. }
  297. return cn, err
  298. }
  299. func dial(d Dialer, o values) (net.Conn, error) {
  300. ntw, addr := network(o)
  301. // SSL is not necessary or supported over UNIX domain sockets
  302. if ntw == "unix" {
  303. o["sslmode"] = "disable"
  304. }
  305. // Zero or not specified means wait indefinitely.
  306. if timeout, ok := o["connect_timeout"]; ok && timeout != "0" {
  307. seconds, err := strconv.ParseInt(timeout, 10, 0)
  308. if err != nil {
  309. return nil, fmt.Errorf("invalid value for parameter connect_timeout: %s", err)
  310. }
  311. duration := time.Duration(seconds) * time.Second
  312. // connect_timeout should apply to the entire connection establishment
  313. // procedure, so we both use a timeout for the TCP connection
  314. // establishment and set a deadline for doing the initial handshake.
  315. // The deadline is then reset after startup() is done.
  316. deadline := time.Now().Add(duration)
  317. conn, err := d.DialTimeout(ntw, addr, duration)
  318. if err != nil {
  319. return nil, err
  320. }
  321. err = conn.SetDeadline(deadline)
  322. return conn, err
  323. }
  324. return d.Dial(ntw, addr)
  325. }
  326. func network(o values) (string, string) {
  327. host := o["host"]
  328. if strings.HasPrefix(host, "/") {
  329. sockPath := path.Join(host, ".s.PGSQL."+o["port"])
  330. return "unix", sockPath
  331. }
  332. return "tcp", net.JoinHostPort(host, o["port"])
  333. }
  334. type values map[string]string
  335. // scanner implements a tokenizer for libpq-style option strings.
  336. type scanner struct {
  337. s []rune
  338. i int
  339. }
  340. // newScanner returns a new scanner initialized with the option string s.
  341. func newScanner(s string) *scanner {
  342. return &scanner{[]rune(s), 0}
  343. }
  344. // Next returns the next rune.
  345. // It returns 0, false if the end of the text has been reached.
  346. func (s *scanner) Next() (rune, bool) {
  347. if s.i >= len(s.s) {
  348. return 0, false
  349. }
  350. r := s.s[s.i]
  351. s.i++
  352. return r, true
  353. }
  354. // SkipSpaces returns the next non-whitespace rune.
  355. // It returns 0, false if the end of the text has been reached.
  356. func (s *scanner) SkipSpaces() (rune, bool) {
  357. r, ok := s.Next()
  358. for unicode.IsSpace(r) && ok {
  359. r, ok = s.Next()
  360. }
  361. return r, ok
  362. }
  363. // parseOpts parses the options from name and adds them to the values.
  364. //
  365. // The parsing code is based on conninfo_parse from libpq's fe-connect.c
  366. func parseOpts(name string, o values) error {
  367. s := newScanner(name)
  368. for {
  369. var (
  370. keyRunes, valRunes []rune
  371. r rune
  372. ok bool
  373. )
  374. if r, ok = s.SkipSpaces(); !ok {
  375. break
  376. }
  377. // Scan the key
  378. for !unicode.IsSpace(r) && r != '=' {
  379. keyRunes = append(keyRunes, r)
  380. if r, ok = s.Next(); !ok {
  381. break
  382. }
  383. }
  384. // Skip any whitespace if we're not at the = yet
  385. if r != '=' {
  386. r, ok = s.SkipSpaces()
  387. }
  388. // The current character should be =
  389. if r != '=' || !ok {
  390. return fmt.Errorf(`missing "=" after %q in connection info string"`, string(keyRunes))
  391. }
  392. // Skip any whitespace after the =
  393. if r, ok = s.SkipSpaces(); !ok {
  394. // If we reach the end here, the last value is just an empty string as per libpq.
  395. o[string(keyRunes)] = ""
  396. break
  397. }
  398. if r != '\'' {
  399. for !unicode.IsSpace(r) {
  400. if r == '\\' {
  401. if r, ok = s.Next(); !ok {
  402. return fmt.Errorf(`missing character after backslash`)
  403. }
  404. }
  405. valRunes = append(valRunes, r)
  406. if r, ok = s.Next(); !ok {
  407. break
  408. }
  409. }
  410. } else {
  411. quote:
  412. for {
  413. if r, ok = s.Next(); !ok {
  414. return fmt.Errorf(`unterminated quoted string literal in connection string`)
  415. }
  416. switch r {
  417. case '\'':
  418. break quote
  419. case '\\':
  420. r, _ = s.Next()
  421. fallthrough
  422. default:
  423. valRunes = append(valRunes, r)
  424. }
  425. }
  426. }
  427. o[string(keyRunes)] = string(valRunes)
  428. }
  429. return nil
  430. }
  431. func (cn *conn) isInTransaction() bool {
  432. return cn.txnStatus == txnStatusIdleInTransaction ||
  433. cn.txnStatus == txnStatusInFailedTransaction
  434. }
  435. func (cn *conn) checkIsInTransaction(intxn bool) {
  436. if cn.isInTransaction() != intxn {
  437. cn.bad = true
  438. errorf("unexpected transaction status %v", cn.txnStatus)
  439. }
  440. }
  441. func (cn *conn) Begin() (_ driver.Tx, err error) {
  442. return cn.begin("")
  443. }
  444. func (cn *conn) begin(mode string) (_ driver.Tx, err error) {
  445. if cn.bad {
  446. return nil, driver.ErrBadConn
  447. }
  448. defer cn.errRecover(&err)
  449. cn.checkIsInTransaction(false)
  450. _, commandTag, err := cn.simpleExec("BEGIN" + mode)
  451. if err != nil {
  452. return nil, err
  453. }
  454. if commandTag != "BEGIN" {
  455. cn.bad = true
  456. return nil, fmt.Errorf("unexpected command tag %s", commandTag)
  457. }
  458. if cn.txnStatus != txnStatusIdleInTransaction {
  459. cn.bad = true
  460. return nil, fmt.Errorf("unexpected transaction status %v", cn.txnStatus)
  461. }
  462. return cn, nil
  463. }
  464. func (cn *conn) closeTxn() {
  465. if finish := cn.txnFinish; finish != nil {
  466. finish()
  467. }
  468. }
  469. func (cn *conn) Commit() (err error) {
  470. defer cn.closeTxn()
  471. if cn.bad {
  472. return driver.ErrBadConn
  473. }
  474. defer cn.errRecover(&err)
  475. cn.checkIsInTransaction(true)
  476. // We don't want the client to think that everything is okay if it tries
  477. // to commit a failed transaction. However, no matter what we return,
  478. // database/sql will release this connection back into the free connection
  479. // pool so we have to abort the current transaction here. Note that you
  480. // would get the same behaviour if you issued a COMMIT in a failed
  481. // transaction, so it's also the least surprising thing to do here.
  482. if cn.txnStatus == txnStatusInFailedTransaction {
  483. if err := cn.Rollback(); err != nil {
  484. return err
  485. }
  486. return ErrInFailedTransaction
  487. }
  488. _, commandTag, err := cn.simpleExec("COMMIT")
  489. if err != nil {
  490. if cn.isInTransaction() {
  491. cn.bad = true
  492. }
  493. return err
  494. }
  495. if commandTag != "COMMIT" {
  496. cn.bad = true
  497. return fmt.Errorf("unexpected command tag %s", commandTag)
  498. }
  499. cn.checkIsInTransaction(false)
  500. return nil
  501. }
  502. func (cn *conn) Rollback() (err error) {
  503. defer cn.closeTxn()
  504. if cn.bad {
  505. return driver.ErrBadConn
  506. }
  507. defer cn.errRecover(&err)
  508. cn.checkIsInTransaction(true)
  509. _, commandTag, err := cn.simpleExec("ROLLBACK")
  510. if err != nil {
  511. if cn.isInTransaction() {
  512. cn.bad = true
  513. }
  514. return err
  515. }
  516. if commandTag != "ROLLBACK" {
  517. return fmt.Errorf("unexpected command tag %s", commandTag)
  518. }
  519. cn.checkIsInTransaction(false)
  520. return nil
  521. }
  522. func (cn *conn) gname() string {
  523. cn.namei++
  524. return strconv.FormatInt(int64(cn.namei), 10)
  525. }
  526. func (cn *conn) simpleExec(q string) (res driver.Result, commandTag string, err error) {
  527. b := cn.writeBuf('Q')
  528. b.string(q)
  529. cn.send(b)
  530. for {
  531. t, r := cn.recv1()
  532. switch t {
  533. case 'C':
  534. res, commandTag = cn.parseComplete(r.string())
  535. case 'Z':
  536. cn.processReadyForQuery(r)
  537. if res == nil && err == nil {
  538. err = errUnexpectedReady
  539. }
  540. // done
  541. return
  542. case 'E':
  543. err = parseError(r)
  544. case 'I':
  545. res = emptyRows
  546. case 'T', 'D':
  547. // ignore any results
  548. default:
  549. cn.bad = true
  550. errorf("unknown response for simple query: %q", t)
  551. }
  552. }
  553. }
  554. func (cn *conn) simpleQuery(q string) (res *rows, err error) {
  555. defer cn.errRecover(&err)
  556. b := cn.writeBuf('Q')
  557. b.string(q)
  558. cn.send(b)
  559. for {
  560. t, r := cn.recv1()
  561. switch t {
  562. case 'C', 'I':
  563. // We allow queries which don't return any results through Query as
  564. // well as Exec. We still have to give database/sql a rows object
  565. // the user can close, though, to avoid connections from being
  566. // leaked. A "rows" with done=true works fine for that purpose.
  567. if err != nil {
  568. cn.bad = true
  569. errorf("unexpected message %q in simple query execution", t)
  570. }
  571. if res == nil {
  572. res = &rows{
  573. cn: cn,
  574. }
  575. }
  576. // Set the result and tag to the last command complete if there wasn't a
  577. // query already run. Although queries usually return from here and cede
  578. // control to Next, a query with zero results does not.
  579. if t == 'C' && res.colNames == nil {
  580. res.result, res.tag = cn.parseComplete(r.string())
  581. }
  582. res.done = true
  583. case 'Z':
  584. cn.processReadyForQuery(r)
  585. // done
  586. return
  587. case 'E':
  588. res = nil
  589. err = parseError(r)
  590. case 'D':
  591. if res == nil {
  592. cn.bad = true
  593. errorf("unexpected DataRow in simple query execution")
  594. }
  595. // the query didn't fail; kick off to Next
  596. cn.saveMessage(t, r)
  597. return
  598. case 'T':
  599. // res might be non-nil here if we received a previous
  600. // CommandComplete, but that's fine; just overwrite it
  601. res = &rows{cn: cn}
  602. res.colNames, res.colFmts, res.colTyps = parsePortalRowDescribe(r)
  603. // To work around a bug in QueryRow in Go 1.2 and earlier, wait
  604. // until the first DataRow has been received.
  605. default:
  606. cn.bad = true
  607. errorf("unknown response for simple query: %q", t)
  608. }
  609. }
  610. }
  611. type noRows struct{}
  612. var emptyRows noRows
  613. var _ driver.Result = noRows{}
  614. func (noRows) LastInsertId() (int64, error) {
  615. return 0, errNoLastInsertID
  616. }
  617. func (noRows) RowsAffected() (int64, error) {
  618. return 0, errNoRowsAffected
  619. }
  620. // Decides which column formats to use for a prepared statement. The input is
  621. // an array of type oids, one element per result column.
  622. func decideColumnFormats(colTyps []fieldDesc, forceText bool) (colFmts []format, colFmtData []byte) {
  623. if len(colTyps) == 0 {
  624. return nil, colFmtDataAllText
  625. }
  626. colFmts = make([]format, len(colTyps))
  627. if forceText {
  628. return colFmts, colFmtDataAllText
  629. }
  630. allBinary := true
  631. allText := true
  632. for i, t := range colTyps {
  633. switch t.OID {
  634. // This is the list of types to use binary mode for when receiving them
  635. // through a prepared statement. If a type appears in this list, it
  636. // must also be implemented in binaryDecode in encode.go.
  637. case oid.T_bytea:
  638. fallthrough
  639. case oid.T_int8:
  640. fallthrough
  641. case oid.T_int4:
  642. fallthrough
  643. case oid.T_int2:
  644. fallthrough
  645. case oid.T_uuid:
  646. colFmts[i] = formatBinary
  647. allText = false
  648. default:
  649. allBinary = false
  650. }
  651. }
  652. if allBinary {
  653. return colFmts, colFmtDataAllBinary
  654. } else if allText {
  655. return colFmts, colFmtDataAllText
  656. } else {
  657. colFmtData = make([]byte, 2+len(colFmts)*2)
  658. binary.BigEndian.PutUint16(colFmtData, uint16(len(colFmts)))
  659. for i, v := range colFmts {
  660. binary.BigEndian.PutUint16(colFmtData[2+i*2:], uint16(v))
  661. }
  662. return colFmts, colFmtData
  663. }
  664. }
  665. func (cn *conn) prepareTo(q, stmtName string) *stmt {
  666. st := &stmt{cn: cn, name: stmtName}
  667. b := cn.writeBuf('P')
  668. b.string(st.name)
  669. b.string(q)
  670. b.int16(0)
  671. b.next('D')
  672. b.byte('S')
  673. b.string(st.name)
  674. b.next('S')
  675. cn.send(b)
  676. cn.readParseResponse()
  677. st.paramTyps, st.colNames, st.colTyps = cn.readStatementDescribeResponse()
  678. st.colFmts, st.colFmtData = decideColumnFormats(st.colTyps, cn.disablePreparedBinaryResult)
  679. cn.readReadyForQuery()
  680. return st
  681. }
  682. func (cn *conn) Prepare(q string) (_ driver.Stmt, err error) {
  683. if cn.bad {
  684. return nil, driver.ErrBadConn
  685. }
  686. defer cn.errRecover(&err)
  687. if len(q) >= 4 && strings.EqualFold(q[:4], "COPY") {
  688. s, err := cn.prepareCopyIn(q)
  689. if err == nil {
  690. cn.inCopy = true
  691. }
  692. return s, err
  693. }
  694. return cn.prepareTo(q, cn.gname()), nil
  695. }
  696. func (cn *conn) Close() (err error) {
  697. // Skip cn.bad return here because we always want to close a connection.
  698. defer cn.errRecover(&err)
  699. // Ensure that cn.c.Close is always run. Since error handling is done with
  700. // panics and cn.errRecover, the Close must be in a defer.
  701. defer func() {
  702. cerr := cn.c.Close()
  703. if err == nil {
  704. err = cerr
  705. }
  706. }()
  707. // Don't go through send(); ListenerConn relies on us not scribbling on the
  708. // scratch buffer of this connection.
  709. return cn.sendSimpleMessage('X')
  710. }
  711. // Implement the "Queryer" interface
  712. func (cn *conn) Query(query string, args []driver.Value) (driver.Rows, error) {
  713. return cn.query(query, args)
  714. }
  715. func (cn *conn) query(query string, args []driver.Value) (_ *rows, err error) {
  716. if cn.bad {
  717. return nil, driver.ErrBadConn
  718. }
  719. if cn.inCopy {
  720. return nil, errCopyInProgress
  721. }
  722. defer cn.errRecover(&err)
  723. // Check to see if we can use the "simpleQuery" interface, which is
  724. // *much* faster than going through prepare/exec
  725. if len(args) == 0 {
  726. return cn.simpleQuery(query)
  727. }
  728. if cn.binaryParameters {
  729. cn.sendBinaryModeQuery(query, args)
  730. cn.readParseResponse()
  731. cn.readBindResponse()
  732. rows := &rows{cn: cn}
  733. rows.colNames, rows.colFmts, rows.colTyps = cn.readPortalDescribeResponse()
  734. cn.postExecuteWorkaround()
  735. return rows, nil
  736. }
  737. st := cn.prepareTo(query, "")
  738. st.exec(args)
  739. return &rows{
  740. cn: cn,
  741. colNames: st.colNames,
  742. colTyps: st.colTyps,
  743. colFmts: st.colFmts,
  744. }, nil
  745. }
  746. // Implement the optional "Execer" interface for one-shot queries
  747. func (cn *conn) Exec(query string, args []driver.Value) (res driver.Result, err error) {
  748. if cn.bad {
  749. return nil, driver.ErrBadConn
  750. }
  751. defer cn.errRecover(&err)
  752. // Check to see if we can use the "simpleExec" interface, which is
  753. // *much* faster than going through prepare/exec
  754. if len(args) == 0 {
  755. // ignore commandTag, our caller doesn't care
  756. r, _, err := cn.simpleExec(query)
  757. return r, err
  758. }
  759. if cn.binaryParameters {
  760. cn.sendBinaryModeQuery(query, args)
  761. cn.readParseResponse()
  762. cn.readBindResponse()
  763. cn.readPortalDescribeResponse()
  764. cn.postExecuteWorkaround()
  765. res, _, err = cn.readExecuteResponse("Execute")
  766. return res, err
  767. }
  768. // Use the unnamed statement to defer planning until bind
  769. // time, or else value-based selectivity estimates cannot be
  770. // used.
  771. st := cn.prepareTo(query, "")
  772. r, err := st.Exec(args)
  773. if err != nil {
  774. panic(err)
  775. }
  776. return r, err
  777. }
  778. func (cn *conn) send(m *writeBuf) {
  779. _, err := cn.c.Write(m.wrap())
  780. if err != nil {
  781. panic(err)
  782. }
  783. }
  784. func (cn *conn) sendStartupPacket(m *writeBuf) error {
  785. _, err := cn.c.Write((m.wrap())[1:])
  786. return err
  787. }
  788. // Send a message of type typ to the server on the other end of cn. The
  789. // message should have no payload. This method does not use the scratch
  790. // buffer.
  791. func (cn *conn) sendSimpleMessage(typ byte) (err error) {
  792. _, err = cn.c.Write([]byte{typ, '\x00', '\x00', '\x00', '\x04'})
  793. return err
  794. }
  795. // saveMessage memorizes a message and its buffer in the conn struct.
  796. // recvMessage will then return these values on the next call to it. This
  797. // method is useful in cases where you have to see what the next message is
  798. // going to be (e.g. to see whether it's an error or not) but you can't handle
  799. // the message yourself.
  800. func (cn *conn) saveMessage(typ byte, buf *readBuf) {
  801. if cn.saveMessageType != 0 {
  802. cn.bad = true
  803. errorf("unexpected saveMessageType %d", cn.saveMessageType)
  804. }
  805. cn.saveMessageType = typ
  806. cn.saveMessageBuffer = *buf
  807. }
  808. // recvMessage receives any message from the backend, or returns an error if
  809. // a problem occurred while reading the message.
  810. func (cn *conn) recvMessage(r *readBuf) (byte, error) {
  811. // workaround for a QueryRow bug, see exec
  812. if cn.saveMessageType != 0 {
  813. t := cn.saveMessageType
  814. *r = cn.saveMessageBuffer
  815. cn.saveMessageType = 0
  816. cn.saveMessageBuffer = nil
  817. return t, nil
  818. }
  819. x := cn.scratch[:5]
  820. _, err := io.ReadFull(cn.buf, x)
  821. if err != nil {
  822. return 0, err
  823. }
  824. // read the type and length of the message that follows
  825. t := x[0]
  826. n := int(binary.BigEndian.Uint32(x[1:])) - 4
  827. var y []byte
  828. if n <= len(cn.scratch) {
  829. y = cn.scratch[:n]
  830. } else {
  831. y = make([]byte, n)
  832. }
  833. _, err = io.ReadFull(cn.buf, y)
  834. if err != nil {
  835. return 0, err
  836. }
  837. *r = y
  838. return t, nil
  839. }
  840. // recv receives a message from the backend, but if an error happened while
  841. // reading the message or the received message was an ErrorResponse, it panics.
  842. // NoticeResponses are ignored. This function should generally be used only
  843. // during the startup sequence.
  844. func (cn *conn) recv() (t byte, r *readBuf) {
  845. for {
  846. var err error
  847. r = &readBuf{}
  848. t, err = cn.recvMessage(r)
  849. if err != nil {
  850. panic(err)
  851. }
  852. switch t {
  853. case 'E':
  854. panic(parseError(r))
  855. case 'N':
  856. // ignore
  857. default:
  858. return
  859. }
  860. }
  861. }
  862. // recv1Buf is exactly equivalent to recv1, except it uses a buffer supplied by
  863. // the caller to avoid an allocation.
  864. func (cn *conn) recv1Buf(r *readBuf) byte {
  865. for {
  866. t, err := cn.recvMessage(r)
  867. if err != nil {
  868. panic(err)
  869. }
  870. switch t {
  871. case 'A', 'N':
  872. // ignore
  873. case 'S':
  874. cn.processParameterStatus(r)
  875. default:
  876. return t
  877. }
  878. }
  879. }
  880. // recv1 receives a message from the backend, panicking if an error occurs
  881. // while attempting to read it. All asynchronous messages are ignored, with
  882. // the exception of ErrorResponse.
  883. func (cn *conn) recv1() (t byte, r *readBuf) {
  884. r = &readBuf{}
  885. t = cn.recv1Buf(r)
  886. return t, r
  887. }
  888. func (cn *conn) ssl(o values) {
  889. upgrade := ssl(o)
  890. if upgrade == nil {
  891. // Nothing to do
  892. return
  893. }
  894. w := cn.writeBuf(0)
  895. w.int32(80877103)
  896. if err := cn.sendStartupPacket(w); err != nil {
  897. panic(err)
  898. }
  899. b := cn.scratch[:1]
  900. _, err := io.ReadFull(cn.c, b)
  901. if err != nil {
  902. panic(err)
  903. }
  904. if b[0] != 'S' {
  905. panic(ErrSSLNotSupported)
  906. }
  907. cn.c = upgrade(cn.c)
  908. }
  909. // isDriverSetting returns true iff a setting is purely for configuring the
  910. // driver's options and should not be sent to the server in the connection
  911. // startup packet.
  912. func isDriverSetting(key string) bool {
  913. switch key {
  914. case "host", "port":
  915. return true
  916. case "password":
  917. return true
  918. case "sslmode", "sslcert", "sslkey", "sslrootcert":
  919. return true
  920. case "fallback_application_name":
  921. return true
  922. case "connect_timeout":
  923. return true
  924. case "disable_prepared_binary_result":
  925. return true
  926. case "binary_parameters":
  927. return true
  928. default:
  929. return false
  930. }
  931. }
  932. func (cn *conn) startup(o values) {
  933. w := cn.writeBuf(0)
  934. w.int32(196608)
  935. // Send the backend the name of the database we want to connect to, and the
  936. // user we want to connect as. Additionally, we send over any run-time
  937. // parameters potentially included in the connection string. If the server
  938. // doesn't recognize any of them, it will reply with an error.
  939. for k, v := range o {
  940. if isDriverSetting(k) {
  941. // skip options which can't be run-time parameters
  942. continue
  943. }
  944. // The protocol requires us to supply the database name as "database"
  945. // instead of "dbname".
  946. if k == "dbname" {
  947. k = "database"
  948. }
  949. w.string(k)
  950. w.string(v)
  951. }
  952. w.string("")
  953. if err := cn.sendStartupPacket(w); err != nil {
  954. panic(err)
  955. }
  956. for {
  957. t, r := cn.recv()
  958. switch t {
  959. case 'K':
  960. cn.processBackendKeyData(r)
  961. case 'S':
  962. cn.processParameterStatus(r)
  963. case 'R':
  964. cn.auth(r, o)
  965. case 'Z':
  966. cn.processReadyForQuery(r)
  967. return
  968. default:
  969. errorf("unknown response for startup: %q", t)
  970. }
  971. }
  972. }
  973. func (cn *conn) auth(r *readBuf, o values) {
  974. switch code := r.int32(); code {
  975. case 0:
  976. // OK
  977. case 3:
  978. w := cn.writeBuf('p')
  979. w.string(o["password"])
  980. cn.send(w)
  981. t, r := cn.recv()
  982. if t != 'R' {
  983. errorf("unexpected password response: %q", t)
  984. }
  985. if r.int32() != 0 {
  986. errorf("unexpected authentication response: %q", t)
  987. }
  988. case 5:
  989. s := string(r.next(4))
  990. w := cn.writeBuf('p')
  991. w.string("md5" + md5s(md5s(o["password"]+o["user"])+s))
  992. cn.send(w)
  993. t, r := cn.recv()
  994. if t != 'R' {
  995. errorf("unexpected password response: %q", t)
  996. }
  997. if r.int32() != 0 {
  998. errorf("unexpected authentication response: %q", t)
  999. }
  1000. default:
  1001. errorf("unknown authentication response: %d", code)
  1002. }
  1003. }
  1004. type format int
  1005. const formatText format = 0
  1006. const formatBinary format = 1
  1007. // One result-column format code with the value 1 (i.e. all binary).
  1008. var colFmtDataAllBinary = []byte{0, 1, 0, 1}
  1009. // No result-column format codes (i.e. all text).
  1010. var colFmtDataAllText = []byte{0, 0}
  1011. type stmt struct {
  1012. cn *conn
  1013. name string
  1014. colNames []string
  1015. colFmts []format
  1016. colFmtData []byte
  1017. colTyps []fieldDesc
  1018. paramTyps []oid.Oid
  1019. closed bool
  1020. }
  1021. func (st *stmt) Close() (err error) {
  1022. if st.closed {
  1023. return nil
  1024. }
  1025. if st.cn.bad {
  1026. return driver.ErrBadConn
  1027. }
  1028. defer st.cn.errRecover(&err)
  1029. w := st.cn.writeBuf('C')
  1030. w.byte('S')
  1031. w.string(st.name)
  1032. st.cn.send(w)
  1033. st.cn.send(st.cn.writeBuf('S'))
  1034. t, _ := st.cn.recv1()
  1035. if t != '3' {
  1036. st.cn.bad = true
  1037. errorf("unexpected close response: %q", t)
  1038. }
  1039. st.closed = true
  1040. t, r := st.cn.recv1()
  1041. if t != 'Z' {
  1042. st.cn.bad = true
  1043. errorf("expected ready for query, but got: %q", t)
  1044. }
  1045. st.cn.processReadyForQuery(r)
  1046. return nil
  1047. }
  1048. func (st *stmt) Query(v []driver.Value) (r driver.Rows, err error) {
  1049. if st.cn.bad {
  1050. return nil, driver.ErrBadConn
  1051. }
  1052. defer st.cn.errRecover(&err)
  1053. st.exec(v)
  1054. return &rows{
  1055. cn: st.cn,
  1056. colNames: st.colNames,
  1057. colTyps: st.colTyps,
  1058. colFmts: st.colFmts,
  1059. }, nil
  1060. }
  1061. func (st *stmt) Exec(v []driver.Value) (res driver.Result, err error) {
  1062. if st.cn.bad {
  1063. return nil, driver.ErrBadConn
  1064. }
  1065. defer st.cn.errRecover(&err)
  1066. st.exec(v)
  1067. res, _, err = st.cn.readExecuteResponse("simple query")
  1068. return res, err
  1069. }
  1070. func (st *stmt) exec(v []driver.Value) {
  1071. if len(v) >= 65536 {
  1072. errorf("got %d parameters but PostgreSQL only supports 65535 parameters", len(v))
  1073. }
  1074. if len(v) != len(st.paramTyps) {
  1075. errorf("got %d parameters but the statement requires %d", len(v), len(st.paramTyps))
  1076. }
  1077. cn := st.cn
  1078. w := cn.writeBuf('B')
  1079. w.byte(0) // unnamed portal
  1080. w.string(st.name)
  1081. if cn.binaryParameters {
  1082. cn.sendBinaryParameters(w, v)
  1083. } else {
  1084. w.int16(0)
  1085. w.int16(len(v))
  1086. for i, x := range v {
  1087. if x == nil {
  1088. w.int32(-1)
  1089. } else {
  1090. b := encode(&cn.parameterStatus, x, st.paramTyps[i])
  1091. w.int32(len(b))
  1092. w.bytes(b)
  1093. }
  1094. }
  1095. }
  1096. w.bytes(st.colFmtData)
  1097. w.next('E')
  1098. w.byte(0)
  1099. w.int32(0)
  1100. w.next('S')
  1101. cn.send(w)
  1102. cn.readBindResponse()
  1103. cn.postExecuteWorkaround()
  1104. }
  1105. func (st *stmt) NumInput() int {
  1106. return len(st.paramTyps)
  1107. }
  1108. // parseComplete parses the "command tag" from a CommandComplete message, and
  1109. // returns the number of rows affected (if applicable) and a string
  1110. // identifying only the command that was executed, e.g. "ALTER TABLE". If the
  1111. // command tag could not be parsed, parseComplete panics.
  1112. func (cn *conn) parseComplete(commandTag string) (driver.Result, string) {
  1113. commandsWithAffectedRows := []string{
  1114. "SELECT ",
  1115. // INSERT is handled below
  1116. "UPDATE ",
  1117. "DELETE ",
  1118. "FETCH ",
  1119. "MOVE ",
  1120. "COPY ",
  1121. }
  1122. var affectedRows *string
  1123. for _, tag := range commandsWithAffectedRows {
  1124. if strings.HasPrefix(commandTag, tag) {
  1125. t := commandTag[len(tag):]
  1126. affectedRows = &t
  1127. commandTag = tag[:len(tag)-1]
  1128. break
  1129. }
  1130. }
  1131. // INSERT also includes the oid of the inserted row in its command tag.
  1132. // Oids in user tables are deprecated, and the oid is only returned when
  1133. // exactly one row is inserted, so it's unlikely to be of value to any
  1134. // real-world application and we can ignore it.
  1135. if affectedRows == nil && strings.HasPrefix(commandTag, "INSERT ") {
  1136. parts := strings.Split(commandTag, " ")
  1137. if len(parts) != 3 {
  1138. cn.bad = true
  1139. errorf("unexpected INSERT command tag %s", commandTag)
  1140. }
  1141. affectedRows = &parts[len(parts)-1]
  1142. commandTag = "INSERT"
  1143. }
  1144. // There should be no affected rows attached to the tag, just return it
  1145. if affectedRows == nil {
  1146. return driver.RowsAffected(0), commandTag
  1147. }
  1148. n, err := strconv.ParseInt(*affectedRows, 10, 64)
  1149. if err != nil {
  1150. cn.bad = true
  1151. errorf("could not parse commandTag: %s", err)
  1152. }
  1153. return driver.RowsAffected(n), commandTag
  1154. }
  1155. type rows struct {
  1156. cn *conn
  1157. finish func()
  1158. colNames []string
  1159. colTyps []fieldDesc
  1160. colFmts []format
  1161. done bool
  1162. rb readBuf
  1163. result driver.Result
  1164. tag string
  1165. }
  1166. func (rs *rows) Close() error {
  1167. if finish := rs.finish; finish != nil {
  1168. defer finish()
  1169. }
  1170. // no need to look at cn.bad as Next() will
  1171. for {
  1172. err := rs.Next(nil)
  1173. switch err {
  1174. case nil:
  1175. case io.EOF:
  1176. // rs.Next can return io.EOF on both 'Z' (ready for query) and 'T' (row
  1177. // description, used with HasNextResultSet). We need to fetch messages until
  1178. // we hit a 'Z', which is done by waiting for done to be set.
  1179. if rs.done {
  1180. return nil
  1181. }
  1182. default:
  1183. return err
  1184. }
  1185. }
  1186. }
  1187. func (rs *rows) Columns() []string {
  1188. return rs.colNames
  1189. }
  1190. func (rs *rows) Result() driver.Result {
  1191. if rs.result == nil {
  1192. return emptyRows
  1193. }
  1194. return rs.result
  1195. }
  1196. func (rs *rows) Tag() string {
  1197. return rs.tag
  1198. }
  1199. func (rs *rows) Next(dest []driver.Value) (err error) {
  1200. if rs.done {
  1201. return io.EOF
  1202. }
  1203. conn := rs.cn
  1204. if conn.bad {
  1205. return driver.ErrBadConn
  1206. }
  1207. defer conn.errRecover(&err)
  1208. for {
  1209. t := conn.recv1Buf(&rs.rb)
  1210. switch t {
  1211. case 'E':
  1212. err = parseError(&rs.rb)
  1213. case 'C', 'I':
  1214. if t == 'C' {
  1215. rs.result, rs.tag = conn.parseComplete(rs.rb.string())
  1216. }
  1217. continue
  1218. case 'Z':
  1219. conn.processReadyForQuery(&rs.rb)
  1220. rs.done = true
  1221. if err != nil {
  1222. return err
  1223. }
  1224. return io.EOF
  1225. case 'D':
  1226. n := rs.rb.int16()
  1227. if err != nil {
  1228. conn.bad = true
  1229. errorf("unexpected DataRow after error %s", err)
  1230. }
  1231. if n < len(dest) {
  1232. dest = dest[:n]
  1233. }
  1234. for i := range dest {
  1235. l := rs.rb.int32()
  1236. if l == -1 {
  1237. dest[i] = nil
  1238. continue
  1239. }
  1240. dest[i] = decode(&conn.parameterStatus, rs.rb.next(l), rs.colTyps[i].OID, rs.colFmts[i])
  1241. }
  1242. return
  1243. case 'T':
  1244. rs.colNames, rs.colFmts, rs.colTyps = parsePortalRowDescribe(&rs.rb)
  1245. return io.EOF
  1246. default:
  1247. errorf("unexpected message after execute: %q", t)
  1248. }
  1249. }
  1250. }
  1251. func (rs *rows) HasNextResultSet() bool {
  1252. return !rs.done
  1253. }
  1254. func (rs *rows) NextResultSet() error {
  1255. return nil
  1256. }
  1257. // QuoteIdentifier quotes an "identifier" (e.g. a table or a column name) to be
  1258. // used as part of an SQL statement. For example:
  1259. //
  1260. // tblname := "my_table"
  1261. // data := "my_data"
  1262. // err = db.Exec(fmt.Sprintf("INSERT INTO %s VALUES ($1)", pq.QuoteIdentifier(tblname)), data)
  1263. //
  1264. // Any double quotes in name will be escaped. The quoted identifier will be
  1265. // case sensitive when used in a query. If the input string contains a zero
  1266. // byte, the result will be truncated immediately before it.
  1267. func QuoteIdentifier(name string) string {
  1268. end := strings.IndexRune(name, 0)
  1269. if end > -1 {
  1270. name = name[:end]
  1271. }
  1272. return `"` + strings.Replace(name, `"`, `""`, -1) + `"`
  1273. }
  1274. func md5s(s string) string {
  1275. h := md5.New()
  1276. h.Write([]byte(s))
  1277. return fmt.Sprintf("%x", h.Sum(nil))
  1278. }
  1279. func (cn *conn) sendBinaryParameters(b *writeBuf, args []driver.Value) {
  1280. // Do one pass over the parameters to see if we're going to send any of
  1281. // them over in binary. If we are, create a paramFormats array at the
  1282. // same time.
  1283. var paramFormats []int
  1284. for i, x := range args {
  1285. _, ok := x.([]byte)
  1286. if ok {
  1287. if paramFormats == nil {
  1288. paramFormats = make([]int, len(args))
  1289. }
  1290. paramFormats[i] = 1
  1291. }
  1292. }
  1293. if paramFormats == nil {
  1294. b.int16(0)
  1295. } else {
  1296. b.int16(len(paramFormats))
  1297. for _, x := range paramFormats {
  1298. b.int16(x)
  1299. }
  1300. }
  1301. b.int16(len(args))
  1302. for _, x := range args {
  1303. if x == nil {
  1304. b.int32(-1)
  1305. } else {
  1306. datum := binaryEncode(&cn.parameterStatus, x)
  1307. b.int32(len(datum))
  1308. b.bytes(datum)
  1309. }
  1310. }
  1311. }
  1312. func (cn *conn) sendBinaryModeQuery(query string, args []driver.Value) {
  1313. if len(args) >= 65536 {
  1314. errorf("got %d parameters but PostgreSQL only supports 65535 parameters", len(args))
  1315. }
  1316. b := cn.writeBuf('P')
  1317. b.byte(0) // unnamed statement
  1318. b.string(query)
  1319. b.int16(0)
  1320. b.next('B')
  1321. b.int16(0) // unnamed portal and statement
  1322. cn.sendBinaryParameters(b, args)
  1323. b.bytes(colFmtDataAllText)
  1324. b.next('D')
  1325. b.byte('P')
  1326. b.byte(0) // unnamed portal
  1327. b.next('E')
  1328. b.byte(0)
  1329. b.int32(0)
  1330. b.next('S')
  1331. cn.send(b)
  1332. }
  1333. func (cn *conn) processParameterStatus(r *readBuf) {
  1334. var err error
  1335. param := r.string()
  1336. switch param {
  1337. case "server_version":
  1338. var major1 int
  1339. var major2 int
  1340. var minor int
  1341. _, err = fmt.Sscanf(r.string(), "%d.%d.%d", &major1, &major2, &minor)
  1342. if err == nil {
  1343. cn.parameterStatus.serverVersion = major1*10000 + major2*100 + minor
  1344. }
  1345. case "TimeZone":
  1346. cn.parameterStatus.currentLocation, err = time.LoadLocation(r.string())
  1347. if err != nil {
  1348. cn.parameterStatus.currentLocation = nil
  1349. }
  1350. default:
  1351. // ignore
  1352. }
  1353. }
  1354. func (cn *conn) processReadyForQuery(r *readBuf) {
  1355. cn.txnStatus = transactionStatus(r.byte())
  1356. }
  1357. func (cn *conn) readReadyForQuery() {
  1358. t, r := cn.recv1()
  1359. switch t {
  1360. case 'Z':
  1361. cn.processReadyForQuery(r)
  1362. return
  1363. default:
  1364. cn.bad = true
  1365. errorf("unexpected message %q; expected ReadyForQuery", t)
  1366. }
  1367. }
  1368. func (cn *conn) processBackendKeyData(r *readBuf) {
  1369. cn.processID = r.int32()
  1370. cn.secretKey = r.int32()
  1371. }
  1372. func (cn *conn) readParseResponse() {
  1373. t, r := cn.recv1()
  1374. switch t {
  1375. case '1':
  1376. return
  1377. case 'E':
  1378. err := parseError(r)
  1379. cn.readReadyForQuery()
  1380. panic(err)
  1381. default:
  1382. cn.bad = true
  1383. errorf("unexpected Parse response %q", t)
  1384. }
  1385. }
  1386. func (cn *conn) readStatementDescribeResponse() (paramTyps []oid.Oid, colNames []string, colTyps []fieldDesc) {
  1387. for {
  1388. t, r := cn.recv1()
  1389. switch t {
  1390. case 't':
  1391. nparams := r.int16()
  1392. paramTyps = make([]oid.Oid, nparams)
  1393. for i := range paramTyps {
  1394. paramTyps[i] = r.oid()
  1395. }
  1396. case 'n':
  1397. return paramTyps, nil, nil
  1398. case 'T':
  1399. colNames, colTyps = parseStatementRowDescribe(r)
  1400. return paramTyps, colNames, colTyps
  1401. case 'E':
  1402. err := parseError(r)
  1403. cn.readReadyForQuery()
  1404. panic(err)
  1405. default:
  1406. cn.bad = true
  1407. errorf("unexpected Describe statement response %q", t)
  1408. }
  1409. }
  1410. }
  1411. func (cn *conn) readPortalDescribeResponse() (colNames []string, colFmts []format, colTyps []fieldDesc) {
  1412. t, r := cn.recv1()
  1413. switch t {
  1414. case 'T':
  1415. return parsePortalRowDescribe(r)
  1416. case 'n':
  1417. return nil, nil, nil
  1418. case 'E':
  1419. err := parseError(r)
  1420. cn.readReadyForQuery()
  1421. panic(err)
  1422. default:
  1423. cn.bad = true
  1424. errorf("unexpected Describe response %q", t)
  1425. }
  1426. panic("not reached")
  1427. }
  1428. func (cn *conn) readBindResponse() {
  1429. t, r := cn.recv1()
  1430. switch t {
  1431. case '2':
  1432. return
  1433. case 'E':
  1434. err := parseError(r)
  1435. cn.readReadyForQuery()
  1436. panic(err)
  1437. default:
  1438. cn.bad = true
  1439. errorf("unexpected Bind response %q", t)
  1440. }
  1441. }
  1442. func (cn *conn) postExecuteWorkaround() {
  1443. // Work around a bug in sql.DB.QueryRow: in Go 1.2 and earlier it ignores
  1444. // any errors from rows.Next, which masks errors that happened during the
  1445. // execution of the query. To avoid the problem in common cases, we wait
  1446. // here for one more message from the database. If it's not an error the
  1447. // query will likely succeed (or perhaps has already, if it's a
  1448. // CommandComplete), so we push the message into the conn struct; recv1
  1449. // will return it as the next message for rows.Next or rows.Close.
  1450. // However, if it's an error, we wait until ReadyForQuery and then return
  1451. // the error to our caller.
  1452. for {
  1453. t, r := cn.recv1()
  1454. switch t {
  1455. case 'E':
  1456. err := parseError(r)
  1457. cn.readReadyForQuery()
  1458. panic(err)
  1459. case 'C', 'D', 'I':
  1460. // the query didn't fail, but we can't process this message
  1461. cn.saveMessage(t, r)
  1462. return
  1463. default:
  1464. cn.bad = true
  1465. errorf("unexpected message during extended query execution: %q", t)
  1466. }
  1467. }
  1468. }
  1469. // Only for Exec(), since we ignore the returned data
  1470. func (cn *conn) readExecuteResponse(protocolState string) (res driver.Result, commandTag string, err error) {
  1471. for {
  1472. t, r := cn.recv1()
  1473. switch t {
  1474. case 'C':
  1475. if err != nil {
  1476. cn.bad = true
  1477. errorf("unexpected CommandComplete after error %s", err)
  1478. }
  1479. res, commandTag = cn.parseComplete(r.string())
  1480. case 'Z':
  1481. cn.processReadyForQuery(r)
  1482. if res == nil && err == nil {
  1483. err = errUnexpectedReady
  1484. }
  1485. return res, commandTag, err
  1486. case 'E':
  1487. err = parseError(r)
  1488. case 'T', 'D', 'I':
  1489. if err != nil {
  1490. cn.bad = true
  1491. errorf("unexpected %q after error %s", t, err)
  1492. }
  1493. if t == 'I' {
  1494. res = emptyRows
  1495. }
  1496. // ignore any results
  1497. default:
  1498. cn.bad = true
  1499. errorf("unknown %s response: %q", protocolState, t)
  1500. }
  1501. }
  1502. }
  1503. func parseStatementRowDescribe(r *readBuf) (colNames []string, colTyps []fieldDesc) {
  1504. n := r.int16()
  1505. colNames = make([]string, n)
  1506. colTyps = make([]fieldDesc, n)
  1507. for i := range colNames {
  1508. colNames[i] = r.string()
  1509. r.next(6)
  1510. colTyps[i].OID = r.oid()
  1511. colTyps[i].Len = r.int16()
  1512. colTyps[i].Mod = r.int32()
  1513. // format code not known when describing a statement; always 0
  1514. r.next(2)
  1515. }
  1516. return
  1517. }
  1518. func parsePortalRowDescribe(r *readBuf) (colNames []string, colFmts []format, colTyps []fieldDesc) {
  1519. n := r.int16()
  1520. colNames = make([]string, n)
  1521. colFmts = make([]format, n)
  1522. colTyps = make([]fieldDesc, n)
  1523. for i := range colNames {
  1524. colNames[i] = r.string()
  1525. r.next(6)
  1526. colTyps[i].OID = r.oid()
  1527. colTyps[i].Len = r.int16()
  1528. colTyps[i].Mod = r.int32()
  1529. colFmts[i] = format(r.int16())
  1530. }
  1531. return
  1532. }
  1533. // parseEnviron tries to mimic some of libpq's environment handling
  1534. //
  1535. // To ease testing, it does not directly reference os.Environ, but is
  1536. // designed to accept its output.
  1537. //
  1538. // Environment-set connection information is intended to have a higher
  1539. // precedence than a library default but lower than any explicitly
  1540. // passed information (such as in the URL or connection string).
  1541. func parseEnviron(env []string) (out map[string]string) {
  1542. out = make(map[string]string)
  1543. for _, v := range env {
  1544. parts := strings.SplitN(v, "=", 2)
  1545. accrue := func(keyname string) {
  1546. out[keyname] = parts[1]
  1547. }
  1548. unsupported := func() {
  1549. panic(fmt.Sprintf("setting %v not supported", parts[0]))
  1550. }
  1551. // The order of these is the same as is seen in the
  1552. // PostgreSQL 9.1 manual. Unsupported but well-defined
  1553. // keys cause a panic; these should be unset prior to
  1554. // execution. Options which pq expects to be set to a
  1555. // certain value are allowed, but must be set to that
  1556. // value if present (they can, of course, be absent).
  1557. switch parts[0] {
  1558. case "PGHOST":
  1559. accrue("host")
  1560. case "PGHOSTADDR":
  1561. unsupported()
  1562. case "PGPORT":
  1563. accrue("port")
  1564. case "PGDATABASE":
  1565. accrue("dbname")
  1566. case "PGUSER":
  1567. accrue("user")
  1568. case "PGPASSWORD":
  1569. accrue("password")
  1570. case "PGSERVICE", "PGSERVICEFILE", "PGREALM":
  1571. unsupported()
  1572. case "PGOPTIONS":
  1573. accrue("options")
  1574. case "PGAPPNAME":
  1575. accrue("application_name")
  1576. case "PGSSLMODE":
  1577. accrue("sslmode")
  1578. case "PGSSLCERT":
  1579. accrue("sslcert")
  1580. case "PGSSLKEY":
  1581. accrue("sslkey")
  1582. case "PGSSLROOTCERT":
  1583. accrue("sslrootcert")
  1584. case "PGREQUIRESSL", "PGSSLCRL":
  1585. unsupported()
  1586. case "PGREQUIREPEER":
  1587. unsupported()
  1588. case "PGKRBSRVNAME", "PGGSSLIB":
  1589. unsupported()
  1590. case "PGCONNECT_TIMEOUT":
  1591. accrue("connect_timeout")
  1592. case "PGCLIENTENCODING":
  1593. accrue("client_encoding")
  1594. case "PGDATESTYLE":
  1595. accrue("datestyle")
  1596. case "PGTZ":
  1597. accrue("timezone")
  1598. case "PGGEQO":
  1599. accrue("geqo")
  1600. case "PGSYSCONFDIR", "PGLOCALEDIR":
  1601. unsupported()
  1602. }
  1603. }
  1604. return out
  1605. }
  1606. // isUTF8 returns whether name is a fuzzy variation of the string "UTF-8".
  1607. func isUTF8(name string) bool {
  1608. // Recognize all sorts of silly things as "UTF-8", like Postgres does
  1609. s := strings.Map(alnumLowerASCII, name)
  1610. return s == "utf8" || s == "unicode"
  1611. }
  1612. func alnumLowerASCII(ch rune) rune {
  1613. if 'A' <= ch && ch <= 'Z' {
  1614. return ch + ('a' - 'A')
  1615. }
  1616. if 'a' <= ch && ch <= 'z' || '0' <= ch && ch <= '9' {
  1617. return ch
  1618. }
  1619. return -1 // discard
  1620. }