tx.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857
  1. package storage
  2. import (
  3. "fmt"
  4. "sort"
  5. "strings"
  6. "sync"
  7. )
  8. // ErrSerialization is returned when a transaction's optimistic validation fails
  9. // because a concurrent transaction committed a conflicting change.
  10. var ErrSerialization = fmt.Errorf("serialization failure: concurrent transaction modified the database")
  11. // Session is a buffered SQL transaction/session wrapper around a TableManager.
  12. // While in a transaction, writes are staged in memory, reads observe that
  13. // staged overlay, and COMMIT issues a single atomic compare-and-swap batch
  14. // write. Rollback simply discards the staged changes (no compensating durable
  15. // writes). A Session is used for both autocommit statements (where operations
  16. // delegate straight to the durable TableManager) and buffered transactions.
  17. type Session struct {
  18. schema *SchemaManager
  19. table *TableManager
  20. mu sync.Mutex
  21. inTx bool
  22. aborted bool
  23. // overlay holds staged writes keyed by lowercased table then data key.
  24. overlay map[string]map[string]*overlayEntry
  25. // log records mutations in order so savepoints can roll back.
  26. log []txMutation
  27. // reads records the durable LSN of every key the transaction observed
  28. // (0 means the key was observed absent). It becomes the compare set at
  29. // commit and also captures the base LSN of every written key.
  30. reads map[string]uint64
  31. // scanGens records the per-table generation captured by the first scan of
  32. // each table, validated at commit to detect phantoms. Indexed equality reads
  33. // use predicateGens so writes to other index values do not cause conflicts.
  34. scanGens map[string]uint64
  35. predicateGens map[string]predicateRead
  36. }
  37. type predicateRead struct {
  38. table string
  39. gen uint64
  40. }
  41. type overlayEntry struct {
  42. row Row
  43. absent bool
  44. }
  45. type txMutation struct {
  46. table string
  47. key string
  48. prev *overlayEntry
  49. }
  50. // NewSession creates a session wrapping the given schema and table managers.
  51. func NewSession(schema *SchemaManager, table *TableManager) *Session {
  52. return &Session{
  53. schema: schema,
  54. table: table,
  55. overlay: make(map[string]map[string]*overlayEntry),
  56. reads: make(map[string]uint64),
  57. scanGens: make(map[string]uint64),
  58. predicateGens: make(map[string]predicateRead),
  59. }
  60. }
  61. // Begin starts a buffered transaction.
  62. func (s *Session) Begin() error {
  63. s.mu.Lock()
  64. defer s.mu.Unlock()
  65. if s.inTx {
  66. return fmt.Errorf("cannot start a transaction within a transaction")
  67. }
  68. s.inTx = true
  69. s.aborted = false
  70. s.overlay = make(map[string]map[string]*overlayEntry)
  71. s.log = nil
  72. s.reads = make(map[string]uint64)
  73. s.scanGens = make(map[string]uint64)
  74. s.predicateGens = make(map[string]predicateRead)
  75. return nil
  76. }
  77. // InTx reports whether a transaction is in progress.
  78. func (s *Session) InTx() bool {
  79. s.mu.Lock()
  80. defer s.mu.Unlock()
  81. return s.inTx
  82. }
  83. // Abort marks the transaction as aborted without discarding state.
  84. func (s *Session) Abort() {
  85. s.mu.Lock()
  86. defer s.mu.Unlock()
  87. if s.inTx {
  88. s.aborted = true
  89. }
  90. }
  91. // Snapshot returns the current mutation-log position for a savepoint.
  92. func (s *Session) Snapshot() int {
  93. s.mu.Lock()
  94. defer s.mu.Unlock()
  95. return len(s.log)
  96. }
  97. // RollbackTo discards mutations after the given savepoint position.
  98. func (s *Session) RollbackTo(pos int) {
  99. s.mu.Lock()
  100. defer s.mu.Unlock()
  101. for i := len(s.log) - 1; i >= pos && i >= 0; i-- {
  102. m := s.log[i]
  103. if m.prev == nil {
  104. delete(s.overlay[m.table], m.key)
  105. if len(s.overlay[m.table]) == 0 {
  106. delete(s.overlay, m.table)
  107. }
  108. } else {
  109. if s.overlay[m.table] == nil {
  110. s.overlay[m.table] = make(map[string]*overlayEntry)
  111. }
  112. s.overlay[m.table][m.key] = m.prev
  113. }
  114. }
  115. if pos < len(s.log) {
  116. s.log = s.log[:pos]
  117. }
  118. s.aborted = false
  119. }
  120. // Rollback discards the transaction without writing anything durable.
  121. func (s *Session) Rollback() error {
  122. s.mu.Lock()
  123. defer s.mu.Unlock()
  124. if !s.inTx {
  125. return fmt.Errorf("cannot rollback: no transaction in progress")
  126. }
  127. s.resetLocked()
  128. return nil
  129. }
  130. // Commit validates and durably applies the staged transaction in one atomic
  131. // compare-and-swap batch write.
  132. func (s *Session) Commit() error {
  133. s.mu.Lock()
  134. defer s.mu.Unlock()
  135. if !s.inTx {
  136. return fmt.Errorf("cannot commit: no transaction in progress")
  137. }
  138. if s.aborted {
  139. s.resetLocked()
  140. return fmt.Errorf("current transaction is aborted")
  141. }
  142. // Acquire the per-table gates, escalating any table whose UNIQUE index
  143. // appeared after the initial unlocked probe. See acquireCommitLocks.
  144. _, _, unlockAll, uniqueTables, err := s.acquireCommitLocks()
  145. if err != nil {
  146. s.resetLocked()
  147. return err
  148. }
  149. defer unlockAll()
  150. // Validate scan generations for phantom detection.
  151. for t, gen := range s.scanGens {
  152. if s.table.generation(t) != gen {
  153. s.resetLocked()
  154. return ErrSerialization
  155. }
  156. }
  157. for key, predicate := range s.predicateGens {
  158. if s.table.predicateGeneration(key) != predicate.gen {
  159. s.resetLocked()
  160. return ErrSerialization
  161. }
  162. }
  163. // Validate the final overlay of every unique-indexed written table against
  164. // durable rows. This runs under the exclusive gate acquired above, so it
  165. // cannot race a concurrent writer. Duplicate unique values fail the commit
  166. // with a "UNIQUE constraint failed" error rather than ErrSerialization.
  167. for t := range uniqueTables {
  168. entries := s.overlay[t]
  169. pending := make([]Row, 0, len(entries))
  170. // Exclude every overlay data key: rows being written, updated, or deleted
  171. // are all replaced by this transaction, so their durable unique values
  172. // must not self-conflict with the staged state (e.g. delete a row and
  173. // re-insert the same primary key with a different rowid).
  174. excluded := make(map[string]bool, len(entries))
  175. for key, e := range entries {
  176. excluded[key] = true
  177. if !e.absent {
  178. pending = append(pending, e.row)
  179. }
  180. }
  181. if err := s.table.validateUniqueRows(t, pending, excluded); err != nil {
  182. s.resetLocked()
  183. return err
  184. }
  185. }
  186. // Build the compare set from every observed key.
  187. checks := make([]CompareCheck, 0, len(s.reads))
  188. for key, lsn := range s.reads {
  189. checks = append(checks, CompareCheck{Key: []byte(key), LSN: lsn})
  190. }
  191. sort.Slice(checks, func(i, j int) bool { return string(checks[i].Key) < string(checks[j].Key) })
  192. // Build the batch ops from the staged overlay.
  193. ops := make([]BatchOp, 0)
  194. for _, entries := range s.overlay {
  195. for key, e := range entries {
  196. if e.absent {
  197. ops = append(ops, BatchOp{Op: batchDelete, Key: []byte(key)})
  198. } else {
  199. data, err := encodeRow(e.row)
  200. if err != nil {
  201. s.resetLocked()
  202. return err
  203. }
  204. ops = append(ops, BatchOp{Op: batchPut, Key: []byte(key), Value: data})
  205. }
  206. }
  207. }
  208. sort.Slice(ops, func(i, j int) bool { return string(ops[i].Key) < string(ops[j].Key) })
  209. // A read-only transaction has nothing to write; commit trivially.
  210. if len(ops) == 0 {
  211. s.resetLocked()
  212. return nil
  213. }
  214. var committed bool
  215. err = s.table.pool.WithClient(func(c *KVClient) error {
  216. _, ok, err := c.CompareBatchWrite(checks, ops, nil)
  217. committed = ok
  218. return err
  219. })
  220. if err != nil {
  221. s.resetLocked()
  222. return err
  223. }
  224. if !committed {
  225. s.resetLocked()
  226. return ErrSerialization
  227. }
  228. // Advance generations and invalidate derived caches for written tables.
  229. for t := range s.overlay {
  230. s.table.InvalidateCache(t)
  231. s.table.bumpIndexPredicateWildcard(t)
  232. s.table.bumpGeneration(t)
  233. }
  234. s.resetLocked()
  235. return nil
  236. }
  237. // acquireCommitLocks acquires the per-table gates needed to commit the staged
  238. // transaction, returning the held locks, whether each is exclusive, an unlock
  239. // function, and the set of written tables that carry a UNIQUE index. A written
  240. // table with a UNIQUE index is locked exclusively so its final overlay can be
  241. // validated against durable rows. The uniqueness probe is re-checked under the
  242. // held gates so a concurrent CREATE UNIQUE INDEX cannot slip in after the probe
  243. // and leave a duplicate unvalidated.
  244. func (s *Session) acquireCommitLocks() (locks []*sync.RWMutex, exclusive []bool, unlockAll func(), uniqueTables map[string]bool, err error) {
  245. unlock := func(ls []*sync.RWMutex, ex []bool) {
  246. for i := len(ls) - 1; i >= 0; i-- {
  247. if ex[i] {
  248. ls[i].Unlock()
  249. } else {
  250. ls[i].RUnlock()
  251. }
  252. }
  253. }
  254. for {
  255. // Tables read through scans or predicates need an exclusive validation
  256. // gate. Tables that are only written use the shared publication gate,
  257. // allowing disjoint optimistic commits to proceed concurrently while
  258. // still excluding scanner commits.
  259. affected := make(map[string]bool)
  260. for t := range s.overlay {
  261. affected[t] = false
  262. }
  263. for t := range s.scanGens {
  264. affected[t] = true
  265. }
  266. for _, predicate := range s.predicateGens {
  267. affected[predicate.table] = true
  268. }
  269. // Tables written in this transaction that carry a UNIQUE index need the
  270. // exclusive gate so their final overlay can be validated against durable
  271. // rows without racing another writer.
  272. uniqueTables = make(map[string]bool)
  273. for t := range s.overlay {
  274. uniq, err := s.table.hasUniqueIndex(t)
  275. if err != nil {
  276. return nil, nil, nil, nil, err
  277. }
  278. if uniq {
  279. uniqueTables[t] = true
  280. affected[t] = true
  281. }
  282. }
  283. tables := make([]string, 0, len(affected))
  284. for t := range affected {
  285. tables = append(tables, t)
  286. }
  287. sort.Strings(tables)
  288. locks = make([]*sync.RWMutex, len(tables))
  289. exclusive = make([]bool, len(tables))
  290. for i, t := range tables {
  291. locks[i] = s.table.tableLock(t)
  292. exclusive[i] = affected[t]
  293. }
  294. for i, lock := range locks {
  295. if exclusive[i] {
  296. lock.Lock()
  297. } else {
  298. lock.RLock()
  299. }
  300. }
  301. // Re-check under the held gates. A concurrent CREATE UNIQUE INDEX can
  302. // only have completed before we acquired the gate (it needs the exclusive
  303. // gate), so this probe is authoritative; escalate and retry if a written
  304. // table gained a unique index.
  305. retry := false
  306. for i, t := range tables {
  307. if exclusive[i] {
  308. continue
  309. }
  310. if _, written := s.overlay[t]; !written {
  311. continue
  312. }
  313. uniq, err := s.table.hasUniqueIndex(t)
  314. if err != nil {
  315. unlock(locks, exclusive)
  316. return nil, nil, nil, nil, err
  317. }
  318. if uniq {
  319. retry = true
  320. break
  321. }
  322. }
  323. if retry {
  324. unlock(locks, exclusive)
  325. continue
  326. }
  327. unlockAll = func() { unlock(locks, exclusive) }
  328. return locks, exclusive, unlockAll, uniqueTables, nil
  329. }
  330. }
  331. func (s *Session) resetLocked() {
  332. s.inTx = false
  333. s.aborted = false
  334. s.overlay = make(map[string]map[string]*overlayEntry)
  335. s.log = nil
  336. s.reads = make(map[string]uint64)
  337. s.scanGens = make(map[string]uint64)
  338. s.predicateGens = make(map[string]predicateRead)
  339. }
  340. func (s *Session) stagePut(table, key string, row Row) {
  341. tl := strings.ToLower(table)
  342. if s.overlay[tl] == nil {
  343. s.overlay[tl] = make(map[string]*overlayEntry)
  344. }
  345. s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
  346. s.overlay[tl][key] = &overlayEntry{row: cloneRow(row)}
  347. if _, ok := s.reads[key]; !ok {
  348. s.reads[key] = 0
  349. }
  350. }
  351. func (s *Session) stageDelete(table, key string, row Row) {
  352. tl := strings.ToLower(table)
  353. if s.overlay[tl] == nil {
  354. s.overlay[tl] = make(map[string]*overlayEntry)
  355. }
  356. s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
  357. // Keep the deleted row so COMMIT can exempt its rowid from the unique-index
  358. // scan and release its unique value.
  359. s.overlay[tl][key] = &overlayEntry{row: cloneRow(row), absent: true}
  360. if _, ok := s.reads[key]; !ok {
  361. s.reads[key] = 0
  362. }
  363. }
  364. // GetByPK retrieves a row by primary key, observing the staged overlay in a
  365. // transaction and recording the observed LSN for validation.
  366. func (s *Session) GetByPK(table, pk string) (Row, error) {
  367. s.mu.Lock()
  368. defer s.mu.Unlock()
  369. if !s.inTx {
  370. return s.table.GetByPK(table, pk)
  371. }
  372. key := s.table.dataKey(table, pk)
  373. tl := strings.ToLower(table)
  374. if e, ok := s.overlay[tl][key]; ok {
  375. if e.absent {
  376. return nil, ErrKeyNotFound
  377. }
  378. return cloneRow(e.row), nil
  379. }
  380. row, lsn, err := s.table.getByPKWithLSN(table, pk)
  381. if err != nil {
  382. if err == ErrKeyNotFound {
  383. s.reads[key] = 0
  384. return nil, ErrKeyNotFound
  385. }
  386. return nil, err
  387. }
  388. s.reads[key] = lsn
  389. return row, nil
  390. }
  391. // Select scans a table, merging the staged overlay so a transaction sees its
  392. // own writes, and records per-row LSNs plus the table generation.
  393. func (s *Session) Select(table string, filter func(Row) bool) ([]Row, error) {
  394. s.mu.Lock()
  395. if !s.inTx {
  396. s.mu.Unlock()
  397. return s.table.Select(table, filter)
  398. }
  399. defer s.mu.Unlock()
  400. return s.selectLocked(table, filter)
  401. }
  402. func (s *Session) selectLocked(table string, filter func(Row) bool) ([]Row, error) {
  403. schema, err := s.schema.GetSchema(table)
  404. if err != nil {
  405. return nil, err
  406. }
  407. tl := strings.ToLower(table)
  408. tableLock := s.table.tableLock(tl)
  409. tableLock.RLock()
  410. defer tableLock.RUnlock()
  411. if _, ok := s.scanGens[tl]; !ok {
  412. s.scanGens[tl] = s.table.generation(tl)
  413. }
  414. overlay := s.overlay[tl]
  415. var rows []Row
  416. err = s.table.scanRowsWithLSN(table, func(row Row, lsn uint64) (bool, error) {
  417. key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  418. if _, ok := overlay[key]; ok {
  419. return false, nil
  420. }
  421. s.reads[key] = lsn
  422. if filter == nil || filter(row) {
  423. rows = append(rows, row)
  424. }
  425. return false, nil
  426. })
  427. if err != nil {
  428. return nil, err
  429. }
  430. for _, e := range overlay {
  431. if e.absent {
  432. continue
  433. }
  434. if filter == nil || filter(e.row) {
  435. rows = append(rows, cloneRow(e.row))
  436. }
  437. }
  438. return rows, nil
  439. }
  440. // SelectByIndex reads only matching durable rows while capturing their LSNs,
  441. // then merges the transaction overlay. The table generation protects against
  442. // matching rows being inserted or removed after the lookup.
  443. func (s *Session) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
  444. s.mu.Lock()
  445. defer s.mu.Unlock()
  446. if !s.inTx {
  447. return s.table.SelectByIndex(table, indexName, colValue)
  448. }
  449. tableKey := strings.ToLower(table)
  450. index, err := s.schema.GetIndex(indexName)
  451. if err != nil {
  452. return nil, err
  453. }
  454. col := index.Columns[0].Name
  455. want := formatIndexValue(colValue)
  456. versions, predicate, err := s.table.selectByIndexWithLSN(table, indexName, colValue)
  457. if err != nil {
  458. return nil, err
  459. }
  460. if _, ok := s.predicateGens[predicate.valueKey]; !ok {
  461. s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
  462. }
  463. if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
  464. s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
  465. }
  466. overlay := s.overlay[tableKey]
  467. rows := make([]Row, 0, len(versions)+len(overlay))
  468. for _, version := range versions {
  469. if _, staged := overlay[version.key]; staged {
  470. continue
  471. }
  472. s.reads[version.key] = version.lsn
  473. rows = append(rows, version.row)
  474. }
  475. for _, entry := range overlay {
  476. if !entry.absent && formatIndexValue(entry.row[col]) == want {
  477. rows = append(rows, cloneRow(entry.row))
  478. }
  479. }
  480. return rows, nil
  481. }
  482. // SelectByIndexKey retrieves the rows whose index key equals valueKey, merging
  483. // the staged overlay so a transaction sees its own writes. valueKey is the
  484. // formatted key from TableManager.IndexRowKey, which supports plain, composite,
  485. // and expression indexes uniformly.
  486. func (s *Session) SelectByIndexKey(table string, idx *Index, valueKey string) ([]Row, error) {
  487. s.mu.Lock()
  488. defer s.mu.Unlock()
  489. if !s.inTx {
  490. versions, _, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
  491. if err != nil {
  492. return nil, err
  493. }
  494. rows := make([]Row, len(versions))
  495. for i := range versions {
  496. rows[i] = versions[i].row
  497. }
  498. return rows, nil
  499. }
  500. tableKey := strings.ToLower(table)
  501. versions, predicate, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
  502. if err != nil {
  503. return nil, err
  504. }
  505. if _, ok := s.predicateGens[predicate.valueKey]; !ok {
  506. s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
  507. }
  508. if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
  509. s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
  510. }
  511. overlay := s.overlay[tableKey]
  512. rows := make([]Row, 0, len(versions)+len(overlay))
  513. for _, version := range versions {
  514. if _, staged := overlay[version.key]; staged {
  515. continue
  516. }
  517. s.reads[version.key] = version.lsn
  518. rows = append(rows, version.row)
  519. }
  520. for _, entry := range overlay {
  521. if entry.absent {
  522. continue
  523. }
  524. key, kerr := s.table.IndexRowKey(idx, entry.row)
  525. if kerr != nil {
  526. return nil, kerr
  527. }
  528. if key == valueKey {
  529. rows = append(rows, cloneRow(entry.row))
  530. }
  531. }
  532. return rows, nil
  533. }
  534. // CountFast returns the exact row count, observing the staged overlay in a
  535. // transaction.
  536. func (s *Session) CountFast(table string) (int, error) {
  537. s.mu.Lock()
  538. defer s.mu.Unlock()
  539. if !s.inTx {
  540. return s.table.CountFast(table)
  541. }
  542. rows, err := s.selectLocked(table, nil)
  543. if err != nil {
  544. return 0, err
  545. }
  546. return len(rows), nil
  547. }
  548. // Insert stages an insert in a transaction, or performs a durable autocommit
  549. // insert otherwise.
  550. func (s *Session) Insert(table string, row Row) error {
  551. _, err := s.InsertWithRowID(table, row)
  552. return err
  553. }
  554. // InsertWithRowID stages an insert in a transaction, or performs a durable
  555. // autocommit insert otherwise, returning the actual generated ROWID.
  556. func (s *Session) InsertWithRowID(table string, row Row) (int64, error) {
  557. s.mu.Lock()
  558. defer s.mu.Unlock()
  559. if !s.inTx {
  560. return s.table.InsertWithRowID(table, row)
  561. }
  562. return s.insertLocked(table, row)
  563. }
  564. // InsertBulk stages or durably bulk-inserts multiple rows.
  565. func (s *Session) InsertBulk(table string, rows []Row) (int, error) {
  566. count, _, err := s.InsertBulkWithLastRowID(table, rows)
  567. return count, err
  568. }
  569. // InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the
  570. // last staged/persisted row (0 when nothing was inserted).
  571. func (s *Session) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) {
  572. s.mu.Lock()
  573. defer s.mu.Unlock()
  574. if !s.inTx {
  575. return s.table.InsertBulkWithLastRowID(table, rows)
  576. }
  577. var lastRowID int64
  578. count := 0
  579. for _, row := range rows {
  580. rid, err := s.insertLocked(table, row)
  581. if err != nil {
  582. return count, lastRowID, err
  583. }
  584. lastRowID = rid
  585. count++
  586. }
  587. return count, lastRowID, nil
  588. }
  589. // insertLocked is the transaction insert helper (caller holds s.mu).
  590. func (s *Session) insertLocked(table string, row Row) (int64, error) {
  591. nr, key, err := s.table.prepareInsert(table, row)
  592. if err != nil {
  593. return 0, err
  594. }
  595. rowid, _ := rowIDFromRow(nr)
  596. schema, err := s.schema.GetSchema(table)
  597. if err != nil {
  598. return 0, err
  599. }
  600. pk := fmt.Sprintf("%v", nr[schema.PrimaryKey])
  601. tl := strings.ToLower(table)
  602. if e, ok := s.overlay[tl][key]; ok {
  603. if !e.absent {
  604. return 0, fmt.Errorf("duplicate primary key: %s", pk)
  605. }
  606. } else {
  607. _, lsn, err := s.table.getByPKWithLSN(table, pk)
  608. if err == nil {
  609. s.reads[key] = lsn
  610. return 0, fmt.Errorf("duplicate primary key: %s", pk)
  611. }
  612. if err != ErrKeyNotFound {
  613. return 0, err
  614. }
  615. s.reads[key] = 0
  616. }
  617. s.stagePut(table, key, nr)
  618. return rowid, nil
  619. }
  620. // UpdateByPK stages or durably applies a single-row update.
  621. func (s *Session) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) {
  622. s.mu.Lock()
  623. if !s.inTx {
  624. s.mu.Unlock()
  625. return s.table.UpdateByPK(table, pk, updateFn)
  626. }
  627. defer s.mu.Unlock()
  628. schema, err := s.schema.GetSchema(table)
  629. if err != nil {
  630. return nil, false, err
  631. }
  632. key := s.table.dataKey(table, pk)
  633. row, err := s.getByPKLocked(table, pk)
  634. if err == ErrKeyNotFound {
  635. return nil, false, nil
  636. }
  637. if err != nil {
  638. return nil, false, err
  639. }
  640. oldRow := cloneRow(row)
  641. updates, err := s.runUpdateFn(updateFn, row)
  642. if err != nil {
  643. return nil, false, err
  644. }
  645. for name, value := range updates {
  646. for _, column := range schema.Columns {
  647. if strings.EqualFold(name, column.Name) {
  648. row[column.Name] = value
  649. break
  650. }
  651. }
  652. }
  653. s.stagePut(table, key, row)
  654. return oldRow, true, nil
  655. }
  656. // runUpdateFn releases the session lock while invoking the update callback so the
  657. // callback can run nested reads (e.g. a scalar subquery in SET) through the same
  658. // session without deadlocking on s.mu. A connection executes one statement at a
  659. // time, so the staged overlay cannot change while the callback runs. The lock is
  660. // re-acquired before returning.
  661. func (s *Session) runUpdateFn(updateFn func(Row) (Row, error), row Row) (Row, error) {
  662. s.mu.Unlock()
  663. defer s.mu.Lock()
  664. return updateFn(row)
  665. }
  666. // DeleteByPK stages or durably applies a single-row delete.
  667. func (s *Session) DeleteByPK(table, pk string) (Row, bool, error) {
  668. s.mu.Lock()
  669. defer s.mu.Unlock()
  670. if !s.inTx {
  671. return s.table.DeleteByPK(table, pk)
  672. }
  673. key := s.table.dataKey(table, pk)
  674. row, err := s.getByPKLocked(table, pk)
  675. if err == ErrKeyNotFound {
  676. return nil, false, nil
  677. }
  678. if err != nil {
  679. return nil, false, err
  680. }
  681. s.stageDelete(table, key, row)
  682. return row, true, nil
  683. }
  684. // getByPKLocked reads a row observing the overlay (caller holds s.mu).
  685. func (s *Session) getByPKLocked(table, pk string) (Row, error) {
  686. key := s.table.dataKey(table, pk)
  687. tl := strings.ToLower(table)
  688. if e, ok := s.overlay[tl][key]; ok {
  689. if e.absent {
  690. return nil, ErrKeyNotFound
  691. }
  692. return cloneRow(e.row), nil
  693. }
  694. row, lsn, err := s.table.getByPKWithLSN(table, pk)
  695. if err != nil {
  696. if err == ErrKeyNotFound {
  697. s.reads[key] = 0
  698. return nil, ErrKeyNotFound
  699. }
  700. return nil, err
  701. }
  702. s.reads[key] = lsn
  703. return row, nil
  704. }
  705. // UpdateFunc stages or durably applies a scan-based update.
  706. func (s *Session) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
  707. s.mu.Lock()
  708. if !s.inTx {
  709. s.mu.Unlock()
  710. return s.table.UpdateFunc(table, updateFn, filter)
  711. }
  712. defer s.mu.Unlock()
  713. schema, err := s.schema.GetSchema(table)
  714. if err != nil {
  715. return 0, err
  716. }
  717. rows, err := s.selectLocked(table, filter)
  718. if err != nil {
  719. return 0, err
  720. }
  721. count := 0
  722. for _, row := range rows {
  723. oldRow := cloneRow(row)
  724. oldKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  725. updates, err := s.runUpdateFn(updateFn, row)
  726. if err != nil {
  727. return count, err
  728. }
  729. for name, value := range updates {
  730. for _, column := range schema.Columns {
  731. if strings.EqualFold(name, column.Name) {
  732. row[column.Name] = value
  733. break
  734. }
  735. }
  736. }
  737. newKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  738. if newKey != oldKey {
  739. // The primary key changed: move the staged row instead of leaving a
  740. // stale copy under the old key, and refuse to overwrite another row.
  741. if _, gerr := s.getByPKLocked(table, fmt.Sprintf("%v", row[schema.PrimaryKey])); gerr == nil {
  742. return count, fmt.Errorf("duplicate primary key: %v", row[schema.PrimaryKey])
  743. } else if gerr != ErrKeyNotFound {
  744. return count, gerr
  745. }
  746. s.stageDelete(table, oldKey, oldRow)
  747. }
  748. s.stagePut(table, newKey, row)
  749. count++
  750. }
  751. return count, nil
  752. }
  753. // Delete stages or durably applies a scan-based delete.
  754. func (s *Session) Delete(table string, filter func(Row) bool) (int, error) {
  755. s.mu.Lock()
  756. if !s.inTx {
  757. s.mu.Unlock()
  758. return s.table.Delete(table, filter)
  759. }
  760. defer s.mu.Unlock()
  761. schema, err := s.schema.GetSchema(table)
  762. if err != nil {
  763. return 0, err
  764. }
  765. rows, err := s.selectLocked(table, filter)
  766. if err != nil {
  767. return 0, err
  768. }
  769. count := 0
  770. for _, row := range rows {
  771. key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  772. s.stageDelete(table, key, row)
  773. count++
  774. }
  775. return count, nil
  776. }
  777. // ClearIndex passes through to the durable table manager.
  778. func (s *Session) ClearIndex(indexName, tableName string, columns []string) error {
  779. return s.table.ClearIndex(indexName, tableName, columns)
  780. }
  781. // InvalidateCache passes through to the durable table manager.
  782. func (s *Session) InvalidateCache(table string) {
  783. s.table.InvalidateCache(table)
  784. }
  785. // BuildIndex passes through to the durable table manager.
  786. func (s *Session) BuildIndex(indexName, tableName string, columns []string) error {
  787. return s.table.BuildIndex(indexName, tableName, columns)
  788. }