conn.go 26 KB

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