conn.go 26 KB

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