conn.go 43 KB

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