conn.go 25 KB

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