conn.go 22 KB

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