conn.go 24 KB

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