conn.go 23 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010
  1. // Copyright (c) 2012 The gocql Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package gocql
  5. import (
  6. "bufio"
  7. "crypto/tls"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "io/ioutil"
  12. "log"
  13. "net"
  14. "strconv"
  15. "strings"
  16. "sync"
  17. "sync/atomic"
  18. "time"
  19. "github.com/gocql/gocql/internal/streams"
  20. )
  21. var (
  22. approvedAuthenticators = [...]string{
  23. "org.apache.cassandra.auth.PasswordAuthenticator",
  24. "com.instaclustr.cassandra.auth.SharedSecretAuthenticator",
  25. }
  26. )
  27. func approve(authenticator string) bool {
  28. for _, s := range approvedAuthenticators {
  29. if authenticator == s {
  30. return true
  31. }
  32. }
  33. return false
  34. }
  35. //JoinHostPort is a utility to return a address string that can be used
  36. //gocql.Conn to form a connection with a host.
  37. func JoinHostPort(addr string, port int) string {
  38. addr = strings.TrimSpace(addr)
  39. if _, _, err := net.SplitHostPort(addr); err != nil {
  40. addr = net.JoinHostPort(addr, strconv.Itoa(port))
  41. }
  42. return addr
  43. }
  44. type Authenticator interface {
  45. Challenge(req []byte) (resp []byte, auth Authenticator, err error)
  46. Success(data []byte) error
  47. }
  48. type PasswordAuthenticator struct {
  49. Username string
  50. Password string
  51. }
  52. func (p PasswordAuthenticator) Challenge(req []byte) ([]byte, Authenticator, error) {
  53. if !approve(string(req)) {
  54. return nil, nil, fmt.Errorf("unexpected authenticator %q", req)
  55. }
  56. resp := make([]byte, 2+len(p.Username)+len(p.Password))
  57. resp[0] = 0
  58. copy(resp[1:], p.Username)
  59. resp[len(p.Username)+1] = 0
  60. copy(resp[2+len(p.Username):], p.Password)
  61. return resp, nil, nil
  62. }
  63. func (p PasswordAuthenticator) Success(data []byte) error {
  64. return nil
  65. }
  66. type SslOptions struct {
  67. tls.Config
  68. // CertPath and KeyPath are optional depending on server
  69. // config, but both fields must be omitted to avoid using a
  70. // client certificate
  71. CertPath string
  72. KeyPath string
  73. CaPath string //optional depending on server config
  74. // If you want to verify the hostname and server cert (like a wildcard for cass cluster) then you should turn this on
  75. // This option is basically the inverse of InSecureSkipVerify
  76. // See InSecureSkipVerify in http://golang.org/pkg/crypto/tls/ for more info
  77. EnableHostVerification bool
  78. }
  79. type ConnConfig struct {
  80. ProtoVersion int
  81. CQLVersion string
  82. Timeout time.Duration
  83. Compressor Compressor
  84. Authenticator Authenticator
  85. Keepalive time.Duration
  86. tlsConfig *tls.Config
  87. }
  88. type ConnErrorHandler interface {
  89. HandleError(conn *Conn, err error, closed bool)
  90. }
  91. // How many timeouts we will allow to occur before the connection is closed
  92. // and restarted. This is to prevent a single query timeout from killing a connection
  93. // which may be serving more queries just fine.
  94. // Default is 10, should not be changed concurrently with queries.
  95. var TimeoutLimit int64 = 10
  96. // Conn is a single connection to a Cassandra node. It can be used to execute
  97. // queries, but users are usually advised to use a more reliable, higher
  98. // level API.
  99. type Conn struct {
  100. conn net.Conn
  101. r *bufio.Reader
  102. timeout time.Duration
  103. cfg *ConnConfig
  104. headerBuf []byte
  105. streams *streams.IDGenerator
  106. mu sync.RWMutex
  107. calls map[int]*callReq
  108. errorHandler ConnErrorHandler
  109. compressor Compressor
  110. auth Authenticator
  111. addr string
  112. version uint8
  113. currentKeyspace string
  114. started bool
  115. session *Session
  116. closed int32
  117. quit chan struct{}
  118. timeouts int64
  119. }
  120. // Connect establishes a connection to a Cassandra node.
  121. func Connect(addr string, cfg *ConnConfig, errorHandler ConnErrorHandler, session *Session) (*Conn, error) {
  122. var (
  123. err error
  124. conn net.Conn
  125. )
  126. dialer := &net.Dialer{
  127. Timeout: cfg.Timeout,
  128. }
  129. if cfg.tlsConfig != nil {
  130. // the TLS config is safe to be reused by connections but it must not
  131. // be modified after being used.
  132. conn, err = tls.DialWithDialer(dialer, "tcp", addr, cfg.tlsConfig)
  133. } else {
  134. conn, err = dialer.Dial("tcp", addr)
  135. }
  136. if err != nil {
  137. return nil, err
  138. }
  139. // going to default to proto 2
  140. if cfg.ProtoVersion < protoVersion1 || cfg.ProtoVersion > protoVersion4 {
  141. log.Printf("unsupported protocol version: %d using 2\n", cfg.ProtoVersion)
  142. cfg.ProtoVersion = 2
  143. }
  144. headerSize := 8
  145. if cfg.ProtoVersion > protoVersion2 {
  146. headerSize = 9
  147. }
  148. c := &Conn{
  149. conn: conn,
  150. r: bufio.NewReader(conn),
  151. cfg: cfg,
  152. calls: make(map[int]*callReq),
  153. timeout: cfg.Timeout,
  154. version: uint8(cfg.ProtoVersion),
  155. addr: conn.RemoteAddr().String(),
  156. errorHandler: errorHandler,
  157. compressor: cfg.Compressor,
  158. auth: cfg.Authenticator,
  159. headerBuf: make([]byte, headerSize),
  160. quit: make(chan struct{}),
  161. session: session,
  162. streams: streams.New(cfg.ProtoVersion),
  163. }
  164. if cfg.Keepalive > 0 {
  165. c.setKeepalive(cfg.Keepalive)
  166. }
  167. go c.serve()
  168. if err := c.startup(); err != nil {
  169. conn.Close()
  170. return nil, err
  171. }
  172. c.started = true
  173. return c, nil
  174. }
  175. func (c *Conn) Write(p []byte) (int, error) {
  176. if c.timeout > 0 {
  177. c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
  178. }
  179. return c.conn.Write(p)
  180. }
  181. func (c *Conn) Read(p []byte) (n int, err error) {
  182. const maxAttempts = 5
  183. for i := 0; i < maxAttempts; i++ {
  184. var nn int
  185. if c.timeout > 0 {
  186. c.conn.SetReadDeadline(time.Now().Add(c.timeout))
  187. }
  188. nn, err = io.ReadFull(c.r, p[n:])
  189. n += nn
  190. if err == nil {
  191. break
  192. }
  193. if verr, ok := err.(net.Error); !ok || !verr.Temporary() {
  194. break
  195. }
  196. }
  197. return
  198. }
  199. func (c *Conn) startup() error {
  200. m := map[string]string{
  201. "CQL_VERSION": c.cfg.CQLVersion,
  202. }
  203. if c.compressor != nil {
  204. m["COMPRESSION"] = c.compressor.Name()
  205. }
  206. framer, err := c.exec(&writeStartupFrame{opts: m}, nil)
  207. if err != nil {
  208. return err
  209. }
  210. frame, err := framer.parseFrame()
  211. if err != nil {
  212. return err
  213. }
  214. switch v := frame.(type) {
  215. case error:
  216. return v
  217. case *readyFrame:
  218. return nil
  219. case *authenticateFrame:
  220. return c.authenticateHandshake(v)
  221. default:
  222. return NewErrProtocol("Unknown type of response to startup frame: %s", v)
  223. }
  224. }
  225. func (c *Conn) authenticateHandshake(authFrame *authenticateFrame) error {
  226. if c.auth == nil {
  227. return fmt.Errorf("authentication required (using %q)", authFrame.class)
  228. }
  229. resp, challenger, err := c.auth.Challenge([]byte(authFrame.class))
  230. if err != nil {
  231. return err
  232. }
  233. req := &writeAuthResponseFrame{data: resp}
  234. for {
  235. framer, err := c.exec(req, nil)
  236. if err != nil {
  237. return err
  238. }
  239. frame, err := framer.parseFrame()
  240. if err != nil {
  241. return err
  242. }
  243. switch v := frame.(type) {
  244. case error:
  245. return v
  246. case *authSuccessFrame:
  247. if challenger != nil {
  248. return challenger.Success(v.data)
  249. }
  250. return nil
  251. case *authChallengeFrame:
  252. resp, challenger, err = challenger.Challenge(v.data)
  253. if err != nil {
  254. return err
  255. }
  256. req = &writeAuthResponseFrame{
  257. data: resp,
  258. }
  259. default:
  260. return fmt.Errorf("unknown frame response during authentication: %v", v)
  261. }
  262. framerPool.Put(framer)
  263. }
  264. }
  265. func (c *Conn) closeWithError(err error) {
  266. if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
  267. return
  268. }
  269. if err != nil {
  270. // we should attempt to deliver the error back to the caller if it
  271. // exists
  272. c.mu.RLock()
  273. for _, req := range c.calls {
  274. // we need to send the error to all waiting queries, put the state
  275. // of this conn into not active so that it can not execute any queries.
  276. if err != nil {
  277. select {
  278. case req.resp <- err:
  279. default:
  280. }
  281. }
  282. }
  283. c.mu.RUnlock()
  284. }
  285. // if error was nil then unblock the quit channel
  286. close(c.quit)
  287. c.conn.Close()
  288. if c.started && err != nil {
  289. c.errorHandler.HandleError(c, err, true)
  290. }
  291. }
  292. func (c *Conn) Close() {
  293. c.closeWithError(nil)
  294. }
  295. // Serve starts the stream multiplexer for this connection, which is required
  296. // to execute any queries. This method runs as long as the connection is
  297. // open and is therefore usually called in a separate goroutine.
  298. func (c *Conn) serve() {
  299. var (
  300. err error
  301. )
  302. for {
  303. err = c.recv()
  304. if err != nil {
  305. break
  306. }
  307. }
  308. c.closeWithError(err)
  309. }
  310. func (c *Conn) discardFrame(head frameHeader) error {
  311. _, err := io.CopyN(ioutil.Discard, c, int64(head.length))
  312. if err != nil {
  313. return err
  314. }
  315. return nil
  316. }
  317. func (c *Conn) recv() error {
  318. // not safe for concurrent reads
  319. // read a full header, ignore timeouts, as this is being ran in a loop
  320. // TODO: TCP level deadlines? or just query level deadlines?
  321. if c.timeout > 0 {
  322. c.conn.SetReadDeadline(time.Time{})
  323. }
  324. // were just reading headers over and over and copy bodies
  325. head, err := readHeader(c.r, c.headerBuf)
  326. if err != nil {
  327. return err
  328. }
  329. if head.stream > c.streams.NumStreams {
  330. return fmt.Errorf("gocql: frame header stream is beyond call exepected bounds: %d", head.stream)
  331. } else if head.stream == -1 {
  332. // TODO: handle cassandra event frames, we shouldnt get any currently
  333. framer := newFramer(c, c, c.compressor, c.version)
  334. if err := framer.readFrame(&head); err != nil {
  335. return err
  336. }
  337. go c.session.handleEvent(framer)
  338. return nil
  339. } else if head.stream <= 0 {
  340. // reserved stream that we dont use, probably due to a protocol error
  341. // or a bug in Cassandra, this should be an error, parse it and return.
  342. framer := newFramer(c, c, c.compressor, c.version)
  343. if err := framer.readFrame(&head); err != nil {
  344. return err
  345. }
  346. defer framerPool.Put(framer)
  347. frame, err := framer.parseFrame()
  348. if err != nil {
  349. return err
  350. }
  351. switch v := frame.(type) {
  352. case error:
  353. return fmt.Errorf("gocql: error on stream %d: %v", head.stream, v)
  354. default:
  355. return fmt.Errorf("gocql: received frame on stream %d: %v", head.stream, frame)
  356. }
  357. }
  358. c.mu.RLock()
  359. call, ok := c.calls[head.stream]
  360. c.mu.RUnlock()
  361. if call == nil || call.framer == nil || !ok {
  362. log.Printf("gocql: received response for stream which has no handler: header=%v\n", head)
  363. return c.discardFrame(head)
  364. }
  365. err = call.framer.readFrame(&head)
  366. if err != nil {
  367. // only net errors should cause the connection to be closed. Though
  368. // cassandra returning corrupt frames will be returned here as well.
  369. if _, ok := err.(net.Error); ok {
  370. return err
  371. }
  372. }
  373. // we either, return a response to the caller, the caller timedout, or the
  374. // connection has closed. Either way we should never block indefinatly here
  375. select {
  376. case call.resp <- err:
  377. case <-call.timeout:
  378. c.releaseStream(head.stream)
  379. case <-c.quit:
  380. }
  381. return nil
  382. }
  383. type callReq struct {
  384. // could use a waitgroup but this allows us to do timeouts on the read/send
  385. resp chan error
  386. framer *framer
  387. timeout chan struct{} // indicates to recv() that a call has timedout
  388. streamID int // current stream in use
  389. }
  390. func (c *Conn) releaseStream(stream int) {
  391. c.mu.Lock()
  392. call := c.calls[stream]
  393. if call != nil && stream != call.streamID {
  394. panic(fmt.Sprintf("attempt to release streamID with ivalid stream: %d -> %+v\n", stream, call))
  395. } else if call == nil {
  396. panic(fmt.Sprintf("releasing a stream not in use: %d", stream))
  397. }
  398. delete(c.calls, stream)
  399. c.mu.Unlock()
  400. streamPool.Put(call)
  401. c.streams.Clear(stream)
  402. }
  403. func (c *Conn) handleTimeout() {
  404. if atomic.AddInt64(&c.timeouts, 1) > TimeoutLimit {
  405. c.closeWithError(ErrTooManyTimeouts)
  406. }
  407. }
  408. var (
  409. streamPool = sync.Pool{
  410. New: func() interface{} {
  411. return &callReq{
  412. resp: make(chan error),
  413. }
  414. },
  415. }
  416. )
  417. func (c *Conn) exec(req frameWriter, tracer Tracer) (*framer, error) {
  418. // TODO: move tracer onto conn
  419. stream, ok := c.streams.GetStream()
  420. if !ok {
  421. fmt.Println(c.streams)
  422. return nil, ErrNoStreams
  423. }
  424. // resp is basically a waiting semaphore protecting the framer
  425. framer := newFramer(c, c, c.compressor, c.version)
  426. c.mu.Lock()
  427. call := c.calls[stream]
  428. if call != nil {
  429. c.mu.Unlock()
  430. return nil, fmt.Errorf("attempting to use stream already in use: %d -> %d", stream, call.streamID)
  431. } else {
  432. call = streamPool.Get().(*callReq)
  433. }
  434. c.calls[stream] = call
  435. c.mu.Unlock()
  436. call.framer = framer
  437. call.timeout = make(chan struct{})
  438. call.streamID = stream
  439. if tracer != nil {
  440. framer.trace()
  441. }
  442. err := req.writeFrame(framer, stream)
  443. if err != nil {
  444. // I think this is the correct thing to do, im not entirely sure. It is not
  445. // ideal as readers might still get some data, but they probably wont.
  446. // Here we need to be careful as the stream is not available and if all
  447. // writes just timeout or fail then the pool might use this connection to
  448. // send a frame on, with all the streams used up and not returned.
  449. c.closeWithError(err)
  450. return nil, err
  451. }
  452. select {
  453. case err := <-call.resp:
  454. if err != nil {
  455. if !c.Closed() {
  456. // if the connection is closed then we cant release the stream,
  457. // this is because the request is still outstanding and we have
  458. // been handed another error from another stream which caused the
  459. // connection to close.
  460. c.releaseStream(stream)
  461. }
  462. return nil, err
  463. }
  464. case <-time.After(c.timeout):
  465. close(call.timeout)
  466. c.handleTimeout()
  467. return nil, ErrTimeoutNoResponse
  468. case <-c.quit:
  469. return nil, ErrConnectionClosed
  470. }
  471. // dont release the stream if detect a timeout as another request can reuse
  472. // that stream and get a response for the old request, which we have no
  473. // easy way of detecting.
  474. //
  475. // Ensure that the stream is not released if there are potentially outstanding
  476. // requests on the stream to prevent nil pointer dereferences in recv().
  477. defer c.releaseStream(stream)
  478. if v := framer.header.version.version(); v != c.version {
  479. return nil, NewErrProtocol("unexpected protocol version in response: got %d expected %d", v, c.version)
  480. }
  481. return framer, nil
  482. }
  483. func (c *Conn) prepareStatement(stmt string, tracer Tracer) (*QueryInfo, error) {
  484. stmtsLRU.Lock()
  485. if stmtsLRU.lru == nil {
  486. initStmtsLRU(defaultMaxPreparedStmts)
  487. }
  488. stmtCacheKey := c.addr + c.currentKeyspace + stmt
  489. if val, ok := stmtsLRU.lru.Get(stmtCacheKey); ok {
  490. stmtsLRU.Unlock()
  491. flight := val.(*inflightPrepare)
  492. flight.wg.Wait()
  493. return &flight.info, flight.err
  494. }
  495. flight := new(inflightPrepare)
  496. flight.wg.Add(1)
  497. stmtsLRU.lru.Add(stmtCacheKey, flight)
  498. stmtsLRU.Unlock()
  499. prep := &writePrepareFrame{
  500. statement: stmt,
  501. }
  502. framer, err := c.exec(prep, tracer)
  503. if err != nil {
  504. flight.err = err
  505. flight.wg.Done()
  506. return nil, err
  507. }
  508. frame, err := framer.parseFrame()
  509. if err != nil {
  510. flight.err = err
  511. flight.wg.Done()
  512. return nil, err
  513. }
  514. // TODO(zariel): tidy this up, simplify handling of frame parsing so its not duplicated
  515. // everytime we need to parse a frame.
  516. if len(framer.traceID) > 0 {
  517. tracer.Trace(framer.traceID)
  518. }
  519. switch x := frame.(type) {
  520. case *resultPreparedFrame:
  521. // defensivly copy as we will recycle the underlying buffer after we
  522. // return.
  523. flight.info.Id = copyBytes(x.preparedID)
  524. // the type info's should _not_ have a reference to the framers read buffer,
  525. // therefore we can just copy them directly.
  526. flight.info.Args = x.reqMeta.columns
  527. flight.info.PKeyColumns = x.reqMeta.pkeyColumns
  528. flight.info.Rval = x.respMeta.columns
  529. case error:
  530. flight.err = x
  531. default:
  532. flight.err = NewErrProtocol("Unknown type in response to prepare frame: %s", x)
  533. }
  534. flight.wg.Done()
  535. if flight.err != nil {
  536. stmtsLRU.Lock()
  537. stmtsLRU.lru.Remove(stmtCacheKey)
  538. stmtsLRU.Unlock()
  539. }
  540. framerPool.Put(framer)
  541. return &flight.info, flight.err
  542. }
  543. func (c *Conn) executeQuery(qry *Query) *Iter {
  544. params := queryParams{
  545. consistency: qry.cons,
  546. }
  547. // frame checks that it is not 0
  548. params.serialConsistency = qry.serialCons
  549. params.defaultTimestamp = qry.defaultTimestamp
  550. if len(qry.pageState) > 0 {
  551. params.pagingState = qry.pageState
  552. }
  553. if qry.pageSize > 0 {
  554. params.pageSize = qry.pageSize
  555. }
  556. var frame frameWriter
  557. if qry.shouldPrepare() {
  558. // Prepare all DML queries. Other queries can not be prepared.
  559. info, err := c.prepareStatement(qry.stmt, qry.trace)
  560. if err != nil {
  561. return &Iter{err: err}
  562. }
  563. var values []interface{}
  564. if qry.binding == nil {
  565. values = qry.values
  566. } else {
  567. values, err = qry.binding(info)
  568. if err != nil {
  569. return &Iter{err: err}
  570. }
  571. }
  572. if len(values) != len(info.Args) {
  573. return &Iter{err: ErrQueryArgLength}
  574. }
  575. params.values = make([]queryValues, len(values))
  576. for i := 0; i < len(values); i++ {
  577. val, err := Marshal(info.Args[i].TypeInfo, values[i])
  578. if err != nil {
  579. return &Iter{err: err}
  580. }
  581. v := &params.values[i]
  582. v.value = val
  583. // TODO: handle query binding names
  584. }
  585. frame = &writeExecuteFrame{
  586. preparedID: info.Id,
  587. params: params,
  588. }
  589. } else {
  590. frame = &writeQueryFrame{
  591. statement: qry.stmt,
  592. params: params,
  593. }
  594. }
  595. framer, err := c.exec(frame, qry.trace)
  596. if err != nil {
  597. return &Iter{err: err}
  598. }
  599. resp, err := framer.parseFrame()
  600. if err != nil {
  601. return &Iter{err: err}
  602. }
  603. if len(framer.traceID) > 0 {
  604. qry.trace.Trace(framer.traceID)
  605. }
  606. switch x := resp.(type) {
  607. case *resultVoidFrame:
  608. return &Iter{framer: framer}
  609. case *resultRowsFrame:
  610. iter := &Iter{
  611. meta: x.meta,
  612. rows: x.rows,
  613. framer: framer,
  614. }
  615. if len(x.meta.pagingState) > 0 && !qry.disableAutoPage {
  616. iter.next = &nextIter{
  617. qry: *qry,
  618. pos: int((1 - qry.prefetch) * float64(len(iter.rows))),
  619. }
  620. iter.next.qry.pageState = x.meta.pagingState
  621. if iter.next.pos < 1 {
  622. iter.next.pos = 1
  623. }
  624. }
  625. return iter
  626. case *resultKeyspaceFrame:
  627. return &Iter{framer: framer}
  628. case *resultSchemaChangeFrame, *schemaChangeKeyspace, *schemaChangeTable, *schemaChangeFunction:
  629. iter := &Iter{framer: framer}
  630. if err := c.awaitSchemaAgreement(); err != nil {
  631. // TODO: should have this behind a flag
  632. log.Println(err)
  633. }
  634. // dont return an error from this, might be a good idea to give a warning
  635. // though. The impact of this returning an error would be that the cluster
  636. // is not consistent with regards to its schema.
  637. return iter
  638. case *RequestErrUnprepared:
  639. stmtsLRU.Lock()
  640. stmtCacheKey := c.addr + c.currentKeyspace + qry.stmt
  641. if _, ok := stmtsLRU.lru.Get(stmtCacheKey); ok {
  642. stmtsLRU.lru.Remove(stmtCacheKey)
  643. stmtsLRU.Unlock()
  644. return c.executeQuery(qry)
  645. }
  646. stmtsLRU.Unlock()
  647. return &Iter{err: x, framer: framer}
  648. case error:
  649. return &Iter{err: x, framer: framer}
  650. default:
  651. return &Iter{
  652. err: NewErrProtocol("Unknown type in response to execute query (%T): %s", x, x),
  653. framer: framer,
  654. }
  655. }
  656. }
  657. func (c *Conn) Pick(qry *Query) *Conn {
  658. if c.Closed() {
  659. return nil
  660. }
  661. return c
  662. }
  663. func (c *Conn) Closed() bool {
  664. return atomic.LoadInt32(&c.closed) == 1
  665. }
  666. func (c *Conn) Address() string {
  667. return c.addr
  668. }
  669. func (c *Conn) AvailableStreams() int {
  670. return c.streams.Available()
  671. }
  672. func (c *Conn) UseKeyspace(keyspace string) error {
  673. q := &writeQueryFrame{statement: `USE "` + keyspace + `"`}
  674. q.params.consistency = Any
  675. framer, err := c.exec(q, nil)
  676. if err != nil {
  677. return err
  678. }
  679. resp, err := framer.parseFrame()
  680. if err != nil {
  681. return err
  682. }
  683. switch x := resp.(type) {
  684. case *resultKeyspaceFrame:
  685. case error:
  686. return x
  687. default:
  688. return NewErrProtocol("unknown frame in response to USE: %v", x)
  689. }
  690. c.currentKeyspace = keyspace
  691. return nil
  692. }
  693. func (c *Conn) executeBatch(batch *Batch) (*Iter, error) {
  694. if c.version == protoVersion1 {
  695. return nil, ErrUnsupported
  696. }
  697. n := len(batch.Entries)
  698. req := &writeBatchFrame{
  699. typ: batch.Type,
  700. statements: make([]batchStatment, n),
  701. consistency: batch.Cons,
  702. serialConsistency: batch.serialCons,
  703. defaultTimestamp: batch.defaultTimestamp,
  704. }
  705. stmts := make(map[string]string)
  706. for i := 0; i < n; i++ {
  707. entry := &batch.Entries[i]
  708. b := &req.statements[i]
  709. if len(entry.Args) > 0 || entry.binding != nil {
  710. info, err := c.prepareStatement(entry.Stmt, nil)
  711. if err != nil {
  712. return nil, err
  713. }
  714. var args []interface{}
  715. if entry.binding == nil {
  716. args = entry.Args
  717. } else {
  718. args, err = entry.binding(info)
  719. if err != nil {
  720. return nil, err
  721. }
  722. }
  723. if len(args) != len(info.Args) {
  724. return nil, ErrQueryArgLength
  725. }
  726. b.preparedID = info.Id
  727. stmts[string(info.Id)] = entry.Stmt
  728. b.values = make([]queryValues, len(info.Args))
  729. for j := 0; j < len(info.Args); j++ {
  730. val, err := Marshal(info.Args[j].TypeInfo, args[j])
  731. if err != nil {
  732. return nil, err
  733. }
  734. b.values[j].value = val
  735. // TODO: add names
  736. }
  737. } else {
  738. b.statement = entry.Stmt
  739. }
  740. }
  741. // TODO: should batch support tracing?
  742. framer, err := c.exec(req, nil)
  743. if err != nil {
  744. return nil, err
  745. }
  746. resp, err := framer.parseFrame()
  747. if err != nil {
  748. return nil, err
  749. }
  750. switch x := resp.(type) {
  751. case *resultVoidFrame:
  752. framerPool.Put(framer)
  753. return nil, nil
  754. case *RequestErrUnprepared:
  755. stmt, found := stmts[string(x.StatementId)]
  756. if found {
  757. stmtsLRU.Lock()
  758. stmtsLRU.lru.Remove(c.addr + c.currentKeyspace + stmt)
  759. stmtsLRU.Unlock()
  760. }
  761. framerPool.Put(framer)
  762. if found {
  763. return c.executeBatch(batch)
  764. } else {
  765. return nil, x
  766. }
  767. case *resultRowsFrame:
  768. iter := &Iter{
  769. meta: x.meta,
  770. rows: x.rows,
  771. framer: framer,
  772. }
  773. return iter, nil
  774. case error:
  775. framerPool.Put(framer)
  776. return nil, x
  777. default:
  778. framerPool.Put(framer)
  779. return nil, NewErrProtocol("Unknown type in response to batch statement: %s", x)
  780. }
  781. }
  782. func (c *Conn) setKeepalive(d time.Duration) error {
  783. if tc, ok := c.conn.(*net.TCPConn); ok {
  784. err := tc.SetKeepAlivePeriod(d)
  785. if err != nil {
  786. return err
  787. }
  788. return tc.SetKeepAlive(true)
  789. }
  790. return nil
  791. }
  792. func (c *Conn) query(statement string, values ...interface{}) (iter *Iter) {
  793. q := c.session.Query(statement, values...).Consistency(One)
  794. return c.executeQuery(q)
  795. }
  796. func (c *Conn) awaitSchemaAgreement() (err error) {
  797. const (
  798. peerSchemas = "SELECT schema_version FROM system.peers"
  799. localSchemas = "SELECT schema_version FROM system.local WHERE key='local'"
  800. )
  801. var versions map[string]struct{}
  802. endDeadline := time.Now().Add(c.session.cfg.MaxWaitSchemaAgreement)
  803. for time.Now().Before(endDeadline) {
  804. iter := c.query(peerSchemas)
  805. versions = make(map[string]struct{})
  806. var schemaVersion string
  807. for iter.Scan(&schemaVersion) {
  808. versions[schemaVersion] = struct{}{}
  809. schemaVersion = ""
  810. }
  811. if err = iter.Close(); err != nil {
  812. goto cont
  813. }
  814. iter = c.query(localSchemas)
  815. for iter.Scan(&schemaVersion) {
  816. versions[schemaVersion] = struct{}{}
  817. schemaVersion = ""
  818. }
  819. if err = iter.Close(); err != nil {
  820. goto cont
  821. }
  822. if len(versions) <= 1 {
  823. return nil
  824. }
  825. cont:
  826. time.Sleep(200 * time.Millisecond)
  827. }
  828. if err != nil {
  829. return
  830. }
  831. schemas := make([]string, 0, len(versions))
  832. for schema := range versions {
  833. schemas = append(schemas, schema)
  834. }
  835. // not exported
  836. return fmt.Errorf("gocql: cluster schema versions not consistent: %+v", schemas)
  837. }
  838. type inflightPrepare struct {
  839. info QueryInfo
  840. err error
  841. wg sync.WaitGroup
  842. }
  843. var (
  844. ErrQueryArgLength = errors.New("gocql: query argument length mismatch")
  845. ErrTimeoutNoResponse = errors.New("gocql: no response received from cassandra within timeout period")
  846. ErrTooManyTimeouts = errors.New("gocql: too many query timeouts on the connection")
  847. ErrConnectionClosed = errors.New("gocql: connection closed waiting for response")
  848. ErrNoStreams = errors.New("gocql: no streams available on connection")
  849. )