conn.go 25 KB

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