session.go 47 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730
  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. "bytes"
  7. "context"
  8. "encoding/binary"
  9. "errors"
  10. "fmt"
  11. "io"
  12. "net"
  13. "strings"
  14. "sync"
  15. "sync/atomic"
  16. "time"
  17. "unicode"
  18. "github.com/gocql/gocql/internal/lru"
  19. )
  20. // Session is the interface used by users to interact with the database.
  21. //
  22. // It's safe for concurrent use by multiple goroutines and a typical usage
  23. // scenario is to have one global session object to interact with the
  24. // whole Cassandra cluster.
  25. //
  26. // This type extends the Node interface by adding a convinient query builder
  27. // and automatically sets a default consistency level on all operations
  28. // that do not have a consistency level set.
  29. type Session struct {
  30. cons Consistency
  31. pageSize int
  32. prefetch float64
  33. routingKeyInfoCache routingKeyInfoLRU
  34. schemaDescriber *schemaDescriber
  35. trace Tracer
  36. queryObserver QueryObserver
  37. batchObserver BatchObserver
  38. hostSource *ringDescriber
  39. stmtsLRU *preparedLRU
  40. connCfg *ConnConfig
  41. executor *queryExecutor
  42. pool *policyConnPool
  43. policy HostSelectionPolicy
  44. ring ring
  45. metadata clusterMetadata
  46. mu sync.RWMutex
  47. control *controlConn
  48. // event handlers
  49. nodeEvents *eventDebouncer
  50. schemaEvents *eventDebouncer
  51. // ring metadata
  52. hosts []HostInfo
  53. useSystemSchema bool
  54. cfg ClusterConfig
  55. quit chan struct{}
  56. closeMu sync.RWMutex
  57. isClosed bool
  58. }
  59. var queryPool = &sync.Pool{
  60. New: func() interface{} {
  61. return new(Query)
  62. },
  63. }
  64. func addrsToHosts(addrs []string, defaultPort int) ([]*HostInfo, error) {
  65. var hosts []*HostInfo
  66. for _, hostport := range addrs {
  67. resolvedHosts, err := hostInfo(hostport, defaultPort)
  68. if err != nil {
  69. // Try other hosts if unable to resolve DNS name
  70. if _, ok := err.(*net.DNSError); ok {
  71. Logger.Printf("gocql: dns error: %v\n", err)
  72. continue
  73. }
  74. return nil, err
  75. }
  76. hosts = append(hosts, resolvedHosts...)
  77. }
  78. if len(hosts) == 0 {
  79. return nil, errors.New("failed to resolve any of the provided hostnames")
  80. }
  81. return hosts, nil
  82. }
  83. // NewSession wraps an existing Node.
  84. func NewSession(cfg ClusterConfig) (*Session, error) {
  85. // Check that hosts in the ClusterConfig is not empty
  86. if len(cfg.Hosts) < 1 {
  87. return nil, ErrNoHosts
  88. }
  89. s := &Session{
  90. cons: cfg.Consistency,
  91. prefetch: 0.25,
  92. cfg: cfg,
  93. pageSize: cfg.PageSize,
  94. stmtsLRU: &preparedLRU{lru: lru.New(cfg.MaxPreparedStmts)},
  95. quit: make(chan struct{}),
  96. }
  97. s.schemaDescriber = newSchemaDescriber(s)
  98. s.nodeEvents = newEventDebouncer("NodeEvents", s.handleNodeEvent)
  99. s.schemaEvents = newEventDebouncer("SchemaEvents", s.handleSchemaEvent)
  100. s.routingKeyInfoCache.lru = lru.New(cfg.MaxRoutingKeyInfo)
  101. s.hostSource = &ringDescriber{session: s}
  102. if cfg.PoolConfig.HostSelectionPolicy == nil {
  103. cfg.PoolConfig.HostSelectionPolicy = RoundRobinHostPolicy()
  104. }
  105. s.pool = cfg.PoolConfig.buildPool(s)
  106. s.policy = cfg.PoolConfig.HostSelectionPolicy
  107. s.policy.Init(s)
  108. s.executor = &queryExecutor{
  109. pool: s.pool,
  110. policy: cfg.PoolConfig.HostSelectionPolicy,
  111. }
  112. s.queryObserver = cfg.QueryObserver
  113. s.batchObserver = cfg.BatchObserver
  114. //Check the TLS Config before trying to connect to anything external
  115. connCfg, err := connConfig(&s.cfg)
  116. if err != nil {
  117. //TODO: Return a typed error
  118. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  119. }
  120. s.connCfg = connCfg
  121. if err := s.init(); err != nil {
  122. s.Close()
  123. if err == ErrNoConnectionsStarted {
  124. //This error used to be generated inside NewSession & returned directly
  125. //Forward it on up to be backwards compatible
  126. return nil, ErrNoConnectionsStarted
  127. } else {
  128. // TODO(zariel): dont wrap this error in fmt.Errorf, return a typed error
  129. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  130. }
  131. }
  132. return s, nil
  133. }
  134. func (s *Session) init() error {
  135. hosts, err := addrsToHosts(s.cfg.Hosts, s.cfg.Port)
  136. if err != nil {
  137. return err
  138. }
  139. s.ring.endpoints = hosts
  140. if !s.cfg.disableControlConn {
  141. s.control = createControlConn(s)
  142. if s.cfg.ProtoVersion == 0 {
  143. proto, err := s.control.discoverProtocol(hosts)
  144. if err != nil {
  145. return fmt.Errorf("unable to discover protocol version: %v", err)
  146. } else if proto == 0 {
  147. return errors.New("unable to discovery protocol version")
  148. }
  149. // TODO(zariel): we really only need this in 1 place
  150. s.cfg.ProtoVersion = proto
  151. s.connCfg.ProtoVersion = proto
  152. }
  153. if err := s.control.connect(hosts); err != nil {
  154. return err
  155. }
  156. if !s.cfg.DisableInitialHostLookup {
  157. var partitioner string
  158. newHosts, partitioner, err := s.hostSource.GetHosts()
  159. if err != nil {
  160. return err
  161. }
  162. s.policy.SetPartitioner(partitioner)
  163. filteredHosts := make([]*HostInfo, 0, len(newHosts))
  164. for _, host := range newHosts {
  165. if !s.cfg.filterHost(host) {
  166. filteredHosts = append(filteredHosts, host)
  167. }
  168. }
  169. hosts = append(hosts, filteredHosts...)
  170. }
  171. }
  172. hostMap := make(map[string]*HostInfo, len(hosts))
  173. for _, host := range hosts {
  174. hostMap[host.ConnectAddress().String()] = host
  175. }
  176. for _, host := range hostMap {
  177. host = s.ring.addOrUpdate(host)
  178. s.addNewNode(host)
  179. }
  180. // TODO(zariel): we probably dont need this any more as we verify that we
  181. // can connect to one of the endpoints supplied by using the control conn.
  182. // See if there are any connections in the pool
  183. if s.cfg.ReconnectInterval > 0 {
  184. go s.reconnectDownedHosts(s.cfg.ReconnectInterval)
  185. }
  186. // If we disable the initial host lookup, we need to still check if the
  187. // cluster is using the newer system schema or not... however, if control
  188. // connection is disable, we really have no choice, so we just make our
  189. // best guess...
  190. if !s.cfg.disableControlConn && s.cfg.DisableInitialHostLookup {
  191. newer, _ := checkSystemSchema(s.control)
  192. s.useSystemSchema = newer
  193. } else {
  194. host := s.ring.rrHost()
  195. s.useSystemSchema = host.Version().Major >= 3
  196. }
  197. if s.pool.Size() == 0 {
  198. return ErrNoConnectionsStarted
  199. }
  200. return nil
  201. }
  202. func (s *Session) reconnectDownedHosts(intv time.Duration) {
  203. reconnectTicker := time.NewTicker(intv)
  204. defer reconnectTicker.Stop()
  205. for {
  206. select {
  207. case <-reconnectTicker.C:
  208. hosts := s.ring.allHosts()
  209. // Print session.ring for debug.
  210. if gocqlDebug {
  211. buf := bytes.NewBufferString("Session.ring:")
  212. for _, h := range hosts {
  213. buf.WriteString("[" + h.ConnectAddress().String() + ":" + h.State().String() + "]")
  214. }
  215. Logger.Println(buf.String())
  216. }
  217. for _, h := range hosts {
  218. if h.IsUp() {
  219. continue
  220. }
  221. s.handleNodeUp(h.ConnectAddress(), h.Port(), true)
  222. }
  223. case <-s.quit:
  224. return
  225. }
  226. }
  227. }
  228. // SetConsistency sets the default consistency level for this session. This
  229. // setting can also be changed on a per-query basis and the default value
  230. // is Quorum.
  231. func (s *Session) SetConsistency(cons Consistency) {
  232. s.mu.Lock()
  233. s.cons = cons
  234. s.mu.Unlock()
  235. }
  236. // SetPageSize sets the default page size for this session. A value <= 0 will
  237. // disable paging. This setting can also be changed on a per-query basis.
  238. func (s *Session) SetPageSize(n int) {
  239. s.mu.Lock()
  240. s.pageSize = n
  241. s.mu.Unlock()
  242. }
  243. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  244. // there are only p*pageSize rows remaining, the next page will be requested
  245. // automatically. This value can also be changed on a per-query basis and
  246. // the default value is 0.25.
  247. func (s *Session) SetPrefetch(p float64) {
  248. s.mu.Lock()
  249. s.prefetch = p
  250. s.mu.Unlock()
  251. }
  252. // SetTrace sets the default tracer for this session. This setting can also
  253. // be changed on a per-query basis.
  254. func (s *Session) SetTrace(trace Tracer) {
  255. s.mu.Lock()
  256. s.trace = trace
  257. s.mu.Unlock()
  258. }
  259. // Query generates a new query object for interacting with the database.
  260. // Further details of the query may be tweaked using the resulting query
  261. // value before the query is executed. Query is automatically prepared
  262. // if it has not previously been executed.
  263. func (s *Session) Query(stmt string, values ...interface{}) *Query {
  264. s.mu.RLock()
  265. qry := queryPool.Get().(*Query)
  266. qry.stmt = stmt
  267. qry.values = values
  268. qry.cons = s.cons
  269. qry.session = s
  270. qry.pageSize = s.pageSize
  271. qry.trace = s.trace
  272. qry.observer = s.queryObserver
  273. qry.prefetch = s.prefetch
  274. qry.rt = s.cfg.RetryPolicy
  275. qry.serialCons = s.cfg.SerialConsistency
  276. qry.defaultTimestamp = s.cfg.DefaultTimestamp
  277. s.mu.RUnlock()
  278. return qry
  279. }
  280. type QueryInfo struct {
  281. Id []byte
  282. Args []ColumnInfo
  283. Rval []ColumnInfo
  284. PKeyColumns []int
  285. }
  286. // Bind generates a new query object based on the query statement passed in.
  287. // The query is automatically prepared if it has not previously been executed.
  288. // The binding callback allows the application to define which query argument
  289. // values will be marshalled as part of the query execution.
  290. // During execution, the meta data of the prepared query will be routed to the
  291. // binding callback, which is responsible for producing the query argument values.
  292. func (s *Session) Bind(stmt string, b func(q *QueryInfo) ([]interface{}, error)) *Query {
  293. s.mu.RLock()
  294. qry := &Query{stmt: stmt, binding: b, cons: s.cons,
  295. session: s, pageSize: s.pageSize, trace: s.trace, observer: s.queryObserver,
  296. prefetch: s.prefetch, rt: s.cfg.RetryPolicy}
  297. s.mu.RUnlock()
  298. return qry
  299. }
  300. // Close closes all connections. The session is unusable after this
  301. // operation.
  302. func (s *Session) Close() {
  303. s.closeMu.Lock()
  304. defer s.closeMu.Unlock()
  305. if s.isClosed {
  306. return
  307. }
  308. s.isClosed = true
  309. if s.pool != nil {
  310. s.pool.Close()
  311. }
  312. if s.control != nil {
  313. s.control.close()
  314. }
  315. if s.nodeEvents != nil {
  316. s.nodeEvents.stop()
  317. }
  318. if s.schemaEvents != nil {
  319. s.schemaEvents.stop()
  320. }
  321. if s.quit != nil {
  322. close(s.quit)
  323. }
  324. }
  325. func (s *Session) Closed() bool {
  326. s.closeMu.RLock()
  327. closed := s.isClosed
  328. s.closeMu.RUnlock()
  329. return closed
  330. }
  331. func (s *Session) executeQuery(qry *Query) (it *Iter) {
  332. // fail fast
  333. if s.Closed() {
  334. return &Iter{err: ErrSessionClosed}
  335. }
  336. iter, err := s.executor.executeQuery(qry)
  337. if err != nil {
  338. return &Iter{err: err}
  339. }
  340. if iter == nil {
  341. panic("nil iter")
  342. }
  343. return iter
  344. }
  345. func (s *Session) removeHost(h *HostInfo) {
  346. s.policy.RemoveHost(h)
  347. s.pool.removeHost(h.ConnectAddress())
  348. s.ring.removeHost(h.ConnectAddress())
  349. }
  350. // KeyspaceMetadata returns the schema metadata for the keyspace specified. Returns an error if the keyspace does not exist.
  351. func (s *Session) KeyspaceMetadata(keyspace string) (*KeyspaceMetadata, error) {
  352. // fail fast
  353. if s.Closed() {
  354. return nil, ErrSessionClosed
  355. } else if keyspace == "" {
  356. return nil, ErrNoKeyspace
  357. }
  358. return s.schemaDescriber.getSchema(keyspace)
  359. }
  360. func (s *Session) getConn() *Conn {
  361. hosts := s.ring.allHosts()
  362. for _, host := range hosts {
  363. if !host.IsUp() {
  364. continue
  365. }
  366. pool, ok := s.pool.getPool(host)
  367. if !ok {
  368. continue
  369. } else if conn := pool.Pick(); conn != nil {
  370. return conn
  371. }
  372. }
  373. return nil
  374. }
  375. // returns routing key indexes and type info
  376. func (s *Session) routingKeyInfo(ctx context.Context, stmt string) (*routingKeyInfo, error) {
  377. s.routingKeyInfoCache.mu.Lock()
  378. entry, cached := s.routingKeyInfoCache.lru.Get(stmt)
  379. if cached {
  380. // done accessing the cache
  381. s.routingKeyInfoCache.mu.Unlock()
  382. // the entry is an inflight struct similar to that used by
  383. // Conn to prepare statements
  384. inflight := entry.(*inflightCachedEntry)
  385. // wait for any inflight work
  386. inflight.wg.Wait()
  387. if inflight.err != nil {
  388. return nil, inflight.err
  389. }
  390. key, _ := inflight.value.(*routingKeyInfo)
  391. return key, nil
  392. }
  393. // create a new inflight entry while the data is created
  394. inflight := new(inflightCachedEntry)
  395. inflight.wg.Add(1)
  396. defer inflight.wg.Done()
  397. s.routingKeyInfoCache.lru.Add(stmt, inflight)
  398. s.routingKeyInfoCache.mu.Unlock()
  399. var (
  400. info *preparedStatment
  401. partitionKey []*ColumnMetadata
  402. )
  403. conn := s.getConn()
  404. if conn == nil {
  405. // TODO: better error?
  406. inflight.err = errors.New("gocql: unable to fetch prepared info: no connection available")
  407. return nil, inflight.err
  408. }
  409. // get the query info for the statement
  410. info, inflight.err = conn.prepareStatement(ctx, stmt, nil)
  411. if inflight.err != nil {
  412. // don't cache this error
  413. s.routingKeyInfoCache.Remove(stmt)
  414. return nil, inflight.err
  415. }
  416. // TODO: it would be nice to mark hosts here but as we are not using the policies
  417. // to fetch hosts we cant
  418. if info.request.colCount == 0 {
  419. // no arguments, no routing key, and no error
  420. return nil, nil
  421. }
  422. if len(info.request.pkeyColumns) > 0 {
  423. // proto v4 dont need to calculate primary key columns
  424. types := make([]TypeInfo, len(info.request.pkeyColumns))
  425. for i, col := range info.request.pkeyColumns {
  426. types[i] = info.request.columns[col].TypeInfo
  427. }
  428. routingKeyInfo := &routingKeyInfo{
  429. indexes: info.request.pkeyColumns,
  430. types: types,
  431. }
  432. inflight.value = routingKeyInfo
  433. return routingKeyInfo, nil
  434. }
  435. // get the table metadata
  436. table := info.request.columns[0].Table
  437. var keyspaceMetadata *KeyspaceMetadata
  438. keyspaceMetadata, inflight.err = s.KeyspaceMetadata(info.request.columns[0].Keyspace)
  439. if inflight.err != nil {
  440. // don't cache this error
  441. s.routingKeyInfoCache.Remove(stmt)
  442. return nil, inflight.err
  443. }
  444. tableMetadata, found := keyspaceMetadata.Tables[table]
  445. if !found {
  446. // unlikely that the statement could be prepared and the metadata for
  447. // the table couldn't be found, but this may indicate either a bug
  448. // in the metadata code, or that the table was just dropped.
  449. inflight.err = ErrNoMetadata
  450. // don't cache this error
  451. s.routingKeyInfoCache.Remove(stmt)
  452. return nil, inflight.err
  453. }
  454. partitionKey = tableMetadata.PartitionKey
  455. size := len(partitionKey)
  456. routingKeyInfo := &routingKeyInfo{
  457. indexes: make([]int, size),
  458. types: make([]TypeInfo, size),
  459. }
  460. for keyIndex, keyColumn := range partitionKey {
  461. // set an indicator for checking if the mapping is missing
  462. routingKeyInfo.indexes[keyIndex] = -1
  463. // find the column in the query info
  464. for argIndex, boundColumn := range info.request.columns {
  465. if keyColumn.Name == boundColumn.Name {
  466. // there may be many such bound columns, pick the first
  467. routingKeyInfo.indexes[keyIndex] = argIndex
  468. routingKeyInfo.types[keyIndex] = boundColumn.TypeInfo
  469. break
  470. }
  471. }
  472. if routingKeyInfo.indexes[keyIndex] == -1 {
  473. // missing a routing key column mapping
  474. // no routing key, and no error
  475. return nil, nil
  476. }
  477. }
  478. // cache this result
  479. inflight.value = routingKeyInfo
  480. return routingKeyInfo, nil
  481. }
  482. func (b *Batch) execute(conn *Conn) *Iter {
  483. return conn.executeBatch(b)
  484. }
  485. func (s *Session) executeBatch(batch *Batch) *Iter {
  486. // fail fast
  487. if s.Closed() {
  488. return &Iter{err: ErrSessionClosed}
  489. }
  490. // Prevent the execution of the batch if greater than the limit
  491. // Currently batches have a limit of 65536 queries.
  492. // https://datastax-oss.atlassian.net/browse/JAVA-229
  493. if batch.Size() > BatchSizeMaximum {
  494. return &Iter{err: ErrTooManyStmts}
  495. }
  496. iter, err := s.executor.executeQuery(batch)
  497. if err != nil {
  498. return &Iter{err: err}
  499. }
  500. return iter
  501. }
  502. // ExecuteBatch executes a batch operation and returns nil if successful
  503. // otherwise an error is returned describing the failure.
  504. func (s *Session) ExecuteBatch(batch *Batch) error {
  505. iter := s.executeBatch(batch)
  506. return iter.Close()
  507. }
  508. // ExecuteBatchCAS executes a batch operation and returns true if successful and
  509. // an iterator (to scan aditional rows if more than one conditional statement)
  510. // was sent.
  511. // Further scans on the interator must also remember to include
  512. // the applied boolean as the first argument to *Iter.Scan
  513. func (s *Session) ExecuteBatchCAS(batch *Batch, dest ...interface{}) (applied bool, iter *Iter, err error) {
  514. iter = s.executeBatch(batch)
  515. if err := iter.checkErrAndNotFound(); err != nil {
  516. iter.Close()
  517. return false, nil, err
  518. }
  519. if len(iter.Columns()) > 1 {
  520. dest = append([]interface{}{&applied}, dest...)
  521. iter.Scan(dest...)
  522. } else {
  523. iter.Scan(&applied)
  524. }
  525. return applied, iter, nil
  526. }
  527. // MapExecuteBatchCAS executes a batch operation much like ExecuteBatchCAS,
  528. // however it accepts a map rather than a list of arguments for the initial
  529. // scan.
  530. func (s *Session) MapExecuteBatchCAS(batch *Batch, dest map[string]interface{}) (applied bool, iter *Iter, err error) {
  531. iter = s.executeBatch(batch)
  532. if err := iter.checkErrAndNotFound(); err != nil {
  533. iter.Close()
  534. return false, nil, err
  535. }
  536. iter.MapScan(dest)
  537. applied = dest["[applied]"].(bool)
  538. delete(dest, "[applied]")
  539. // we usually close here, but instead of closing, just returin an error
  540. // if MapScan failed. Although Close just returns err, using Close
  541. // here might be confusing as we are not actually closing the iter
  542. return applied, iter, iter.err
  543. }
  544. func (s *Session) connect(host *HostInfo, errorHandler ConnErrorHandler) (*Conn, error) {
  545. return s.dial(host.ConnectAddress(), host.Port(), s.connCfg, errorHandler)
  546. }
  547. // Query represents a CQL statement that can be executed.
  548. type Query struct {
  549. stmt string
  550. values []interface{}
  551. cons Consistency
  552. pageSize int
  553. routingKey []byte
  554. routingKeyBuffer []byte
  555. pageState []byte
  556. prefetch float64
  557. trace Tracer
  558. observer QueryObserver
  559. session *Session
  560. rt RetryPolicy
  561. binding func(q *QueryInfo) ([]interface{}, error)
  562. attempts int
  563. totalLatency int64
  564. serialCons SerialConsistency
  565. defaultTimestamp bool
  566. defaultTimestampValue int64
  567. disableSkipMetadata bool
  568. context context.Context
  569. disableAutoPage bool
  570. }
  571. // String implements the stringer interface.
  572. func (q Query) String() string {
  573. return fmt.Sprintf("[query statement=%q values=%+v consistency=%s]", q.stmt, q.values, q.cons)
  574. }
  575. //Attempts returns the number of times the query was executed.
  576. func (q *Query) Attempts() int {
  577. return q.attempts
  578. }
  579. //Latency returns the average amount of nanoseconds per attempt of the query.
  580. func (q *Query) Latency() int64 {
  581. if q.attempts > 0 {
  582. return q.totalLatency / int64(q.attempts)
  583. }
  584. return 0
  585. }
  586. // Consistency sets the consistency level for this query. If no consistency
  587. // level have been set, the default consistency level of the cluster
  588. // is used.
  589. func (q *Query) Consistency(c Consistency) *Query {
  590. q.cons = c
  591. return q
  592. }
  593. // GetConsistency returns the currently configured consistency level for
  594. // the query.
  595. func (q *Query) GetConsistency() Consistency {
  596. return q.cons
  597. }
  598. // Trace enables tracing of this query. Look at the documentation of the
  599. // Tracer interface to learn more about tracing.
  600. func (q *Query) Trace(trace Tracer) *Query {
  601. q.trace = trace
  602. return q
  603. }
  604. // Observer enables query-level observer on this query.
  605. // The provided observer will be called every time this query is executed.
  606. func (q *Query) Observer(observer QueryObserver) *Query {
  607. q.observer = observer
  608. return q
  609. }
  610. // PageSize will tell the iterator to fetch the result in pages of size n.
  611. // This is useful for iterating over large result sets, but setting the
  612. // page size too low might decrease the performance. This feature is only
  613. // available in Cassandra 2 and onwards.
  614. func (q *Query) PageSize(n int) *Query {
  615. q.pageSize = n
  616. return q
  617. }
  618. // DefaultTimestamp will enable the with default timestamp flag on the query.
  619. // If enable, this will replace the server side assigned
  620. // timestamp as default timestamp. Note that a timestamp in the query itself
  621. // will still override this timestamp. This is entirely optional.
  622. //
  623. // Only available on protocol >= 3
  624. func (q *Query) DefaultTimestamp(enable bool) *Query {
  625. q.defaultTimestamp = enable
  626. return q
  627. }
  628. // WithTimestamp will enable the with default timestamp flag on the query
  629. // like DefaultTimestamp does. But also allows to define value for timestamp.
  630. // It works the same way as USING TIMESTAMP in the query itself, but
  631. // should not break prepared query optimization
  632. //
  633. // Only available on protocol >= 3
  634. func (q *Query) WithTimestamp(timestamp int64) *Query {
  635. q.DefaultTimestamp(true)
  636. q.defaultTimestampValue = timestamp
  637. return q
  638. }
  639. // RoutingKey sets the routing key to use when a token aware connection
  640. // pool is used to optimize the routing of this query.
  641. func (q *Query) RoutingKey(routingKey []byte) *Query {
  642. q.routingKey = routingKey
  643. return q
  644. }
  645. // WithContext will set the context to use during a query, it will be used to
  646. // timeout when waiting for responses from Cassandra.
  647. func (q *Query) WithContext(ctx context.Context) *Query {
  648. q.context = ctx
  649. return q
  650. }
  651. func (q *Query) execute(conn *Conn) *Iter {
  652. return conn.executeQuery(q)
  653. }
  654. func (q *Query) attempt(keyspace string, end, start time.Time, iter *Iter) {
  655. q.attempts++
  656. q.totalLatency += end.Sub(start).Nanoseconds()
  657. // TODO: track latencies per host and things as well instead of just total
  658. if q.observer != nil {
  659. q.observer.ObserveQuery(q.context, ObservedQuery{
  660. Keyspace: keyspace,
  661. Statement: q.stmt,
  662. Start: start,
  663. End: end,
  664. Rows: iter.numRows,
  665. Err: iter.err,
  666. })
  667. }
  668. }
  669. func (q *Query) retryPolicy() RetryPolicy {
  670. return q.rt
  671. }
  672. // Keyspace returns the keyspace the query will be executed against.
  673. func (q *Query) Keyspace() string {
  674. if q.session == nil {
  675. return ""
  676. }
  677. // TODO(chbannis): this should be parsed from the query or we should let
  678. // this be set by users.
  679. return q.session.cfg.Keyspace
  680. }
  681. // GetRoutingKey gets the routing key to use for routing this query. If
  682. // a routing key has not been explicitly set, then the routing key will
  683. // be constructed if possible using the keyspace's schema and the query
  684. // info for this query statement. If the routing key cannot be determined
  685. // then nil will be returned with no error. On any error condition,
  686. // an error description will be returned.
  687. func (q *Query) GetRoutingKey() ([]byte, error) {
  688. if q.routingKey != nil {
  689. return q.routingKey, nil
  690. } else if q.binding != nil && len(q.values) == 0 {
  691. // If this query was created using session.Bind we wont have the query
  692. // values yet, so we have to pass down to the next policy.
  693. // TODO: Remove this and handle this case
  694. return nil, nil
  695. }
  696. // try to determine the routing key
  697. routingKeyInfo, err := q.session.routingKeyInfo(q.context, q.stmt)
  698. if err != nil {
  699. return nil, err
  700. }
  701. if routingKeyInfo == nil {
  702. return nil, nil
  703. }
  704. if len(routingKeyInfo.indexes) == 1 {
  705. // single column routing key
  706. routingKey, err := Marshal(
  707. routingKeyInfo.types[0],
  708. q.values[routingKeyInfo.indexes[0]],
  709. )
  710. if err != nil {
  711. return nil, err
  712. }
  713. return routingKey, nil
  714. }
  715. // We allocate that buffer only once, so that further re-bind/exec of the
  716. // same query don't allocate more memory.
  717. if q.routingKeyBuffer == nil {
  718. q.routingKeyBuffer = make([]byte, 0, 256)
  719. }
  720. // composite routing key
  721. buf := bytes.NewBuffer(q.routingKeyBuffer)
  722. for i := range routingKeyInfo.indexes {
  723. encoded, err := Marshal(
  724. routingKeyInfo.types[i],
  725. q.values[routingKeyInfo.indexes[i]],
  726. )
  727. if err != nil {
  728. return nil, err
  729. }
  730. lenBuf := []byte{0x00, 0x00}
  731. binary.BigEndian.PutUint16(lenBuf, uint16(len(encoded)))
  732. buf.Write(lenBuf)
  733. buf.Write(encoded)
  734. buf.WriteByte(0x00)
  735. }
  736. routingKey := buf.Bytes()
  737. return routingKey, nil
  738. }
  739. func (q *Query) shouldPrepare() bool {
  740. stmt := strings.TrimLeftFunc(strings.TrimRightFunc(q.stmt, func(r rune) bool {
  741. return unicode.IsSpace(r) || r == ';'
  742. }), unicode.IsSpace)
  743. var stmtType string
  744. if n := strings.IndexFunc(stmt, unicode.IsSpace); n >= 0 {
  745. stmtType = strings.ToLower(stmt[:n])
  746. }
  747. if stmtType == "begin" {
  748. if n := strings.LastIndexFunc(stmt, unicode.IsSpace); n >= 0 {
  749. stmtType = strings.ToLower(stmt[n+1:])
  750. }
  751. }
  752. switch stmtType {
  753. case "select", "insert", "update", "delete", "batch":
  754. return true
  755. }
  756. return false
  757. }
  758. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  759. // there are only p*pageSize rows remaining, the next page will be requested
  760. // automatically.
  761. func (q *Query) Prefetch(p float64) *Query {
  762. q.prefetch = p
  763. return q
  764. }
  765. // RetryPolicy sets the policy to use when retrying the query.
  766. func (q *Query) RetryPolicy(r RetryPolicy) *Query {
  767. q.rt = r
  768. return q
  769. }
  770. // Bind sets query arguments of query. This can also be used to rebind new query arguments
  771. // to an existing query instance.
  772. func (q *Query) Bind(v ...interface{}) *Query {
  773. q.values = v
  774. return q
  775. }
  776. // SerialConsistency sets the consistency level for the
  777. // serial phase of conditional updates. That consistency can only be
  778. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  779. // SERIAL. This option will be ignored for anything else that a
  780. // conditional update/insert.
  781. func (q *Query) SerialConsistency(cons SerialConsistency) *Query {
  782. q.serialCons = cons
  783. return q
  784. }
  785. // PageState sets the paging state for the query to resume paging from a specific
  786. // point in time. Setting this will disable to query paging for this query, and
  787. // must be used for all subsequent pages.
  788. func (q *Query) PageState(state []byte) *Query {
  789. q.pageState = state
  790. q.disableAutoPage = true
  791. return q
  792. }
  793. // NoSkipMetadata will override the internal result metadata cache so that the driver does not
  794. // send skip_metadata for queries, this means that the result will always contain
  795. // the metadata to parse the rows and will not reuse the metadata from the prepared
  796. // staement. This should only be used to work around cassandra bugs, such as when using
  797. // CAS operations which do not end in Cas.
  798. //
  799. // See https://issues.apache.org/jira/browse/CASSANDRA-11099
  800. // https://github.com/gocql/gocql/issues/612
  801. func (q *Query) NoSkipMetadata() *Query {
  802. q.disableSkipMetadata = true
  803. return q
  804. }
  805. // Exec executes the query without returning any rows.
  806. func (q *Query) Exec() error {
  807. return q.Iter().Close()
  808. }
  809. func isUseStatement(stmt string) bool {
  810. if len(stmt) < 3 {
  811. return false
  812. }
  813. return strings.EqualFold(stmt[0:3], "use")
  814. }
  815. // Iter executes the query and returns an iterator capable of iterating
  816. // over all results.
  817. func (q *Query) Iter() *Iter {
  818. if isUseStatement(q.stmt) {
  819. return &Iter{err: ErrUseStmt}
  820. }
  821. return q.session.executeQuery(q)
  822. }
  823. // MapScan executes the query, copies the columns of the first selected
  824. // row into the map pointed at by m and discards the rest. If no rows
  825. // were selected, ErrNotFound is returned.
  826. func (q *Query) MapScan(m map[string]interface{}) error {
  827. iter := q.Iter()
  828. if err := iter.checkErrAndNotFound(); err != nil {
  829. return err
  830. }
  831. iter.MapScan(m)
  832. return iter.Close()
  833. }
  834. // Scan executes the query, copies the columns of the first selected
  835. // row into the values pointed at by dest and discards the rest. If no rows
  836. // were selected, ErrNotFound is returned.
  837. func (q *Query) Scan(dest ...interface{}) error {
  838. iter := q.Iter()
  839. if err := iter.checkErrAndNotFound(); err != nil {
  840. return err
  841. }
  842. iter.Scan(dest...)
  843. return iter.Close()
  844. }
  845. // ScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  846. // statement containing an IF clause). If the transaction fails because
  847. // the existing values did not match, the previous values will be stored
  848. // in dest.
  849. func (q *Query) ScanCAS(dest ...interface{}) (applied bool, err error) {
  850. q.disableSkipMetadata = true
  851. iter := q.Iter()
  852. if err := iter.checkErrAndNotFound(); err != nil {
  853. return false, err
  854. }
  855. if len(iter.Columns()) > 1 {
  856. dest = append([]interface{}{&applied}, dest...)
  857. iter.Scan(dest...)
  858. } else {
  859. iter.Scan(&applied)
  860. }
  861. return applied, iter.Close()
  862. }
  863. // MapScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  864. // statement containing an IF clause). If the transaction fails because
  865. // the existing values did not match, the previous values will be stored
  866. // in dest map.
  867. //
  868. // As for INSERT .. IF NOT EXISTS, previous values will be returned as if
  869. // SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
  870. // column mismatching. MapScanCAS is added to capture them safely.
  871. func (q *Query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
  872. q.disableSkipMetadata = true
  873. iter := q.Iter()
  874. if err := iter.checkErrAndNotFound(); err != nil {
  875. return false, err
  876. }
  877. iter.MapScan(dest)
  878. applied = dest["[applied]"].(bool)
  879. delete(dest, "[applied]")
  880. return applied, iter.Close()
  881. }
  882. // Release releases a query back into a pool of queries. Released Queries
  883. // cannot be reused.
  884. //
  885. // Example:
  886. // qry := session.Query("SELECT * FROM my_table")
  887. // qry.Exec()
  888. // qry.Release()
  889. func (q *Query) Release() {
  890. q.reset()
  891. queryPool.Put(q)
  892. }
  893. // reset zeroes out all fields of a query so that it can be safely pooled.
  894. func (q *Query) reset() {
  895. q.stmt = ""
  896. q.values = nil
  897. q.cons = 0
  898. q.pageSize = 0
  899. q.routingKey = nil
  900. q.routingKeyBuffer = nil
  901. q.pageState = nil
  902. q.prefetch = 0
  903. q.trace = nil
  904. q.session = nil
  905. q.rt = nil
  906. q.binding = nil
  907. q.attempts = 0
  908. q.totalLatency = 0
  909. q.serialCons = 0
  910. q.defaultTimestamp = false
  911. q.disableSkipMetadata = false
  912. q.disableAutoPage = false
  913. q.context = nil
  914. }
  915. // Iter represents an iterator that can be used to iterate over all rows that
  916. // were returned by a query. The iterator might send additional queries to the
  917. // database during the iteration if paging was enabled.
  918. type Iter struct {
  919. err error
  920. pos int
  921. meta resultMetadata
  922. numRows int
  923. next *nextIter
  924. host *HostInfo
  925. framer *framer
  926. closed int32
  927. }
  928. // Host returns the host which the query was sent to.
  929. func (iter *Iter) Host() *HostInfo {
  930. return iter.host
  931. }
  932. // Columns returns the name and type of the selected columns.
  933. func (iter *Iter) Columns() []ColumnInfo {
  934. return iter.meta.columns
  935. }
  936. type Scanner interface {
  937. // Next advances the row pointer to point at the next row, the row is valid until
  938. // the next call of Next. It returns true if there is a row which is available to be
  939. // scanned into with Scan.
  940. // Next must be called before every call to Scan.
  941. Next() bool
  942. // Scan copies the current row's columns into dest. If the length of dest does not equal
  943. // the number of columns returned in the row an error is returned. If an error is encountered
  944. // when unmarshalling a column into the value in dest an error is returned and the row is invalidated
  945. // until the next call to Next.
  946. // Next must be called before calling Scan, if it is not an error is returned.
  947. Scan(...interface{}) error
  948. // Err returns the if there was one during iteration that resulted in iteration being unable to complete.
  949. // Err will also release resources held by the iterator, the Scanner should not used after being called.
  950. Err() error
  951. }
  952. type iterScanner struct {
  953. iter *Iter
  954. cols [][]byte
  955. }
  956. func (is *iterScanner) Next() bool {
  957. iter := is.iter
  958. if iter.err != nil {
  959. return false
  960. }
  961. if iter.pos >= iter.numRows {
  962. if iter.next != nil {
  963. is.iter = iter.next.fetch()
  964. return is.Next()
  965. }
  966. return false
  967. }
  968. cols := make([][]byte, len(iter.meta.columns))
  969. for i := 0; i < len(cols); i++ {
  970. col, err := iter.readColumn()
  971. if err != nil {
  972. iter.err = err
  973. return false
  974. }
  975. cols[i] = col
  976. }
  977. is.cols = cols
  978. iter.pos++
  979. return true
  980. }
  981. func scanColumn(p []byte, col ColumnInfo, dest []interface{}) (int, error) {
  982. if dest[0] == nil {
  983. return 1, nil
  984. }
  985. if col.TypeInfo.Type() == TypeTuple {
  986. // this will panic, actually a bug, please report
  987. tuple := col.TypeInfo.(TupleTypeInfo)
  988. count := len(tuple.Elems)
  989. // here we pass in a slice of the struct which has the number number of
  990. // values as elements in the tuple
  991. if err := Unmarshal(col.TypeInfo, p, dest[:count]); err != nil {
  992. return 0, err
  993. }
  994. return count, nil
  995. } else {
  996. if err := Unmarshal(col.TypeInfo, p, dest[0]); err != nil {
  997. return 0, err
  998. }
  999. return 1, nil
  1000. }
  1001. }
  1002. func (is *iterScanner) Scan(dest ...interface{}) error {
  1003. if is.cols == nil {
  1004. return errors.New("gocql: Scan called without calling Next")
  1005. }
  1006. iter := is.iter
  1007. // currently only support scanning into an expand tuple, such that its the same
  1008. // as scanning in more values from a single column
  1009. if len(dest) != iter.meta.actualColCount {
  1010. return fmt.Errorf("gocql: not enough columns to scan into: have %d want %d", len(dest), iter.meta.actualColCount)
  1011. }
  1012. // i is the current position in dest, could posible replace it and just use
  1013. // slices of dest
  1014. i := 0
  1015. var err error
  1016. for _, col := range iter.meta.columns {
  1017. var n int
  1018. n, err = scanColumn(is.cols[i], col, dest[i:])
  1019. if err != nil {
  1020. break
  1021. }
  1022. i += n
  1023. }
  1024. is.cols = nil
  1025. return err
  1026. }
  1027. func (is *iterScanner) Err() error {
  1028. iter := is.iter
  1029. is.iter = nil
  1030. is.cols = nil
  1031. return iter.Close()
  1032. }
  1033. // Scanner returns a row Scanner which provides an interface to scan rows in a manner which is
  1034. // similar to database/sql. The iter should NOT be used again after calling this method.
  1035. func (iter *Iter) Scanner() Scanner {
  1036. if iter == nil {
  1037. return nil
  1038. }
  1039. return &iterScanner{iter: iter}
  1040. }
  1041. func (iter *Iter) readColumn() ([]byte, error) {
  1042. return iter.framer.readBytesInternal()
  1043. }
  1044. // Scan consumes the next row of the iterator and copies the columns of the
  1045. // current row into the values pointed at by dest. Use nil as a dest value
  1046. // to skip the corresponding column. Scan might send additional queries
  1047. // to the database to retrieve the next set of rows if paging was enabled.
  1048. //
  1049. // Scan returns true if the row was successfully unmarshaled or false if the
  1050. // end of the result set was reached or if an error occurred. Close should
  1051. // be called afterwards to retrieve any potential errors.
  1052. func (iter *Iter) Scan(dest ...interface{}) bool {
  1053. if iter.err != nil {
  1054. return false
  1055. }
  1056. if iter.pos >= iter.numRows {
  1057. if iter.next != nil {
  1058. *iter = *iter.next.fetch()
  1059. return iter.Scan(dest...)
  1060. }
  1061. return false
  1062. }
  1063. if iter.next != nil && iter.pos == iter.next.pos {
  1064. go iter.next.fetch()
  1065. }
  1066. // currently only support scanning into an expand tuple, such that its the same
  1067. // as scanning in more values from a single column
  1068. if len(dest) != iter.meta.actualColCount {
  1069. iter.err = fmt.Errorf("gocql: not enough columns to scan into: have %d want %d", len(dest), iter.meta.actualColCount)
  1070. return false
  1071. }
  1072. // i is the current position in dest, could posible replace it and just use
  1073. // slices of dest
  1074. i := 0
  1075. for _, col := range iter.meta.columns {
  1076. colBytes, err := iter.readColumn()
  1077. if err != nil {
  1078. iter.err = err
  1079. return false
  1080. }
  1081. n, err := scanColumn(colBytes, col, dest[i:])
  1082. if err != nil {
  1083. iter.err = err
  1084. return false
  1085. }
  1086. i += n
  1087. }
  1088. iter.pos++
  1089. return true
  1090. }
  1091. // GetCustomPayload returns any parsed custom payload results if given in the
  1092. // response from Cassandra. Note that the result is not a copy.
  1093. //
  1094. // This additional feature of CQL Protocol v4
  1095. // allows additional results and query information to be returned by
  1096. // custom QueryHandlers running in your C* cluster.
  1097. // See https://datastax.github.io/java-driver/manual/custom_payloads/
  1098. func (iter *Iter) GetCustomPayload() map[string][]byte {
  1099. return iter.framer.header.customPayload
  1100. }
  1101. // Warnings returns any warnings generated if given in the response from Cassandra.
  1102. //
  1103. // This is only available starting with CQL Protocol v4.
  1104. func (iter *Iter) Warnings() []string {
  1105. if iter.framer != nil {
  1106. return iter.framer.header.warnings
  1107. }
  1108. return nil
  1109. }
  1110. // Close closes the iterator and returns any errors that happened during
  1111. // the query or the iteration.
  1112. func (iter *Iter) Close() error {
  1113. if atomic.CompareAndSwapInt32(&iter.closed, 0, 1) {
  1114. if iter.framer != nil {
  1115. framerPool.Put(iter.framer)
  1116. iter.framer = nil
  1117. }
  1118. }
  1119. return iter.err
  1120. }
  1121. // WillSwitchPage detects if iterator reached end of current page
  1122. // and the next page is available.
  1123. func (iter *Iter) WillSwitchPage() bool {
  1124. return iter.pos >= iter.numRows && iter.next != nil
  1125. }
  1126. // checkErrAndNotFound handle error and NotFound in one method.
  1127. func (iter *Iter) checkErrAndNotFound() error {
  1128. if iter.err != nil {
  1129. return iter.err
  1130. } else if iter.numRows == 0 {
  1131. return ErrNotFound
  1132. }
  1133. return nil
  1134. }
  1135. // PageState return the current paging state for a query which can be used for
  1136. // subsequent quries to resume paging this point.
  1137. func (iter *Iter) PageState() []byte {
  1138. return iter.meta.pagingState
  1139. }
  1140. // NumRows returns the number of rows in this pagination, it will update when new
  1141. // pages are fetched, it is not the value of the total number of rows this iter
  1142. // will return unless there is only a single page returned.
  1143. func (iter *Iter) NumRows() int {
  1144. return iter.numRows
  1145. }
  1146. type nextIter struct {
  1147. qry Query
  1148. pos int
  1149. once sync.Once
  1150. next *Iter
  1151. conn *Conn
  1152. }
  1153. func (n *nextIter) fetch() *Iter {
  1154. n.once.Do(func() {
  1155. iter := n.qry.session.executor.attemptQuery(&n.qry, n.conn)
  1156. if iter != nil && iter.err == nil {
  1157. n.next = iter
  1158. } else {
  1159. n.next = n.qry.session.executeQuery(&n.qry)
  1160. }
  1161. })
  1162. return n.next
  1163. }
  1164. type Batch struct {
  1165. Type BatchType
  1166. Entries []BatchEntry
  1167. Cons Consistency
  1168. rt RetryPolicy
  1169. observer BatchObserver
  1170. attempts int
  1171. totalLatency int64
  1172. serialCons SerialConsistency
  1173. defaultTimestamp bool
  1174. defaultTimestampValue int64
  1175. context context.Context
  1176. keyspace string
  1177. }
  1178. // NewBatch creates a new batch operation without defaults from the cluster
  1179. //
  1180. // Depreicated: use session.NewBatch instead
  1181. func NewBatch(typ BatchType) *Batch {
  1182. return &Batch{Type: typ}
  1183. }
  1184. // NewBatch creates a new batch operation using defaults defined in the cluster
  1185. func (s *Session) NewBatch(typ BatchType) *Batch {
  1186. s.mu.RLock()
  1187. batch := &Batch{
  1188. Type: typ,
  1189. rt: s.cfg.RetryPolicy,
  1190. serialCons: s.cfg.SerialConsistency,
  1191. observer: s.batchObserver,
  1192. Cons: s.cons,
  1193. defaultTimestamp: s.cfg.DefaultTimestamp,
  1194. keyspace: s.cfg.Keyspace,
  1195. }
  1196. s.mu.RUnlock()
  1197. return batch
  1198. }
  1199. // Observer enables batch-level observer on this batch.
  1200. // The provided observer will be called every time this batched query is executed.
  1201. func (b *Batch) Observer(observer BatchObserver) *Batch {
  1202. b.observer = observer
  1203. return b
  1204. }
  1205. func (b *Batch) Keyspace() string {
  1206. return b.keyspace
  1207. }
  1208. // Attempts returns the number of attempts made to execute the batch.
  1209. func (b *Batch) Attempts() int {
  1210. return b.attempts
  1211. }
  1212. //Latency returns the average number of nanoseconds to execute a single attempt of the batch.
  1213. func (b *Batch) Latency() int64 {
  1214. if b.attempts > 0 {
  1215. return b.totalLatency / int64(b.attempts)
  1216. }
  1217. return 0
  1218. }
  1219. // GetConsistency returns the currently configured consistency level for the batch
  1220. // operation.
  1221. func (b *Batch) GetConsistency() Consistency {
  1222. return b.Cons
  1223. }
  1224. // Query adds the query to the batch operation
  1225. func (b *Batch) Query(stmt string, args ...interface{}) {
  1226. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, Args: args})
  1227. }
  1228. // Bind adds the query to the batch operation and correlates it with a binding callback
  1229. // that will be invoked when the batch is executed. The binding callback allows the application
  1230. // to define which query argument values will be marshalled as part of the batch execution.
  1231. func (b *Batch) Bind(stmt string, bind func(q *QueryInfo) ([]interface{}, error)) {
  1232. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, binding: bind})
  1233. }
  1234. func (b *Batch) retryPolicy() RetryPolicy {
  1235. return b.rt
  1236. }
  1237. // RetryPolicy sets the retry policy to use when executing the batch operation
  1238. func (b *Batch) RetryPolicy(r RetryPolicy) *Batch {
  1239. b.rt = r
  1240. return b
  1241. }
  1242. // WithContext will set the context to use during a query, it will be used to
  1243. // timeout when waiting for responses from Cassandra.
  1244. func (b *Batch) WithContext(ctx context.Context) *Batch {
  1245. b.context = ctx
  1246. return b
  1247. }
  1248. // Size returns the number of batch statements to be executed by the batch operation.
  1249. func (b *Batch) Size() int {
  1250. return len(b.Entries)
  1251. }
  1252. // SerialConsistency sets the consistency level for the
  1253. // serial phase of conditional updates. That consistency can only be
  1254. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  1255. // SERIAL. This option will be ignored for anything else that a
  1256. // conditional update/insert.
  1257. //
  1258. // Only available for protocol 3 and above
  1259. func (b *Batch) SerialConsistency(cons SerialConsistency) *Batch {
  1260. b.serialCons = cons
  1261. return b
  1262. }
  1263. // DefaultTimestamp will enable the with default timestamp flag on the query.
  1264. // If enable, this will replace the server side assigned
  1265. // timestamp as default timestamp. Note that a timestamp in the query itself
  1266. // will still override this timestamp. This is entirely optional.
  1267. //
  1268. // Only available on protocol >= 3
  1269. func (b *Batch) DefaultTimestamp(enable bool) *Batch {
  1270. b.defaultTimestamp = enable
  1271. return b
  1272. }
  1273. // WithTimestamp will enable the with default timestamp flag on the query
  1274. // like DefaultTimestamp does. But also allows to define value for timestamp.
  1275. // It works the same way as USING TIMESTAMP in the query itself, but
  1276. // should not break prepared query optimization
  1277. //
  1278. // Only available on protocol >= 3
  1279. func (b *Batch) WithTimestamp(timestamp int64) *Batch {
  1280. b.DefaultTimestamp(true)
  1281. b.defaultTimestampValue = timestamp
  1282. return b
  1283. }
  1284. func (b *Batch) attempt(keyspace string, end, start time.Time, iter *Iter) {
  1285. b.attempts++
  1286. b.totalLatency += end.Sub(start).Nanoseconds()
  1287. // TODO: track latencies per host and things as well instead of just total
  1288. if b.observer == nil {
  1289. return
  1290. }
  1291. statements := make([]string, len(b.Entries))
  1292. for i, entry := range b.Entries {
  1293. statements[i] = entry.Stmt
  1294. }
  1295. b.observer.ObserveBatch(b.context, ObservedBatch{
  1296. Keyspace: keyspace,
  1297. Statements: statements,
  1298. Start: start,
  1299. End: end,
  1300. // Rows not used in batch observations // TODO - might be able to support it when using BatchCAS
  1301. Err: iter.err,
  1302. })
  1303. }
  1304. func (b *Batch) GetRoutingKey() ([]byte, error) {
  1305. // TODO: use the first statement in the batch as the routing key?
  1306. return nil, nil
  1307. }
  1308. type BatchType byte
  1309. const (
  1310. LoggedBatch BatchType = 0
  1311. UnloggedBatch BatchType = 1
  1312. CounterBatch BatchType = 2
  1313. )
  1314. type BatchEntry struct {
  1315. Stmt string
  1316. Args []interface{}
  1317. binding func(q *QueryInfo) ([]interface{}, error)
  1318. }
  1319. type ColumnInfo struct {
  1320. Keyspace string
  1321. Table string
  1322. Name string
  1323. TypeInfo TypeInfo
  1324. }
  1325. func (c ColumnInfo) String() string {
  1326. return fmt.Sprintf("[column keyspace=%s table=%s name=%s type=%v]", c.Keyspace, c.Table, c.Name, c.TypeInfo)
  1327. }
  1328. // routing key indexes LRU cache
  1329. type routingKeyInfoLRU struct {
  1330. lru *lru.Cache
  1331. mu sync.Mutex
  1332. }
  1333. type routingKeyInfo struct {
  1334. indexes []int
  1335. types []TypeInfo
  1336. }
  1337. func (r *routingKeyInfo) String() string {
  1338. return fmt.Sprintf("routing key index=%v types=%v", r.indexes, r.types)
  1339. }
  1340. func (r *routingKeyInfoLRU) Remove(key string) {
  1341. r.mu.Lock()
  1342. r.lru.Remove(key)
  1343. r.mu.Unlock()
  1344. }
  1345. //Max adjusts the maximum size of the cache and cleans up the oldest records if
  1346. //the new max is lower than the previous value. Not concurrency safe.
  1347. func (r *routingKeyInfoLRU) Max(max int) {
  1348. r.mu.Lock()
  1349. for r.lru.Len() > max {
  1350. r.lru.RemoveOldest()
  1351. }
  1352. r.lru.MaxEntries = max
  1353. r.mu.Unlock()
  1354. }
  1355. type inflightCachedEntry struct {
  1356. wg sync.WaitGroup
  1357. err error
  1358. value interface{}
  1359. }
  1360. // Tracer is the interface implemented by query tracers. Tracers have the
  1361. // ability to obtain a detailed event log of all events that happened during
  1362. // the execution of a query from Cassandra. Gathering this information might
  1363. // be essential for debugging and optimizing queries, but this feature should
  1364. // not be used on production systems with very high load.
  1365. type Tracer interface {
  1366. Trace(traceId []byte)
  1367. }
  1368. type traceWriter struct {
  1369. session *Session
  1370. w io.Writer
  1371. mu sync.Mutex
  1372. }
  1373. // NewTraceWriter returns a simple Tracer implementation that outputs
  1374. // the event log in a textual format.
  1375. func NewTraceWriter(session *Session, w io.Writer) Tracer {
  1376. return &traceWriter{session: session, w: w}
  1377. }
  1378. func (t *traceWriter) Trace(traceId []byte) {
  1379. var (
  1380. coordinator string
  1381. duration int
  1382. )
  1383. iter := t.session.control.query(`SELECT coordinator, duration
  1384. FROM system_traces.sessions
  1385. WHERE session_id = ?`, traceId)
  1386. iter.Scan(&coordinator, &duration)
  1387. if err := iter.Close(); err != nil {
  1388. t.mu.Lock()
  1389. fmt.Fprintln(t.w, "Error:", err)
  1390. t.mu.Unlock()
  1391. return
  1392. }
  1393. var (
  1394. timestamp time.Time
  1395. activity string
  1396. source string
  1397. elapsed int
  1398. )
  1399. fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
  1400. traceId, coordinator, time.Duration(duration)*time.Microsecond)
  1401. t.mu.Lock()
  1402. defer t.mu.Unlock()
  1403. iter = t.session.control.query(`SELECT event_id, activity, source, source_elapsed
  1404. FROM system_traces.events
  1405. WHERE session_id = ?`, traceId)
  1406. for iter.Scan(&timestamp, &activity, &source, &elapsed) {
  1407. fmt.Fprintf(t.w, "%s: %s (source: %s, elapsed: %d)\n",
  1408. timestamp.Format("2006/01/02 15:04:05.999999"), activity, source, elapsed)
  1409. }
  1410. if err := iter.Close(); err != nil {
  1411. fmt.Fprintln(t.w, "Error:", err)
  1412. }
  1413. }
  1414. type ObservedQuery struct {
  1415. Keyspace string
  1416. Statement string
  1417. Start time.Time // time immediately before the query was called
  1418. End time.Time // time immediately after the query returned
  1419. // Rows is the number of rows in the current iter.
  1420. // In paginated queries, rows from previous scans are not counted.
  1421. // Rows is not used in batch queries and remains at the default value
  1422. Rows int
  1423. // Err is the error in the query.
  1424. // It only tracks network errors or errors of bad cassandra syntax, in particular selects with no match return nil error
  1425. Err error
  1426. }
  1427. // QueryObserver is the interface implemented by query observers / stat collectors.
  1428. //
  1429. // Experimental, this interface and use may change
  1430. type QueryObserver interface {
  1431. // ObserveQuery gets called on every query to cassandra, including all queries in an iterator when paging is enabled.
  1432. // It doesn't get called if there is no query because the session is closed or there are no connections available.
  1433. // The error reported only shows query errors, i.e. if a SELECT is valid but finds no matches it will be nil.
  1434. ObserveQuery(context.Context, ObservedQuery)
  1435. }
  1436. type ObservedBatch struct {
  1437. Keyspace string
  1438. Statements []string
  1439. Start time.Time // time immediately before the batch query was called
  1440. End time.Time // time immediately after the batch query returned
  1441. // Err is the error in the batch query.
  1442. // It only tracks network errors or errors of bad cassandra syntax, in particular selects with no match return nil error
  1443. Err error
  1444. }
  1445. // BatchObserver is the interface implemented by batch observers / stat collectors.
  1446. type BatchObserver interface {
  1447. // ObserveBatch gets called on every batch query to cassandra.
  1448. // It also gets called once for each query in a batch.
  1449. // It doesn't get called if there is no query because the session is closed or there are no connections available.
  1450. // The error reported only shows query errors, i.e. if a SELECT is valid but finds no matches it will be nil.
  1451. // Unlike QueryObserver.ObserveQuery it does no reporting on rows read.
  1452. ObserveBatch(context.Context, ObservedBatch)
  1453. }
  1454. type Error struct {
  1455. Code int
  1456. Message string
  1457. }
  1458. func (e Error) Error() string {
  1459. return e.Message
  1460. }
  1461. var (
  1462. ErrNotFound = errors.New("not found")
  1463. ErrUnavailable = errors.New("unavailable")
  1464. ErrUnsupported = errors.New("feature not supported")
  1465. ErrTooManyStmts = errors.New("too many statements")
  1466. ErrUseStmt = errors.New("use statements aren't supported. Please see https://github.com/gocql/gocql for explanation.")
  1467. ErrSessionClosed = errors.New("session has been closed")
  1468. ErrNoConnections = errors.New("gocql: no hosts available in the pool")
  1469. ErrNoKeyspace = errors.New("no keyspace provided")
  1470. ErrKeyspaceDoesNotExist = errors.New("keyspace does not exist")
  1471. ErrNoMetadata = errors.New("no metadata available")
  1472. )
  1473. type ErrProtocol struct{ error }
  1474. func NewErrProtocol(format string, args ...interface{}) error {
  1475. return ErrProtocol{fmt.Errorf(format, args...)}
  1476. }
  1477. // BatchSizeMaximum is the maximum number of statements a batch operation can have.
  1478. // This limit is set by cassandra and could change in the future.
  1479. const BatchSizeMaximum = 65535