session.go 36 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334
  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. "encoding/binary"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "net"
  12. "strconv"
  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 consinstency 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. hostSource *ringDescriber
  37. stmtsLRU *preparedLRU
  38. connCfg *ConnConfig
  39. executor *queryExecutor
  40. pool *policyConnPool
  41. policy HostSelectionPolicy
  42. ring ring
  43. metadata clusterMetadata
  44. mu sync.RWMutex
  45. control *controlConn
  46. // event handlers
  47. nodeEvents *eventDeouncer
  48. schemaEvents *eventDeouncer
  49. // ring metadata
  50. hosts []HostInfo
  51. useSystemSchema bool
  52. cfg ClusterConfig
  53. closeMu sync.RWMutex
  54. isClosed bool
  55. }
  56. func addrsToHosts(addrs []string, defaultPort int) ([]*HostInfo, error) {
  57. hosts := make([]*HostInfo, len(addrs))
  58. for i, hostport := range addrs {
  59. // TODO: remove duplication
  60. addr, portStr, err := net.SplitHostPort(JoinHostPort(hostport, defaultPort))
  61. if err != nil {
  62. return nil, fmt.Errorf("NewSession: unable to parse hostport of addr %q: %v", hostport, err)
  63. }
  64. port, err := strconv.Atoi(portStr)
  65. if err != nil {
  66. return nil, fmt.Errorf("NewSession: invalid port for hostport of addr %q: %v", hostport, err)
  67. }
  68. hosts[i] = &HostInfo{peer: addr, port: port, state: NodeUp}
  69. }
  70. return hosts, nil
  71. }
  72. // NewSession wraps an existing Node.
  73. func NewSession(cfg ClusterConfig) (*Session, error) {
  74. //Check that hosts in the ClusterConfig is not empty
  75. if len(cfg.Hosts) < 1 {
  76. return nil, ErrNoHosts
  77. }
  78. s := &Session{
  79. cons: cfg.Consistency,
  80. prefetch: 0.25,
  81. cfg: cfg,
  82. pageSize: cfg.PageSize,
  83. stmtsLRU: &preparedLRU{lru: lru.New(cfg.MaxPreparedStmts)},
  84. }
  85. connCfg, err := connConfig(s)
  86. if err != nil {
  87. s.Close()
  88. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  89. }
  90. s.connCfg = connCfg
  91. s.nodeEvents = newEventDeouncer("NodeEvents", s.handleNodeEvent)
  92. s.schemaEvents = newEventDeouncer("SchemaEvents", s.handleSchemaEvent)
  93. s.routingKeyInfoCache.lru = lru.New(cfg.MaxRoutingKeyInfo)
  94. // I think it might be a good idea to simplify this and make it always discover
  95. // hosts, maybe with more filters.
  96. s.hostSource = &ringDescriber{
  97. session: s,
  98. closeChan: make(chan bool),
  99. }
  100. pool := cfg.PoolConfig.buildPool(s)
  101. if cfg.PoolConfig.HostSelectionPolicy == nil {
  102. cfg.PoolConfig.HostSelectionPolicy = RoundRobinHostPolicy()
  103. }
  104. s.pool = pool
  105. s.policy = cfg.PoolConfig.HostSelectionPolicy
  106. s.executor = &queryExecutor{
  107. pool: pool,
  108. policy: cfg.PoolConfig.HostSelectionPolicy,
  109. }
  110. var hosts []*HostInfo
  111. if !cfg.disableControlConn {
  112. s.control = createControlConn(s)
  113. if err := s.control.connect(cfg.Hosts); err != nil {
  114. s.Close()
  115. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  116. }
  117. // need to setup host source to check for broadcast_address in system.local
  118. localHasRPCAddr, _ := checkSystemLocal(s.control)
  119. s.hostSource.localHasRpcAddr = localHasRPCAddr
  120. var err error
  121. if cfg.DisableInitialHostLookup {
  122. // TODO: we could look at system.local to get token and other metadata
  123. // in this case.
  124. hosts, err = addrsToHosts(cfg.Hosts, cfg.Port)
  125. } else {
  126. hosts, _, err = s.hostSource.GetHosts()
  127. }
  128. if err != nil {
  129. s.Close()
  130. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  131. }
  132. } else {
  133. // we dont get host info
  134. hosts, err = addrsToHosts(cfg.Hosts, cfg.Port)
  135. }
  136. for _, host := range hosts {
  137. if s.cfg.HostFilter == nil || s.cfg.HostFilter.Accept(host) {
  138. if existingHost, ok := s.ring.addHostIfMissing(host); ok {
  139. existingHost.update(host)
  140. }
  141. s.handleNodeUp(net.ParseIP(host.Peer()), host.Port(), false)
  142. }
  143. }
  144. // TODO(zariel): we probably dont need this any more as we verify that we
  145. // can connect to one of the endpoints supplied by using the control conn.
  146. // See if there are any connections in the pool
  147. if s.pool.Size() == 0 {
  148. s.Close()
  149. return nil, ErrNoConnectionsStarted
  150. }
  151. s.useSystemSchema = hosts[0].Version().Major >= 3
  152. return s, nil
  153. }
  154. // SetConsistency sets the default consistency level for this session. This
  155. // setting can also be changed on a per-query basis and the default value
  156. // is Quorum.
  157. func (s *Session) SetConsistency(cons Consistency) {
  158. s.mu.Lock()
  159. s.cons = cons
  160. s.mu.Unlock()
  161. }
  162. // SetPageSize sets the default page size for this session. A value <= 0 will
  163. // disable paging. This setting can also be changed on a per-query basis.
  164. func (s *Session) SetPageSize(n int) {
  165. s.mu.Lock()
  166. s.pageSize = n
  167. s.mu.Unlock()
  168. }
  169. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  170. // there are only p*pageSize rows remaining, the next page will be requested
  171. // automatically. This value can also be changed on a per-query basis and
  172. // the default value is 0.25.
  173. func (s *Session) SetPrefetch(p float64) {
  174. s.mu.Lock()
  175. s.prefetch = p
  176. s.mu.Unlock()
  177. }
  178. // SetTrace sets the default tracer for this session. This setting can also
  179. // be changed on a per-query basis.
  180. func (s *Session) SetTrace(trace Tracer) {
  181. s.mu.Lock()
  182. s.trace = trace
  183. s.mu.Unlock()
  184. }
  185. // Query generates a new query object for interacting with the database.
  186. // Further details of the query may be tweaked using the resulting query
  187. // value before the query is executed. Query is automatically prepared
  188. // if it has not previously been executed.
  189. func (s *Session) Query(stmt string, values ...interface{}) *Query {
  190. s.mu.RLock()
  191. qry := &Query{stmt: stmt, values: values, cons: s.cons,
  192. session: s, pageSize: s.pageSize, trace: s.trace,
  193. prefetch: s.prefetch, rt: s.cfg.RetryPolicy, serialCons: s.cfg.SerialConsistency,
  194. defaultTimestamp: s.cfg.DefaultTimestamp,
  195. }
  196. s.mu.RUnlock()
  197. return qry
  198. }
  199. type QueryInfo struct {
  200. Id []byte
  201. Args []ColumnInfo
  202. Rval []ColumnInfo
  203. PKeyColumns []int
  204. }
  205. // Bind generates a new query object based on the query statement passed in.
  206. // The query is automatically prepared if it has not previously been executed.
  207. // The binding callback allows the application to define which query argument
  208. // values will be marshalled as part of the query execution.
  209. // During execution, the meta data of the prepared query will be routed to the
  210. // binding callback, which is responsible for producing the query argument values.
  211. func (s *Session) Bind(stmt string, b func(q *QueryInfo) ([]interface{}, error)) *Query {
  212. s.mu.RLock()
  213. qry := &Query{stmt: stmt, binding: b, cons: s.cons,
  214. session: s, pageSize: s.pageSize, trace: s.trace,
  215. prefetch: s.prefetch, rt: s.cfg.RetryPolicy}
  216. s.mu.RUnlock()
  217. return qry
  218. }
  219. // Close closes all connections. The session is unusable after this
  220. // operation.
  221. func (s *Session) Close() {
  222. s.closeMu.Lock()
  223. defer s.closeMu.Unlock()
  224. if s.isClosed {
  225. return
  226. }
  227. s.isClosed = true
  228. if s.pool != nil {
  229. s.pool.Close()
  230. }
  231. if s.hostSource != nil {
  232. close(s.hostSource.closeChan)
  233. }
  234. if s.control != nil {
  235. s.control.close()
  236. }
  237. if s.nodeEvents != nil {
  238. s.nodeEvents.stop()
  239. }
  240. if s.schemaEvents != nil {
  241. s.schemaEvents.stop()
  242. }
  243. }
  244. func (s *Session) Closed() bool {
  245. s.closeMu.RLock()
  246. closed := s.isClosed
  247. s.closeMu.RUnlock()
  248. return closed
  249. }
  250. func (s *Session) executeQuery(qry *Query) *Iter {
  251. // fail fast
  252. if s.Closed() {
  253. return &Iter{err: ErrSessionClosed}
  254. }
  255. iter, err := s.executor.executeQuery(qry)
  256. if err != nil {
  257. return &Iter{err: err}
  258. }
  259. if iter == nil {
  260. panic("nil iter")
  261. }
  262. return iter
  263. }
  264. // KeyspaceMetadata returns the schema metadata for the keyspace specified.
  265. func (s *Session) KeyspaceMetadata(keyspace string) (*KeyspaceMetadata, error) {
  266. // fail fast
  267. if s.Closed() {
  268. return nil, ErrSessionClosed
  269. }
  270. if keyspace == "" {
  271. return nil, ErrNoKeyspace
  272. }
  273. s.mu.Lock()
  274. // lazy-init schemaDescriber
  275. if s.schemaDescriber == nil {
  276. s.schemaDescriber = newSchemaDescriber(s)
  277. }
  278. s.mu.Unlock()
  279. return s.schemaDescriber.getSchema(keyspace)
  280. }
  281. func (s *Session) getConn() *Conn {
  282. hosts := s.ring.allHosts()
  283. var conn *Conn
  284. for _, host := range hosts {
  285. if !host.IsUp() {
  286. continue
  287. }
  288. pool, ok := s.pool.getPool(host.Peer())
  289. if !ok {
  290. continue
  291. }
  292. conn = pool.Pick()
  293. if conn != nil {
  294. return conn
  295. }
  296. }
  297. return nil
  298. }
  299. // returns routing key indexes and type info
  300. func (s *Session) routingKeyInfo(stmt string) (*routingKeyInfo, error) {
  301. s.routingKeyInfoCache.mu.Lock()
  302. entry, cached := s.routingKeyInfoCache.lru.Get(stmt)
  303. if cached {
  304. // done accessing the cache
  305. s.routingKeyInfoCache.mu.Unlock()
  306. // the entry is an inflight struct similiar to that used by
  307. // Conn to prepare statements
  308. inflight := entry.(*inflightCachedEntry)
  309. // wait for any inflight work
  310. inflight.wg.Wait()
  311. if inflight.err != nil {
  312. return nil, inflight.err
  313. }
  314. key, _ := inflight.value.(*routingKeyInfo)
  315. return key, nil
  316. }
  317. // create a new inflight entry while the data is created
  318. inflight := new(inflightCachedEntry)
  319. inflight.wg.Add(1)
  320. defer inflight.wg.Done()
  321. s.routingKeyInfoCache.lru.Add(stmt, inflight)
  322. s.routingKeyInfoCache.mu.Unlock()
  323. var (
  324. info *preparedStatment
  325. partitionKey []*ColumnMetadata
  326. )
  327. conn := s.getConn()
  328. if conn == nil {
  329. // TODO: better error?
  330. inflight.err = errors.New("gocql: unable to fetch preapred info: no connection avilable")
  331. return nil, inflight.err
  332. }
  333. // get the query info for the statement
  334. info, inflight.err = conn.prepareStatement(stmt, nil)
  335. if inflight.err != nil {
  336. // don't cache this error
  337. s.routingKeyInfoCache.Remove(stmt)
  338. return nil, inflight.err
  339. }
  340. // TODO: it would be nice to mark hosts here but as we are not using the policies
  341. // to fetch hosts we cant
  342. if info.request.colCount == 0 {
  343. // no arguments, no routing key, and no error
  344. return nil, nil
  345. }
  346. if len(info.request.pkeyColumns) > 0 {
  347. // proto v4 dont need to calculate primary key columns
  348. types := make([]TypeInfo, len(info.request.pkeyColumns))
  349. for i, col := range info.request.pkeyColumns {
  350. types[i] = info.request.columns[col].TypeInfo
  351. }
  352. routingKeyInfo := &routingKeyInfo{
  353. indexes: info.request.pkeyColumns,
  354. types: types,
  355. }
  356. inflight.value = routingKeyInfo
  357. return routingKeyInfo, nil
  358. }
  359. // get the table metadata
  360. table := info.request.columns[0].Table
  361. var keyspaceMetadata *KeyspaceMetadata
  362. keyspaceMetadata, inflight.err = s.KeyspaceMetadata(info.request.columns[0].Keyspace)
  363. if inflight.err != nil {
  364. // don't cache this error
  365. s.routingKeyInfoCache.Remove(stmt)
  366. return nil, inflight.err
  367. }
  368. tableMetadata, found := keyspaceMetadata.Tables[table]
  369. if !found {
  370. // unlikely that the statement could be prepared and the metadata for
  371. // the table couldn't be found, but this may indicate either a bug
  372. // in the metadata code, or that the table was just dropped.
  373. inflight.err = ErrNoMetadata
  374. // don't cache this error
  375. s.routingKeyInfoCache.Remove(stmt)
  376. return nil, inflight.err
  377. }
  378. partitionKey = tableMetadata.PartitionKey
  379. size := len(partitionKey)
  380. routingKeyInfo := &routingKeyInfo{
  381. indexes: make([]int, size),
  382. types: make([]TypeInfo, size),
  383. }
  384. for keyIndex, keyColumn := range partitionKey {
  385. // set an indicator for checking if the mapping is missing
  386. routingKeyInfo.indexes[keyIndex] = -1
  387. // find the column in the query info
  388. for argIndex, boundColumn := range info.request.columns {
  389. if keyColumn.Name == boundColumn.Name {
  390. // there may be many such bound columns, pick the first
  391. routingKeyInfo.indexes[keyIndex] = argIndex
  392. routingKeyInfo.types[keyIndex] = boundColumn.TypeInfo
  393. break
  394. }
  395. }
  396. if routingKeyInfo.indexes[keyIndex] == -1 {
  397. // missing a routing key column mapping
  398. // no routing key, and no error
  399. return nil, nil
  400. }
  401. }
  402. // cache this result
  403. inflight.value = routingKeyInfo
  404. return routingKeyInfo, nil
  405. }
  406. func (b *Batch) execute(conn *Conn) *Iter {
  407. return conn.executeBatch(b)
  408. }
  409. func (s *Session) executeBatch(batch *Batch) *Iter {
  410. // fail fast
  411. if s.Closed() {
  412. return &Iter{err: ErrSessionClosed}
  413. }
  414. // Prevent the execution of the batch if greater than the limit
  415. // Currently batches have a limit of 65536 queries.
  416. // https://datastax-oss.atlassian.net/browse/JAVA-229
  417. if batch.Size() > BatchSizeMaximum {
  418. return &Iter{err: ErrTooManyStmts}
  419. }
  420. iter, err := s.executor.executeQuery(batch)
  421. if err != nil {
  422. return &Iter{err: err}
  423. }
  424. return iter
  425. }
  426. // ExecuteBatch executes a batch operation and returns nil if successful
  427. // otherwise an error is returned describing the failure.
  428. func (s *Session) ExecuteBatch(batch *Batch) error {
  429. iter := s.executeBatch(batch)
  430. return iter.Close()
  431. }
  432. // ExecuteBatchCAS executes a batch operation and returns true if successful and
  433. // an iterator (to scan aditional rows if more than one conditional statement)
  434. // was sent.
  435. // Further scans on the interator must also remember to include
  436. // the applied boolean as the first argument to *Iter.Scan
  437. func (s *Session) ExecuteBatchCAS(batch *Batch, dest ...interface{}) (applied bool, iter *Iter, err error) {
  438. iter = s.executeBatch(batch)
  439. if err := iter.checkErrAndNotFound(); err != nil {
  440. iter.Close()
  441. return false, nil, err
  442. }
  443. if len(iter.Columns()) > 1 {
  444. dest = append([]interface{}{&applied}, dest...)
  445. iter.Scan(dest...)
  446. } else {
  447. iter.Scan(&applied)
  448. }
  449. return applied, iter, nil
  450. }
  451. // MapExecuteBatchCAS executes a batch operation much like ExecuteBatchCAS,
  452. // however it accepts a map rather than a list of arguments for the initial
  453. // scan.
  454. func (s *Session) MapExecuteBatchCAS(batch *Batch, dest map[string]interface{}) (applied bool, iter *Iter, err error) {
  455. iter = s.executeBatch(batch)
  456. if err := iter.checkErrAndNotFound(); err != nil {
  457. iter.Close()
  458. return false, nil, err
  459. }
  460. iter.MapScan(dest)
  461. applied = dest["[applied]"].(bool)
  462. delete(dest, "[applied]")
  463. // we usually close here, but instead of closing, just returin an error
  464. // if MapScan failed. Although Close just returns err, using Close
  465. // here might be confusing as we are not actually closing the iter
  466. return applied, iter, iter.err
  467. }
  468. func (s *Session) connect(addr string, errorHandler ConnErrorHandler, host *HostInfo) (*Conn, error) {
  469. return Connect(host, addr, s.connCfg, errorHandler, s)
  470. }
  471. // Query represents a CQL statement that can be executed.
  472. type Query struct {
  473. stmt string
  474. values []interface{}
  475. cons Consistency
  476. pageSize int
  477. routingKey []byte
  478. routingKeyBuffer []byte
  479. pageState []byte
  480. prefetch float64
  481. trace Tracer
  482. session *Session
  483. rt RetryPolicy
  484. binding func(q *QueryInfo) ([]interface{}, error)
  485. attempts int
  486. totalLatency int64
  487. serialCons SerialConsistency
  488. defaultTimestamp bool
  489. defaultTimestampValue int64
  490. disableSkipMetadata bool
  491. disableAutoPage bool
  492. }
  493. // String implements the stringer interface.
  494. func (q Query) String() string {
  495. return fmt.Sprintf("[query statement=%q values=%+v consistency=%s]", q.stmt, q.values, q.cons)
  496. }
  497. //Attempts returns the number of times the query was executed.
  498. func (q *Query) Attempts() int {
  499. return q.attempts
  500. }
  501. //Latency returns the average amount of nanoseconds per attempt of the query.
  502. func (q *Query) Latency() int64 {
  503. if q.attempts > 0 {
  504. return q.totalLatency / int64(q.attempts)
  505. }
  506. return 0
  507. }
  508. // Consistency sets the consistency level for this query. If no consistency
  509. // level have been set, the default consistency level of the cluster
  510. // is used.
  511. func (q *Query) Consistency(c Consistency) *Query {
  512. q.cons = c
  513. return q
  514. }
  515. // GetConsistency returns the currently configured consistency level for
  516. // the query.
  517. func (q *Query) GetConsistency() Consistency {
  518. return q.cons
  519. }
  520. // Trace enables tracing of this query. Look at the documentation of the
  521. // Tracer interface to learn more about tracing.
  522. func (q *Query) Trace(trace Tracer) *Query {
  523. q.trace = trace
  524. return q
  525. }
  526. // PageSize will tell the iterator to fetch the result in pages of size n.
  527. // This is useful for iterating over large result sets, but setting the
  528. // page size too low might decrease the performance. This feature is only
  529. // available in Cassandra 2 and onwards.
  530. func (q *Query) PageSize(n int) *Query {
  531. q.pageSize = n
  532. return q
  533. }
  534. // DefaultTimestamp will enable the with default timestamp flag on the query.
  535. // If enable, this will replace the server side assigned
  536. // timestamp as default timestamp. Note that a timestamp in the query itself
  537. // will still override this timestamp. This is entirely optional.
  538. //
  539. // Only available on protocol >= 3
  540. func (q *Query) DefaultTimestamp(enable bool) *Query {
  541. q.defaultTimestamp = enable
  542. return q
  543. }
  544. // WithTimestamp will enable the with default timestamp flag on the query
  545. // like DefaultTimestamp does. But also allows to define value for timestamp.
  546. // It works the same way as USING TIMESTAMP in the query itself, but
  547. // should not break prepared query optimization
  548. //
  549. // Only available on protocol >= 3
  550. func (q *Query) WithTimestamp(timestamp int64) *Query {
  551. q.DefaultTimestamp(true)
  552. q.defaultTimestampValue = timestamp
  553. return q
  554. }
  555. // RoutingKey sets the routing key to use when a token aware connection
  556. // pool is used to optimize the routing of this query.
  557. func (q *Query) RoutingKey(routingKey []byte) *Query {
  558. q.routingKey = routingKey
  559. return q
  560. }
  561. func (q *Query) execute(conn *Conn) *Iter {
  562. return conn.executeQuery(q)
  563. }
  564. func (q *Query) attempt(d time.Duration) {
  565. q.attempts++
  566. q.totalLatency += d.Nanoseconds()
  567. // TODO: track latencies per host and things as well instead of just total
  568. }
  569. func (q *Query) retryPolicy() RetryPolicy {
  570. return q.rt
  571. }
  572. // GetRoutingKey gets the routing key to use for routing this query. If
  573. // a routing key has not been explicitly set, then the routing key will
  574. // be constructed if possible using the keyspace's schema and the query
  575. // info for this query statement. If the routing key cannot be determined
  576. // then nil will be returned with no error. On any error condition,
  577. // an error description will be returned.
  578. func (q *Query) GetRoutingKey() ([]byte, error) {
  579. if q.routingKey != nil {
  580. return q.routingKey, nil
  581. } else if q.binding != nil && len(q.values) == 0 {
  582. // If this query was created using session.Bind we wont have the query
  583. // values yet, so we have to pass down to the next policy.
  584. // TODO: Remove this and handle this case
  585. return nil, nil
  586. }
  587. // try to determine the routing key
  588. routingKeyInfo, err := q.session.routingKeyInfo(q.stmt)
  589. if err != nil {
  590. return nil, err
  591. }
  592. if routingKeyInfo == nil {
  593. return nil, nil
  594. }
  595. if len(routingKeyInfo.indexes) == 1 {
  596. // single column routing key
  597. routingKey, err := Marshal(
  598. routingKeyInfo.types[0],
  599. q.values[routingKeyInfo.indexes[0]],
  600. )
  601. if err != nil {
  602. return nil, err
  603. }
  604. return routingKey, nil
  605. }
  606. // We allocate that buffer only once, so that further re-bind/exec of the
  607. // same query don't allocate more memory.
  608. if q.routingKeyBuffer == nil {
  609. q.routingKeyBuffer = make([]byte, 0, 256)
  610. }
  611. // composite routing key
  612. buf := bytes.NewBuffer(q.routingKeyBuffer)
  613. for i := range routingKeyInfo.indexes {
  614. encoded, err := Marshal(
  615. routingKeyInfo.types[i],
  616. q.values[routingKeyInfo.indexes[i]],
  617. )
  618. if err != nil {
  619. return nil, err
  620. }
  621. lenBuf := []byte{0x00, 0x00}
  622. binary.BigEndian.PutUint16(lenBuf, uint16(len(encoded)))
  623. buf.Write(lenBuf)
  624. buf.Write(encoded)
  625. buf.WriteByte(0x00)
  626. }
  627. routingKey := buf.Bytes()
  628. return routingKey, nil
  629. }
  630. func (q *Query) shouldPrepare() bool {
  631. stmt := strings.TrimLeftFunc(strings.TrimRightFunc(q.stmt, func(r rune) bool {
  632. return unicode.IsSpace(r) || r == ';'
  633. }), unicode.IsSpace)
  634. var stmtType string
  635. if n := strings.IndexFunc(stmt, unicode.IsSpace); n >= 0 {
  636. stmtType = strings.ToLower(stmt[:n])
  637. }
  638. if stmtType == "begin" {
  639. if n := strings.LastIndexFunc(stmt, unicode.IsSpace); n >= 0 {
  640. stmtType = strings.ToLower(stmt[n+1:])
  641. }
  642. }
  643. switch stmtType {
  644. case "select", "insert", "update", "delete", "batch":
  645. return true
  646. }
  647. return false
  648. }
  649. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  650. // there are only p*pageSize rows remaining, the next page will be requested
  651. // automatically.
  652. func (q *Query) Prefetch(p float64) *Query {
  653. q.prefetch = p
  654. return q
  655. }
  656. // RetryPolicy sets the policy to use when retrying the query.
  657. func (q *Query) RetryPolicy(r RetryPolicy) *Query {
  658. q.rt = r
  659. return q
  660. }
  661. // Bind sets query arguments of query. This can also be used to rebind new query arguments
  662. // to an existing query instance.
  663. func (q *Query) Bind(v ...interface{}) *Query {
  664. q.values = v
  665. return q
  666. }
  667. // SerialConsistency sets the consistency level for the
  668. // serial phase of conditional updates. That consistency can only be
  669. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  670. // SERIAL. This option will be ignored for anything else that a
  671. // conditional update/insert.
  672. func (q *Query) SerialConsistency(cons SerialConsistency) *Query {
  673. q.serialCons = cons
  674. return q
  675. }
  676. // PageState sets the paging state for the query to resume paging from a specific
  677. // point in time. Setting this will disable to query paging for this query, and
  678. // must be used for all subsequent pages.
  679. func (q *Query) PageState(state []byte) *Query {
  680. q.pageState = state
  681. q.disableAutoPage = true
  682. return q
  683. }
  684. // NoSkipMetadata will override the internal result metadata cache so that the driver does not
  685. // send skip_metadata for queries, this means that the result will always contain
  686. // the metadata to parse the rows and will not reuse the metadata from the prepared
  687. // staement. This should only be used to work around cassandra bugs, such as when using
  688. // CAS operations which do not end in Cas.
  689. //
  690. // See https://issues.apache.org/jira/browse/CASSANDRA-11099
  691. // https://github.com/gocql/gocql/issues/612
  692. func (q *Query) NoSkipMetadata() *Query {
  693. q.disableSkipMetadata = true
  694. return q
  695. }
  696. // Exec executes the query without returning any rows.
  697. func (q *Query) Exec() error {
  698. return q.Iter().Close()
  699. }
  700. func isUseStatement(stmt string) bool {
  701. if len(stmt) < 3 {
  702. return false
  703. }
  704. return strings.ToLower(stmt[0:3]) == "use"
  705. }
  706. // Iter executes the query and returns an iterator capable of iterating
  707. // over all results.
  708. func (q *Query) Iter() *Iter {
  709. if isUseStatement(q.stmt) {
  710. return &Iter{err: ErrUseStmt}
  711. }
  712. return q.session.executeQuery(q)
  713. }
  714. // MapScan executes the query, copies the columns of the first selected
  715. // row into the map pointed at by m and discards the rest. If no rows
  716. // were selected, ErrNotFound is returned.
  717. func (q *Query) MapScan(m map[string]interface{}) error {
  718. iter := q.Iter()
  719. if err := iter.checkErrAndNotFound(); err != nil {
  720. return err
  721. }
  722. iter.MapScan(m)
  723. return iter.Close()
  724. }
  725. // Scan executes the query, copies the columns of the first selected
  726. // row into the values pointed at by dest and discards the rest. If no rows
  727. // were selected, ErrNotFound is returned.
  728. func (q *Query) Scan(dest ...interface{}) error {
  729. iter := q.Iter()
  730. if err := iter.checkErrAndNotFound(); err != nil {
  731. return err
  732. }
  733. iter.Scan(dest...)
  734. return iter.Close()
  735. }
  736. // ScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  737. // statement containing an IF clause). If the transaction fails because
  738. // the existing values did not match, the previous values will be stored
  739. // in dest.
  740. func (q *Query) ScanCAS(dest ...interface{}) (applied bool, err error) {
  741. q.disableSkipMetadata = true
  742. iter := q.Iter()
  743. if err := iter.checkErrAndNotFound(); err != nil {
  744. return false, err
  745. }
  746. if len(iter.Columns()) > 1 {
  747. dest = append([]interface{}{&applied}, dest...)
  748. iter.Scan(dest...)
  749. } else {
  750. iter.Scan(&applied)
  751. }
  752. return applied, iter.Close()
  753. }
  754. // MapScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  755. // statement containing an IF clause). If the transaction fails because
  756. // the existing values did not match, the previous values will be stored
  757. // in dest map.
  758. //
  759. // As for INSERT .. IF NOT EXISTS, previous values will be returned as if
  760. // SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
  761. // column mismatching. MapScanCAS is added to capture them safely.
  762. func (q *Query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
  763. q.disableSkipMetadata = true
  764. iter := q.Iter()
  765. if err := iter.checkErrAndNotFound(); err != nil {
  766. return false, err
  767. }
  768. iter.MapScan(dest)
  769. applied = dest["[applied]"].(bool)
  770. delete(dest, "[applied]")
  771. return applied, iter.Close()
  772. }
  773. // Iter represents an iterator that can be used to iterate over all rows that
  774. // were returned by a query. The iterator might send additional queries to the
  775. // database during the iteration if paging was enabled.
  776. type Iter struct {
  777. err error
  778. pos int
  779. meta resultMetadata
  780. numRows int
  781. next *nextIter
  782. host *HostInfo
  783. framer *framer
  784. closed int32
  785. }
  786. // Host returns the host which the query was sent to.
  787. func (iter *Iter) Host() *HostInfo {
  788. return iter.host
  789. }
  790. // Columns returns the name and type of the selected columns.
  791. func (iter *Iter) Columns() []ColumnInfo {
  792. return iter.meta.columns
  793. }
  794. func (iter *Iter) readColumn() ([]byte, error) {
  795. return iter.framer.readBytesInternal()
  796. }
  797. // Scan consumes the next row of the iterator and copies the columns of the
  798. // current row into the values pointed at by dest. Use nil as a dest value
  799. // to skip the corresponding column. Scan might send additional queries
  800. // to the database to retrieve the next set of rows if paging was enabled.
  801. //
  802. // Scan returns true if the row was successfully unmarshaled or false if the
  803. // end of the result set was reached or if an error occurred. Close should
  804. // be called afterwards to retrieve any potential errors.
  805. func (iter *Iter) Scan(dest ...interface{}) bool {
  806. if iter.err != nil {
  807. return false
  808. }
  809. if iter.pos >= iter.numRows {
  810. if iter.next != nil {
  811. *iter = *iter.next.fetch()
  812. return iter.Scan(dest...)
  813. }
  814. return false
  815. }
  816. if iter.next != nil && iter.pos == iter.next.pos {
  817. go iter.next.fetch()
  818. }
  819. // currently only support scanning into an expand tuple, such that its the same
  820. // as scanning in more values from a single column
  821. if len(dest) != iter.meta.actualColCount {
  822. iter.err = fmt.Errorf("gocql: not enough columns to scan into: have %d want %d", len(dest), iter.meta.actualColCount)
  823. return false
  824. }
  825. // i is the current position in dest, could posible replace it and just use
  826. // slices of dest
  827. i := 0
  828. for c := range iter.meta.columns {
  829. col := &iter.meta.columns[c]
  830. colBytes, err := iter.readColumn()
  831. if err != nil {
  832. iter.err = err
  833. return false
  834. }
  835. if dest[i] == nil {
  836. i++
  837. continue
  838. }
  839. switch col.TypeInfo.Type() {
  840. case TypeTuple:
  841. // this will panic, actually a bug, please report
  842. tuple := col.TypeInfo.(TupleTypeInfo)
  843. count := len(tuple.Elems)
  844. // here we pass in a slice of the struct which has the number number of
  845. // values as elements in the tuple
  846. iter.err = Unmarshal(col.TypeInfo, colBytes, dest[i:i+count])
  847. i += count
  848. default:
  849. iter.err = Unmarshal(col.TypeInfo, colBytes, dest[i])
  850. i++
  851. }
  852. if iter.err != nil {
  853. return false
  854. }
  855. }
  856. iter.pos++
  857. return true
  858. }
  859. // GetCustomPayload returns any parsed custom payload results if given in the
  860. // response from Cassandra. Note that the result is not a copy.
  861. //
  862. // This additional feature of CQL Protocol v4
  863. // allows additional results and query information to be returned by
  864. // custom QueryHandlers running in your C* cluster.
  865. // See https://datastax.github.io/java-driver/manual/custom_payloads/
  866. func (iter *Iter) GetCustomPayload() map[string][]byte {
  867. return iter.framer.header.customPayload
  868. }
  869. // Close closes the iterator and returns any errors that happened during
  870. // the query or the iteration.
  871. func (iter *Iter) Close() error {
  872. if atomic.CompareAndSwapInt32(&iter.closed, 0, 1) {
  873. if iter.framer != nil {
  874. framerPool.Put(iter.framer)
  875. iter.framer = nil
  876. }
  877. }
  878. return iter.err
  879. }
  880. // WillSwitchPage detects if iterator reached end of current page
  881. // and the next page is available.
  882. func (iter *Iter) WillSwitchPage() bool {
  883. return iter.pos >= iter.numRows && iter.next != nil
  884. }
  885. // checkErrAndNotFound handle error and NotFound in one method.
  886. func (iter *Iter) checkErrAndNotFound() error {
  887. if iter.err != nil {
  888. return iter.err
  889. } else if iter.numRows == 0 {
  890. return ErrNotFound
  891. }
  892. return nil
  893. }
  894. // PageState return the current paging state for a query which can be used for
  895. // subsequent quries to resume paging this point.
  896. func (iter *Iter) PageState() []byte {
  897. return iter.meta.pagingState
  898. }
  899. type nextIter struct {
  900. qry Query
  901. pos int
  902. once sync.Once
  903. next *Iter
  904. }
  905. func (n *nextIter) fetch() *Iter {
  906. n.once.Do(func() {
  907. n.next = n.qry.session.executeQuery(&n.qry)
  908. })
  909. return n.next
  910. }
  911. type Batch struct {
  912. Type BatchType
  913. Entries []BatchEntry
  914. Cons Consistency
  915. rt RetryPolicy
  916. attempts int
  917. totalLatency int64
  918. serialCons SerialConsistency
  919. defaultTimestamp bool
  920. }
  921. // NewBatch creates a new batch operation without defaults from the cluster
  922. func NewBatch(typ BatchType) *Batch {
  923. return &Batch{Type: typ}
  924. }
  925. // NewBatch creates a new batch operation using defaults defined in the cluster
  926. func (s *Session) NewBatch(typ BatchType) *Batch {
  927. s.mu.RLock()
  928. batch := &Batch{Type: typ, rt: s.cfg.RetryPolicy, serialCons: s.cfg.SerialConsistency,
  929. Cons: s.cons, defaultTimestamp: s.cfg.DefaultTimestamp}
  930. s.mu.RUnlock()
  931. return batch
  932. }
  933. // Attempts returns the number of attempts made to execute the batch.
  934. func (b *Batch) Attempts() int {
  935. return b.attempts
  936. }
  937. //Latency returns the average number of nanoseconds to execute a single attempt of the batch.
  938. func (b *Batch) Latency() int64 {
  939. if b.attempts > 0 {
  940. return b.totalLatency / int64(b.attempts)
  941. }
  942. return 0
  943. }
  944. // GetConsistency returns the currently configured consistency level for the batch
  945. // operation.
  946. func (b *Batch) GetConsistency() Consistency {
  947. return b.Cons
  948. }
  949. // Query adds the query to the batch operation
  950. func (b *Batch) Query(stmt string, args ...interface{}) {
  951. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, Args: args})
  952. }
  953. // Bind adds the query to the batch operation and correlates it with a binding callback
  954. // that will be invoked when the batch is executed. The binding callback allows the application
  955. // to define which query argument values will be marshalled as part of the batch execution.
  956. func (b *Batch) Bind(stmt string, bind func(q *QueryInfo) ([]interface{}, error)) {
  957. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, binding: bind})
  958. }
  959. func (b *Batch) retryPolicy() RetryPolicy {
  960. return b.rt
  961. }
  962. // RetryPolicy sets the retry policy to use when executing the batch operation
  963. func (b *Batch) RetryPolicy(r RetryPolicy) *Batch {
  964. b.rt = r
  965. return b
  966. }
  967. // Size returns the number of batch statements to be executed by the batch operation.
  968. func (b *Batch) Size() int {
  969. return len(b.Entries)
  970. }
  971. // SerialConsistency sets the consistency level for the
  972. // serial phase of conditional updates. That consistency can only be
  973. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  974. // SERIAL. This option will be ignored for anything else that a
  975. // conditional update/insert.
  976. //
  977. // Only available for protocol 3 and above
  978. func (b *Batch) SerialConsistency(cons SerialConsistency) *Batch {
  979. b.serialCons = cons
  980. return b
  981. }
  982. // DefaultTimestamp will enable the with default timestamp flag on the query.
  983. // If enable, this will replace the server side assigned
  984. // timestamp as default timestamp. Note that a timestamp in the query itself
  985. // will still override this timestamp. This is entirely optional.
  986. //
  987. // Only available on protocol >= 3
  988. func (b *Batch) DefaultTimestamp(enable bool) *Batch {
  989. b.defaultTimestamp = enable
  990. return b
  991. }
  992. func (b *Batch) attempt(d time.Duration) {
  993. b.attempts++
  994. b.totalLatency += d.Nanoseconds()
  995. // TODO: track latencies per host and things as well instead of just total
  996. }
  997. func (b *Batch) GetRoutingKey() ([]byte, error) {
  998. // TODO: use the first statement in the batch as the routing key?
  999. return nil, nil
  1000. }
  1001. type BatchType byte
  1002. const (
  1003. LoggedBatch BatchType = 0
  1004. UnloggedBatch BatchType = 1
  1005. CounterBatch BatchType = 2
  1006. )
  1007. type BatchEntry struct {
  1008. Stmt string
  1009. Args []interface{}
  1010. binding func(q *QueryInfo) ([]interface{}, error)
  1011. }
  1012. type ColumnInfo struct {
  1013. Keyspace string
  1014. Table string
  1015. Name string
  1016. TypeInfo TypeInfo
  1017. }
  1018. func (c ColumnInfo) String() string {
  1019. return fmt.Sprintf("[column keyspace=%s table=%s name=%s type=%v]", c.Keyspace, c.Table, c.Name, c.TypeInfo)
  1020. }
  1021. // routing key indexes LRU cache
  1022. type routingKeyInfoLRU struct {
  1023. lru *lru.Cache
  1024. mu sync.Mutex
  1025. }
  1026. type routingKeyInfo struct {
  1027. indexes []int
  1028. types []TypeInfo
  1029. }
  1030. func (r *routingKeyInfo) String() string {
  1031. return fmt.Sprintf("routing key index=%v types=%v", r.indexes, r.types)
  1032. }
  1033. func (r *routingKeyInfoLRU) Remove(key string) {
  1034. r.mu.Lock()
  1035. r.lru.Remove(key)
  1036. r.mu.Unlock()
  1037. }
  1038. //Max adjusts the maximum size of the cache and cleans up the oldest records if
  1039. //the new max is lower than the previous value. Not concurrency safe.
  1040. func (r *routingKeyInfoLRU) Max(max int) {
  1041. r.mu.Lock()
  1042. for r.lru.Len() > max {
  1043. r.lru.RemoveOldest()
  1044. }
  1045. r.lru.MaxEntries = max
  1046. r.mu.Unlock()
  1047. }
  1048. type inflightCachedEntry struct {
  1049. wg sync.WaitGroup
  1050. err error
  1051. value interface{}
  1052. }
  1053. // Tracer is the interface implemented by query tracers. Tracers have the
  1054. // ability to obtain a detailed event log of all events that happened during
  1055. // the execution of a query from Cassandra. Gathering this information might
  1056. // be essential for debugging and optimizing queries, but this feature should
  1057. // not be used on production systems with very high load.
  1058. type Tracer interface {
  1059. Trace(traceId []byte)
  1060. }
  1061. type traceWriter struct {
  1062. session *Session
  1063. w io.Writer
  1064. mu sync.Mutex
  1065. }
  1066. // NewTraceWriter returns a simple Tracer implementation that outputs
  1067. // the event log in a textual format.
  1068. func NewTraceWriter(session *Session, w io.Writer) Tracer {
  1069. return &traceWriter{session: session, w: w}
  1070. }
  1071. func (t *traceWriter) Trace(traceId []byte) {
  1072. var (
  1073. coordinator string
  1074. duration int
  1075. )
  1076. iter := t.session.control.query(`SELECT coordinator, duration
  1077. FROM system_traces.sessions
  1078. WHERE session_id = ?`, traceId)
  1079. iter.Scan(&coordinator, &duration)
  1080. if err := iter.Close(); err != nil {
  1081. t.mu.Lock()
  1082. fmt.Fprintln(t.w, "Error:", err)
  1083. t.mu.Unlock()
  1084. return
  1085. }
  1086. var (
  1087. timestamp time.Time
  1088. activity string
  1089. source string
  1090. elapsed int
  1091. )
  1092. fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
  1093. traceId, coordinator, time.Duration(duration)*time.Microsecond)
  1094. t.mu.Lock()
  1095. defer t.mu.Unlock()
  1096. iter = t.session.control.query(`SELECT event_id, activity, source, source_elapsed
  1097. FROM system_traces.events
  1098. WHERE session_id = ?`, traceId)
  1099. for iter.Scan(&timestamp, &activity, &source, &elapsed) {
  1100. fmt.Fprintf(t.w, "%s: %s (source: %s, elapsed: %d)\n",
  1101. timestamp.Format("2006/01/02 15:04:05.999999"), activity, source, elapsed)
  1102. }
  1103. if err := iter.Close(); err != nil {
  1104. fmt.Fprintln(t.w, "Error:", err)
  1105. }
  1106. }
  1107. type Error struct {
  1108. Code int
  1109. Message string
  1110. }
  1111. func (e Error) Error() string {
  1112. return e.Message
  1113. }
  1114. var (
  1115. ErrNotFound = errors.New("not found")
  1116. ErrUnavailable = errors.New("unavailable")
  1117. ErrUnsupported = errors.New("feature not supported")
  1118. ErrTooManyStmts = errors.New("too many statements")
  1119. ErrUseStmt = errors.New("use statements aren't supported. Please see https://github.com/gocql/gocql for explaination.")
  1120. ErrSessionClosed = errors.New("session has been closed")
  1121. ErrNoConnections = errors.New("qocql: no hosts available in the pool")
  1122. ErrNoKeyspace = errors.New("no keyspace provided")
  1123. ErrNoMetadata = errors.New("no metadata available")
  1124. )
  1125. type ErrProtocol struct{ error }
  1126. func NewErrProtocol(format string, args ...interface{}) error {
  1127. return ErrProtocol{fmt.Errorf(format, args...)}
  1128. }
  1129. // BatchSizeMaximum is the maximum number of statements a batch operation can have.
  1130. // This limit is set by cassandra and could change in the future.
  1131. const BatchSizeMaximum = 65535