conn.go 43 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846
  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. }