conn.go 26 KB

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