conn.go 47 KB

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