conn.go 42 KB

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