conn.go 24 KB

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