本站源代码
Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.

603 linhas
15KB

  1. package pq
  2. import (
  3. "bytes"
  4. "database/sql/driver"
  5. "encoding/binary"
  6. "encoding/hex"
  7. "errors"
  8. "fmt"
  9. "math"
  10. "strconv"
  11. "strings"
  12. "sync"
  13. "time"
  14. "github.com/lib/pq/oid"
  15. )
  16. func binaryEncode(parameterStatus *parameterStatus, x interface{}) []byte {
  17. switch v := x.(type) {
  18. case []byte:
  19. return v
  20. default:
  21. return encode(parameterStatus, x, oid.T_unknown)
  22. }
  23. }
  24. func encode(parameterStatus *parameterStatus, x interface{}, pgtypOid oid.Oid) []byte {
  25. switch v := x.(type) {
  26. case int64:
  27. return strconv.AppendInt(nil, v, 10)
  28. case float64:
  29. return strconv.AppendFloat(nil, v, 'f', -1, 64)
  30. case []byte:
  31. if pgtypOid == oid.T_bytea {
  32. return encodeBytea(parameterStatus.serverVersion, v)
  33. }
  34. return v
  35. case string:
  36. if pgtypOid == oid.T_bytea {
  37. return encodeBytea(parameterStatus.serverVersion, []byte(v))
  38. }
  39. return []byte(v)
  40. case bool:
  41. return strconv.AppendBool(nil, v)
  42. case time.Time:
  43. return formatTs(v)
  44. default:
  45. errorf("encode: unknown type for %T", v)
  46. }
  47. panic("not reached")
  48. }
  49. func decode(parameterStatus *parameterStatus, s []byte, typ oid.Oid, f format) interface{} {
  50. switch f {
  51. case formatBinary:
  52. return binaryDecode(parameterStatus, s, typ)
  53. case formatText:
  54. return textDecode(parameterStatus, s, typ)
  55. default:
  56. panic("not reached")
  57. }
  58. }
  59. func binaryDecode(parameterStatus *parameterStatus, s []byte, typ oid.Oid) interface{} {
  60. switch typ {
  61. case oid.T_bytea:
  62. return s
  63. case oid.T_int8:
  64. return int64(binary.BigEndian.Uint64(s))
  65. case oid.T_int4:
  66. return int64(int32(binary.BigEndian.Uint32(s)))
  67. case oid.T_int2:
  68. return int64(int16(binary.BigEndian.Uint16(s)))
  69. case oid.T_uuid:
  70. b, err := decodeUUIDBinary(s)
  71. if err != nil {
  72. panic(err)
  73. }
  74. return b
  75. default:
  76. errorf("don't know how to decode binary parameter of type %d", uint32(typ))
  77. }
  78. panic("not reached")
  79. }
  80. func textDecode(parameterStatus *parameterStatus, s []byte, typ oid.Oid) interface{} {
  81. switch typ {
  82. case oid.T_char, oid.T_varchar, oid.T_text:
  83. return string(s)
  84. case oid.T_bytea:
  85. b, err := parseBytea(s)
  86. if err != nil {
  87. errorf("%s", err)
  88. }
  89. return b
  90. case oid.T_timestamptz:
  91. return parseTs(parameterStatus.currentLocation, string(s))
  92. case oid.T_timestamp, oid.T_date:
  93. return parseTs(nil, string(s))
  94. case oid.T_time:
  95. return mustParse("15:04:05", typ, s)
  96. case oid.T_timetz:
  97. return mustParse("15:04:05-07", typ, s)
  98. case oid.T_bool:
  99. return s[0] == 't'
  100. case oid.T_int8, oid.T_int4, oid.T_int2:
  101. i, err := strconv.ParseInt(string(s), 10, 64)
  102. if err != nil {
  103. errorf("%s", err)
  104. }
  105. return i
  106. case oid.T_float4, oid.T_float8:
  107. // We always use 64 bit parsing, regardless of whether the input text is for
  108. // a float4 or float8, because clients expect float64s for all float datatypes
  109. // and returning a 32-bit parsed float64 produces lossy results.
  110. f, err := strconv.ParseFloat(string(s), 64)
  111. if err != nil {
  112. errorf("%s", err)
  113. }
  114. return f
  115. }
  116. return s
  117. }
  118. // appendEncodedText encodes item in text format as required by COPY
  119. // and appends to buf
  120. func appendEncodedText(parameterStatus *parameterStatus, buf []byte, x interface{}) []byte {
  121. switch v := x.(type) {
  122. case int64:
  123. return strconv.AppendInt(buf, v, 10)
  124. case float64:
  125. return strconv.AppendFloat(buf, v, 'f', -1, 64)
  126. case []byte:
  127. encodedBytea := encodeBytea(parameterStatus.serverVersion, v)
  128. return appendEscapedText(buf, string(encodedBytea))
  129. case string:
  130. return appendEscapedText(buf, v)
  131. case bool:
  132. return strconv.AppendBool(buf, v)
  133. case time.Time:
  134. return append(buf, formatTs(v)...)
  135. case nil:
  136. return append(buf, "\\N"...)
  137. default:
  138. errorf("encode: unknown type for %T", v)
  139. }
  140. panic("not reached")
  141. }
  142. func appendEscapedText(buf []byte, text string) []byte {
  143. escapeNeeded := false
  144. startPos := 0
  145. var c byte
  146. // check if we need to escape
  147. for i := 0; i < len(text); i++ {
  148. c = text[i]
  149. if c == '\\' || c == '\n' || c == '\r' || c == '\t' {
  150. escapeNeeded = true
  151. startPos = i
  152. break
  153. }
  154. }
  155. if !escapeNeeded {
  156. return append(buf, text...)
  157. }
  158. // copy till first char to escape, iterate the rest
  159. result := append(buf, text[:startPos]...)
  160. for i := startPos; i < len(text); i++ {
  161. c = text[i]
  162. switch c {
  163. case '\\':
  164. result = append(result, '\\', '\\')
  165. case '\n':
  166. result = append(result, '\\', 'n')
  167. case '\r':
  168. result = append(result, '\\', 'r')
  169. case '\t':
  170. result = append(result, '\\', 't')
  171. default:
  172. result = append(result, c)
  173. }
  174. }
  175. return result
  176. }
  177. func mustParse(f string, typ oid.Oid, s []byte) time.Time {
  178. str := string(s)
  179. // check for a 30-minute-offset timezone
  180. if (typ == oid.T_timestamptz || typ == oid.T_timetz) &&
  181. str[len(str)-3] == ':' {
  182. f += ":00"
  183. }
  184. t, err := time.Parse(f, str)
  185. if err != nil {
  186. errorf("decode: %s", err)
  187. }
  188. return t
  189. }
  190. var errInvalidTimestamp = errors.New("invalid timestamp")
  191. type timestampParser struct {
  192. err error
  193. }
  194. func (p *timestampParser) expect(str string, char byte, pos int) {
  195. if p.err != nil {
  196. return
  197. }
  198. if pos+1 > len(str) {
  199. p.err = errInvalidTimestamp
  200. return
  201. }
  202. if c := str[pos]; c != char && p.err == nil {
  203. p.err = fmt.Errorf("expected '%v' at position %v; got '%v'", char, pos, c)
  204. }
  205. }
  206. func (p *timestampParser) mustAtoi(str string, begin int, end int) int {
  207. if p.err != nil {
  208. return 0
  209. }
  210. if begin < 0 || end < 0 || begin > end || end > len(str) {
  211. p.err = errInvalidTimestamp
  212. return 0
  213. }
  214. result, err := strconv.Atoi(str[begin:end])
  215. if err != nil {
  216. if p.err == nil {
  217. p.err = fmt.Errorf("expected number; got '%v'", str)
  218. }
  219. return 0
  220. }
  221. return result
  222. }
  223. // The location cache caches the time zones typically used by the client.
  224. type locationCache struct {
  225. cache map[int]*time.Location
  226. lock sync.Mutex
  227. }
  228. // All connections share the same list of timezones. Benchmarking shows that
  229. // about 5% speed could be gained by putting the cache in the connection and
  230. // losing the mutex, at the cost of a small amount of memory and a somewhat
  231. // significant increase in code complexity.
  232. var globalLocationCache = newLocationCache()
  233. func newLocationCache() *locationCache {
  234. return &locationCache{cache: make(map[int]*time.Location)}
  235. }
  236. // Returns the cached timezone for the specified offset, creating and caching
  237. // it if necessary.
  238. func (c *locationCache) getLocation(offset int) *time.Location {
  239. c.lock.Lock()
  240. defer c.lock.Unlock()
  241. location, ok := c.cache[offset]
  242. if !ok {
  243. location = time.FixedZone("", offset)
  244. c.cache[offset] = location
  245. }
  246. return location
  247. }
  248. var infinityTsEnabled = false
  249. var infinityTsNegative time.Time
  250. var infinityTsPositive time.Time
  251. const (
  252. infinityTsEnabledAlready = "pq: infinity timestamp enabled already"
  253. infinityTsNegativeMustBeSmaller = "pq: infinity timestamp: negative value must be smaller (before) than positive"
  254. )
  255. // EnableInfinityTs controls the handling of Postgres' "-infinity" and
  256. // "infinity" "timestamp"s.
  257. //
  258. // If EnableInfinityTs is not called, "-infinity" and "infinity" will return
  259. // []byte("-infinity") and []byte("infinity") respectively, and potentially
  260. // cause error "sql: Scan error on column index 0: unsupported driver -> Scan
  261. // pair: []uint8 -> *time.Time", when scanning into a time.Time value.
  262. //
  263. // Once EnableInfinityTs has been called, all connections created using this
  264. // driver will decode Postgres' "-infinity" and "infinity" for "timestamp",
  265. // "timestamp with time zone" and "date" types to the predefined minimum and
  266. // maximum times, respectively. When encoding time.Time values, any time which
  267. // equals or precedes the predefined minimum time will be encoded to
  268. // "-infinity". Any values at or past the maximum time will similarly be
  269. // encoded to "infinity".
  270. //
  271. // If EnableInfinityTs is called with negative >= positive, it will panic.
  272. // Calling EnableInfinityTs after a connection has been established results in
  273. // undefined behavior. If EnableInfinityTs is called more than once, it will
  274. // panic.
  275. func EnableInfinityTs(negative time.Time, positive time.Time) {
  276. if infinityTsEnabled {
  277. panic(infinityTsEnabledAlready)
  278. }
  279. if !negative.Before(positive) {
  280. panic(infinityTsNegativeMustBeSmaller)
  281. }
  282. infinityTsEnabled = true
  283. infinityTsNegative = negative
  284. infinityTsPositive = positive
  285. }
  286. /*
  287. * Testing might want to toggle infinityTsEnabled
  288. */
  289. func disableInfinityTs() {
  290. infinityTsEnabled = false
  291. }
  292. // This is a time function specific to the Postgres default DateStyle
  293. // setting ("ISO, MDY"), the only one we currently support. This
  294. // accounts for the discrepancies between the parsing available with
  295. // time.Parse and the Postgres date formatting quirks.
  296. func parseTs(currentLocation *time.Location, str string) interface{} {
  297. switch str {
  298. case "-infinity":
  299. if infinityTsEnabled {
  300. return infinityTsNegative
  301. }
  302. return []byte(str)
  303. case "infinity":
  304. if infinityTsEnabled {
  305. return infinityTsPositive
  306. }
  307. return []byte(str)
  308. }
  309. t, err := ParseTimestamp(currentLocation, str)
  310. if err != nil {
  311. panic(err)
  312. }
  313. return t
  314. }
  315. // ParseTimestamp parses Postgres' text format. It returns a time.Time in
  316. // currentLocation iff that time's offset agrees with the offset sent from the
  317. // Postgres server. Otherwise, ParseTimestamp returns a time.Time with the
  318. // fixed offset offset provided by the Postgres server.
  319. func ParseTimestamp(currentLocation *time.Location, str string) (time.Time, error) {
  320. p := timestampParser{}
  321. monSep := strings.IndexRune(str, '-')
  322. // this is Gregorian year, not ISO Year
  323. // In Gregorian system, the year 1 BC is followed by AD 1
  324. year := p.mustAtoi(str, 0, monSep)
  325. daySep := monSep + 3
  326. month := p.mustAtoi(str, monSep+1, daySep)
  327. p.expect(str, '-', daySep)
  328. timeSep := daySep + 3
  329. day := p.mustAtoi(str, daySep+1, timeSep)
  330. minLen := monSep + len("01-01") + 1
  331. isBC := strings.HasSuffix(str, " BC")
  332. if isBC {
  333. minLen += 3
  334. }
  335. var hour, minute, second int
  336. if len(str) > minLen {
  337. p.expect(str, ' ', timeSep)
  338. minSep := timeSep + 3
  339. p.expect(str, ':', minSep)
  340. hour = p.mustAtoi(str, timeSep+1, minSep)
  341. secSep := minSep + 3
  342. p.expect(str, ':', secSep)
  343. minute = p.mustAtoi(str, minSep+1, secSep)
  344. secEnd := secSep + 3
  345. second = p.mustAtoi(str, secSep+1, secEnd)
  346. }
  347. remainderIdx := monSep + len("01-01 00:00:00") + 1
  348. // Three optional (but ordered) sections follow: the
  349. // fractional seconds, the time zone offset, and the BC
  350. // designation. We set them up here and adjust the other
  351. // offsets if the preceding sections exist.
  352. nanoSec := 0
  353. tzOff := 0
  354. if remainderIdx < len(str) && str[remainderIdx] == '.' {
  355. fracStart := remainderIdx + 1
  356. fracOff := strings.IndexAny(str[fracStart:], "-+ ")
  357. if fracOff < 0 {
  358. fracOff = len(str) - fracStart
  359. }
  360. fracSec := p.mustAtoi(str, fracStart, fracStart+fracOff)
  361. nanoSec = fracSec * (1000000000 / int(math.Pow(10, float64(fracOff))))
  362. remainderIdx += fracOff + 1
  363. }
  364. if tzStart := remainderIdx; tzStart < len(str) && (str[tzStart] == '-' || str[tzStart] == '+') {
  365. // time zone separator is always '-' or '+' (UTC is +00)
  366. var tzSign int
  367. switch c := str[tzStart]; c {
  368. case '-':
  369. tzSign = -1
  370. case '+':
  371. tzSign = +1
  372. default:
  373. return time.Time{}, fmt.Errorf("expected '-' or '+' at position %v; got %v", tzStart, c)
  374. }
  375. tzHours := p.mustAtoi(str, tzStart+1, tzStart+3)
  376. remainderIdx += 3
  377. var tzMin, tzSec int
  378. if remainderIdx < len(str) && str[remainderIdx] == ':' {
  379. tzMin = p.mustAtoi(str, remainderIdx+1, remainderIdx+3)
  380. remainderIdx += 3
  381. }
  382. if remainderIdx < len(str) && str[remainderIdx] == ':' {
  383. tzSec = p.mustAtoi(str, remainderIdx+1, remainderIdx+3)
  384. remainderIdx += 3
  385. }
  386. tzOff = tzSign * ((tzHours * 60 * 60) + (tzMin * 60) + tzSec)
  387. }
  388. var isoYear int
  389. if isBC {
  390. isoYear = 1 - year
  391. remainderIdx += 3
  392. } else {
  393. isoYear = year
  394. }
  395. if remainderIdx < len(str) {
  396. return time.Time{}, fmt.Errorf("expected end of input, got %v", str[remainderIdx:])
  397. }
  398. t := time.Date(isoYear, time.Month(month), day,
  399. hour, minute, second, nanoSec,
  400. globalLocationCache.getLocation(tzOff))
  401. if currentLocation != nil {
  402. // Set the location of the returned Time based on the session's
  403. // TimeZone value, but only if the local time zone database agrees with
  404. // the remote database on the offset.
  405. lt := t.In(currentLocation)
  406. _, newOff := lt.Zone()
  407. if newOff == tzOff {
  408. t = lt
  409. }
  410. }
  411. return t, p.err
  412. }
  413. // formatTs formats t into a format postgres understands.
  414. func formatTs(t time.Time) []byte {
  415. if infinityTsEnabled {
  416. // t <= -infinity : ! (t > -infinity)
  417. if !t.After(infinityTsNegative) {
  418. return []byte("-infinity")
  419. }
  420. // t >= infinity : ! (!t < infinity)
  421. if !t.Before(infinityTsPositive) {
  422. return []byte("infinity")
  423. }
  424. }
  425. return FormatTimestamp(t)
  426. }
  427. // FormatTimestamp formats t into Postgres' text format for timestamps.
  428. func FormatTimestamp(t time.Time) []byte {
  429. // Need to send dates before 0001 A.D. with " BC" suffix, instead of the
  430. // minus sign preferred by Go.
  431. // Beware, "0000" in ISO is "1 BC", "-0001" is "2 BC" and so on
  432. bc := false
  433. if t.Year() <= 0 {
  434. // flip year sign, and add 1, e.g: "0" will be "1", and "-10" will be "11"
  435. t = t.AddDate((-t.Year())*2+1, 0, 0)
  436. bc = true
  437. }
  438. b := []byte(t.Format("2006-01-02 15:04:05.999999999Z07:00"))
  439. _, offset := t.Zone()
  440. offset = offset % 60
  441. if offset != 0 {
  442. // RFC3339Nano already printed the minus sign
  443. if offset < 0 {
  444. offset = -offset
  445. }
  446. b = append(b, ':')
  447. if offset < 10 {
  448. b = append(b, '0')
  449. }
  450. b = strconv.AppendInt(b, int64(offset), 10)
  451. }
  452. if bc {
  453. b = append(b, " BC"...)
  454. }
  455. return b
  456. }
  457. // Parse a bytea value received from the server. Both "hex" and the legacy
  458. // "escape" format are supported.
  459. func parseBytea(s []byte) (result []byte, err error) {
  460. if len(s) >= 2 && bytes.Equal(s[:2], []byte("\\x")) {
  461. // bytea_output = hex
  462. s = s[2:] // trim off leading "\\x"
  463. result = make([]byte, hex.DecodedLen(len(s)))
  464. _, err := hex.Decode(result, s)
  465. if err != nil {
  466. return nil, err
  467. }
  468. } else {
  469. // bytea_output = escape
  470. for len(s) > 0 {
  471. if s[0] == '\\' {
  472. // escaped '\\'
  473. if len(s) >= 2 && s[1] == '\\' {
  474. result = append(result, '\\')
  475. s = s[2:]
  476. continue
  477. }
  478. // '\\' followed by an octal number
  479. if len(s) < 4 {
  480. return nil, fmt.Errorf("invalid bytea sequence %v", s)
  481. }
  482. r, err := strconv.ParseInt(string(s[1:4]), 8, 9)
  483. if err != nil {
  484. return nil, fmt.Errorf("could not parse bytea value: %s", err.Error())
  485. }
  486. result = append(result, byte(r))
  487. s = s[4:]
  488. } else {
  489. // We hit an unescaped, raw byte. Try to read in as many as
  490. // possible in one go.
  491. i := bytes.IndexByte(s, '\\')
  492. if i == -1 {
  493. result = append(result, s...)
  494. break
  495. }
  496. result = append(result, s[:i]...)
  497. s = s[i:]
  498. }
  499. }
  500. }
  501. return result, nil
  502. }
  503. func encodeBytea(serverVersion int, v []byte) (result []byte) {
  504. if serverVersion >= 90000 {
  505. // Use the hex format if we know that the server supports it
  506. result = make([]byte, 2+hex.EncodedLen(len(v)))
  507. result[0] = '\\'
  508. result[1] = 'x'
  509. hex.Encode(result[2:], v)
  510. } else {
  511. // .. or resort to "escape"
  512. for _, b := range v {
  513. if b == '\\' {
  514. result = append(result, '\\', '\\')
  515. } else if b < 0x20 || b > 0x7e {
  516. result = append(result, []byte(fmt.Sprintf("\\%03o", b))...)
  517. } else {
  518. result = append(result, b)
  519. }
  520. }
  521. }
  522. return result
  523. }
  524. // NullTime represents a time.Time that may be null. NullTime implements the
  525. // sql.Scanner interface so it can be used as a scan destination, similar to
  526. // sql.NullString.
  527. type NullTime struct {
  528. Time time.Time
  529. Valid bool // Valid is true if Time is not NULL
  530. }
  531. // Scan implements the Scanner interface.
  532. func (nt *NullTime) Scan(value interface{}) error {
  533. nt.Time, nt.Valid = value.(time.Time)
  534. return nil
  535. }
  536. // Value implements the driver Valuer interface.
  537. func (nt NullTime) Value() (driver.Value, error) {
  538. if !nt.Valid {
  539. return nil, nil
  540. }
  541. return nt.Time, nil
  542. }
上海开阖软件有限公司 沪ICP备12045867号-1