table.go 55 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943
  1. package storage
  2. import (
  3. "errors"
  4. "fmt"
  5. "hash/fnv"
  6. "math"
  7. "strconv"
  8. "strings"
  9. "sync"
  10. "time"
  11. )
  12. // Row represents a database row.
  13. type Row map[string]interface{}
  14. // UniqueViolationError reports a duplicate value on a UNIQUE index. The
  15. // constraint name identifies the violated index so protocol layers can surface a
  16. // machine-detectable error.
  17. type UniqueViolationError struct {
  18. Constraint string
  19. }
  20. func (e *UniqueViolationError) Error() string {
  21. return fmt.Sprintf("UNIQUE constraint failed: %s", e.Constraint)
  22. }
  23. // TableManager manages table data operations.
  24. type TableManager struct {
  25. pool *KVPool
  26. schema *SchemaManager
  27. database string
  28. cacheMu sync.RWMutex
  29. uniqueCache map[string]*uniqueIndexCache
  30. indexRowValues map[string]map[int64]string // index name → rowid → candidate key
  31. keyRowCache map[string]map[string]int64 // table name → primary key → rowid
  32. indexCache map[string]map[string][]int64 // index name → indexed value → rowids
  33. indexTable map[string]string // index name → table name
  34. rowKeyCache map[string]map[int64]string // table name → rowid → primary key
  35. // disabledIndexes prevents a concurrent lookup from rebuilding an index
  36. // after DROP has cleared it but before the schema entry is removed.
  37. disabledIndexes map[string]bool
  38. // counts holds exact per-table row counts for the COUNT(*) fast path.
  39. // It is derived lazily from durable rows on first use and maintained
  40. // incrementally by Insert/InsertBulk/Delete thereafter.
  41. counts map[string]int
  42. countsInit map[string]bool
  43. countGeneration map[string]time.Time
  44. // stripes are deterministic per-key locks used by point operations
  45. // (GetByPK/Insert/UpdateByPK/DeleteByPK). They replace the table-wide lock
  46. // so point operations on different keys of the same table proceed
  47. // concurrently. Index is derived from the full data key (database+table+pk).
  48. stripes [64]sync.Mutex
  49. // generations tracks full-table scans. predicateGenerations narrows indexed
  50. // equality reads to one index value so unrelated writes do not conflict.
  51. genMu sync.Mutex
  52. generations map[string]uint64
  53. predicateGenerations map[string]uint64
  54. // exprMu guards exprEval, the pluggable evaluator for expression-index
  55. // columns. It is set by the SQL executor so storage stays independent of
  56. // the SQL front end.
  57. exprMu sync.RWMutex
  58. exprEval ExpressionEvaluator
  59. }
  60. // ExpressionEvaluator computes the value of a persisted index expression for a
  61. // row. The SQL executor registers one so the storage layer can enforce
  62. // expression indexes without importing the parser/executor.
  63. type ExpressionEvaluator func(expression string, row Row) (interface{}, error)
  64. // SetExpressionEvaluator installs the evaluator used to compute expression-index
  65. // values during uniqueness validation and index maintenance.
  66. func (m *TableManager) SetExpressionEvaluator(eval ExpressionEvaluator) {
  67. m.exprMu.Lock()
  68. m.exprEval = eval
  69. m.exprMu.Unlock()
  70. }
  71. func (m *TableManager) expressionEvaluator() ExpressionEvaluator {
  72. m.exprMu.RLock()
  73. defer m.exprMu.RUnlock()
  74. return m.exprEval
  75. }
  76. // NewTableManager creates a new table manager.
  77. func NewTableManager(pool *KVPool, schema *SchemaManager, database string) *TableManager {
  78. return &TableManager{
  79. pool: pool,
  80. schema: schema,
  81. database: database,
  82. uniqueCache: make(map[string]*uniqueIndexCache),
  83. indexRowValues: make(map[string]map[int64]string),
  84. keyRowCache: make(map[string]map[string]int64),
  85. indexCache: make(map[string]map[string][]int64),
  86. indexTable: make(map[string]string),
  87. rowKeyCache: make(map[string]map[int64]string),
  88. disabledIndexes: make(map[string]bool),
  89. counts: make(map[string]int),
  90. countsInit: make(map[string]bool),
  91. countGeneration: make(map[string]time.Time),
  92. generations: make(map[string]uint64),
  93. predicateGenerations: make(map[string]uint64),
  94. }
  95. }
  96. func (m *TableManager) tableLock(key string) *sync.RWMutex {
  97. return m.schema.tableLock(key)
  98. }
  99. // stripeKey returns the deterministic striped lock for a point operation on the
  100. // given full data key. Point operations on different keys therefore serialize
  101. // independently, while operations on the same key are mutually exclusive.
  102. func (m *TableManager) stripeKey(key string) *sync.Mutex {
  103. h := fnv.New32a()
  104. h.Write([]byte(key))
  105. return &m.stripes[h.Sum32()%uint32(len(m.stripes))]
  106. }
  107. // generation returns the current in-process generation for a table. It is
  108. // bumped on every committed write and captured by transaction scans.
  109. func (m *TableManager) generation(table string) uint64 {
  110. key := strings.ToLower(table)
  111. m.genMu.Lock()
  112. g := m.generations[key]
  113. m.genMu.Unlock()
  114. return g
  115. }
  116. // bumpGeneration advances a table's generation. Callers hold the table gate in
  117. // shared mode for point writes or exclusive mode for scan-based writes.
  118. func (m *TableManager) bumpGeneration(table string) {
  119. key := strings.ToLower(table)
  120. m.genMu.Lock()
  121. m.generations[key]++
  122. m.genMu.Unlock()
  123. }
  124. func indexPredicateKey(table, indexName, value string) string {
  125. return strings.ToLower(table) + "\x00" + strings.ToLower(indexName) + "\x00" + value
  126. }
  127. func indexPredicateWildcardKey(table string) string {
  128. return strings.ToLower(table) + "\x00*"
  129. }
  130. func (m *TableManager) predicateSnapshot(table, indexName, value string) (string, uint64, string, uint64) {
  131. valueKey := indexPredicateKey(table, indexName, value)
  132. wildcardKey := indexPredicateWildcardKey(table)
  133. m.genMu.Lock()
  134. valueGen := m.predicateGenerations[valueKey]
  135. wildcardGen := m.predicateGenerations[wildcardKey]
  136. m.genMu.Unlock()
  137. return valueKey, valueGen, wildcardKey, wildcardGen
  138. }
  139. func (m *TableManager) predicateGeneration(key string) uint64 {
  140. m.genMu.Lock()
  141. gen := m.predicateGenerations[key]
  142. m.genMu.Unlock()
  143. return gen
  144. }
  145. func (m *TableManager) bumpIndexPredicates(table string, rows ...Row) error {
  146. indexes, err := m.schema.ListTableIndexes(table)
  147. if err != nil || len(indexes) == 0 {
  148. return err
  149. }
  150. m.genMu.Lock()
  151. defer m.genMu.Unlock()
  152. for _, row := range rows {
  153. if row == nil {
  154. continue
  155. }
  156. for _, index := range indexes {
  157. value, err := m.buildIndexValue(row, index.Columns)
  158. if err != nil {
  159. return err
  160. }
  161. m.predicateGenerations[indexPredicateKey(table, index.Name, value)]++
  162. }
  163. }
  164. return nil
  165. }
  166. func (m *TableManager) bumpIndexPredicateWildcard(table string) {
  167. m.genMu.Lock()
  168. m.predicateGenerations[indexPredicateWildcardKey(table)]++
  169. m.genMu.Unlock()
  170. }
  171. // compareWritePoint issues a single CompareBatchWrite against the pooled KV.
  172. // It returns whether the compare checks held and the ops committed.
  173. func (m *TableManager) compareWritePoint(checks []CompareCheck, ops []BatchOp) (bool, error) {
  174. var committed bool
  175. err := m.pool.WithClient(func(c *KVClient) error {
  176. _, ok, err := c.CompareBatchWrite(checks, ops, nil)
  177. committed = ok
  178. return err
  179. })
  180. return committed, err
  181. }
  182. // invalidateCache removes a table's derived in-memory indexes.
  183. func (m *TableManager) invalidateCache(table string) {
  184. m.cacheMu.Lock()
  185. key := strings.ToLower(table)
  186. delete(m.rowKeyCache, key)
  187. delete(m.keyRowCache, key)
  188. delete(m.counts, key)
  189. delete(m.countsInit, key)
  190. delete(m.countGeneration, key)
  191. for indexName, tableName := range m.indexTable {
  192. if tableName == key {
  193. delete(m.indexCache, indexName)
  194. delete(m.indexRowValues, indexName)
  195. delete(m.indexTable, indexName)
  196. }
  197. }
  198. m.cacheMu.Unlock()
  199. }
  200. // InvalidateCache is the exported version for use by the executor.
  201. func (m *TableManager) InvalidateCache(table string) {
  202. m.invalidateCache(table)
  203. m.cacheMu.Lock()
  204. for name, cache := range m.uniqueCache {
  205. if cache.table == strings.ToLower(table) {
  206. delete(m.uniqueCache, name)
  207. }
  208. }
  209. m.cacheMu.Unlock()
  210. }
  211. // rowVisitFunc is invoked for each decoded row in a streaming scan. Return
  212. // stop=true to end the scan early; a non-nil error aborts the scan.
  213. type rowVisitFunc func(Row) (stop bool, err error)
  214. // scanRows streams the rows of table by scanning durable KV rows one page at a
  215. // time under a single pooled client. Each page is decoded as it arrives and
  216. // passed to fn, which may stop the scan early. The cursor is always closed and
  217. // the client always returned to the pool, even on error.
  218. func (m *TableManager) scanRows(table string, fn rowVisitFunc) error {
  219. return m.scanRowsWithPageSize(table, scanPageSize, fn)
  220. }
  221. func (m *TableManager) scanRowsWithPageSize(table string, pageSize uint32, fn rowVisitFunc) error {
  222. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  223. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), pageSize)
  224. if err != nil {
  225. return err
  226. }
  227. defer func() {
  228. if err := cursor.Close(); retErr == nil {
  229. retErr = err
  230. }
  231. }()
  232. for {
  233. entries, done, err := cursor.Next()
  234. if err != nil {
  235. return err
  236. }
  237. for _, e := range entries {
  238. row, err := decodeRow(e.Value)
  239. if err != nil {
  240. return err
  241. }
  242. stop, err := fn(row)
  243. if err != nil {
  244. return err
  245. }
  246. if stop {
  247. return nil
  248. }
  249. }
  250. if done {
  251. return nil
  252. }
  253. }
  254. })
  255. }
  256. // rowWithLSNVisitFunc is like rowVisitFunc but also passes the durable row LSN.
  257. type rowWithLSNVisitFunc func(row Row, lsn uint64) (stop bool, err error)
  258. // scanRowsWithLSN streams a table's rows together with their durable LSNs. It
  259. // is used by buffered transactions to capture a per-row read set for
  260. // compare-and-swap validation at commit.
  261. func (m *TableManager) scanRowsWithLSN(table string, fn rowWithLSNVisitFunc) error {
  262. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  263. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), scanPageSize)
  264. if err != nil {
  265. return err
  266. }
  267. defer func() {
  268. if err := cursor.Close(); retErr == nil {
  269. retErr = err
  270. }
  271. }()
  272. for {
  273. entries, done, err := cursor.Next()
  274. if err != nil {
  275. return err
  276. }
  277. for _, e := range entries {
  278. row, err := decodeRow(e.Value)
  279. if err != nil {
  280. return err
  281. }
  282. stop, err := fn(row, e.LSN)
  283. if err != nil {
  284. return err
  285. }
  286. if stop {
  287. return nil
  288. }
  289. }
  290. if done {
  291. return nil
  292. }
  293. }
  294. })
  295. }
  296. // scanCountKeys counts the durable rows of table using a key-only scan so row
  297. // values are never pulled across the wire. It is used for first-time COUNT(*)
  298. // derivation.
  299. func (m *TableManager) scanCountKeys(table string) (int, error) {
  300. count := 0
  301. err := m.pool.WithClient(func(client *KVClient) (retErr error) {
  302. cursor, err := client.ScanKeys([]byte(m.dataPrefix(table)))
  303. if err != nil {
  304. return err
  305. }
  306. defer func() {
  307. if err := cursor.Close(); retErr == nil {
  308. retErr = err
  309. }
  310. }()
  311. for {
  312. entries, done, err := cursor.Next()
  313. if err != nil {
  314. return err
  315. }
  316. count += len(entries)
  317. if done {
  318. return nil
  319. }
  320. }
  321. })
  322. return count, err
  323. }
  324. // CountFast returns the exact number of rows in a table. The count is derived
  325. // from durable rows on first use (recovering across restarts) and then
  326. // maintained incrementally by the write paths, so repeated COUNT(*) queries
  327. // avoid a full table scan. It intentionally does not persist a counter to KV:
  328. // the KV layer has no atomic increment primitive, and a durable counter that
  329. // could diverge from the rows on crash would be worse than a lazily-derived,
  330. // always-exact value. The cost is one key-only scan the first time COUNT(*) is
  331. // issued after startup.
  332. func (m *TableManager) CountFast(table string) (int, error) {
  333. key := strings.ToLower(table)
  334. tableSchema, err := m.schema.GetSchema(table)
  335. if err != nil {
  336. return 0, err
  337. }
  338. // Cached read path: if the count is initialized for the current table
  339. // generation, return it without taking any table lock.
  340. m.cacheMu.Lock()
  341. if m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  342. if m.countsInit[key] {
  343. n := m.counts[key]
  344. m.cacheMu.Unlock()
  345. return n, nil
  346. }
  347. } else {
  348. delete(m.counts, key)
  349. delete(m.countsInit, key)
  350. m.countGeneration[key] = tableSchema.CreatedAt
  351. }
  352. m.cacheMu.Unlock()
  353. // First derivation: hold the table write lock so the key-only scan is
  354. // exact against concurrent writes.
  355. tl := m.tableLock(key)
  356. tl.Lock()
  357. defer tl.Unlock()
  358. m.cacheMu.Lock()
  359. if !m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  360. delete(m.counts, key)
  361. delete(m.countsInit, key)
  362. m.countGeneration[key] = tableSchema.CreatedAt
  363. }
  364. if m.countsInit[key] {
  365. n := m.counts[key]
  366. m.cacheMu.Unlock()
  367. return n, nil
  368. }
  369. m.cacheMu.Unlock()
  370. count, err := m.scanCountKeys(table)
  371. if err != nil {
  372. return 0, err
  373. }
  374. m.cacheMu.Lock()
  375. m.counts[key] = count
  376. m.countsInit[key] = true
  377. m.cacheMu.Unlock()
  378. return count, nil
  379. }
  380. // countInitialized reports whether the derived count cache is initialized for
  381. // the table at this instant. Write paths capture it before their KV write so a
  382. // concurrent first derivation does not double-count the just-written row.
  383. func (m *TableManager) countInitialized(table string) bool {
  384. key := strings.ToLower(table)
  385. m.cacheMu.Lock()
  386. init := m.countsInit[key]
  387. m.cacheMu.Unlock()
  388. return init
  389. }
  390. // incrCount adjusts the derived per-table row count. wasInit reports whether
  391. // the count was already initialized before the corresponding write, so a count
  392. // that was not yet initialized is left to be re-derived from durable rows
  393. // (which already reflect the write) on next use.
  394. func (m *TableManager) incrCount(table string, generation time.Time, delta int, wasInit bool) {
  395. key := strings.ToLower(table)
  396. m.cacheMu.Lock()
  397. if !m.countGeneration[key].Equal(generation) {
  398. delete(m.counts, key)
  399. delete(m.countsInit, key)
  400. m.countGeneration[key] = generation
  401. }
  402. if wasInit && m.countsInit[key] {
  403. m.counts[key] += delta
  404. }
  405. if !wasInit {
  406. // The count was not initialized before the write, so a concurrent first
  407. // derivation may have missed the just-written row. Invalidate to force
  408. // an exact re-derivation on the next COUNT(*).
  409. delete(m.countsInit, key)
  410. }
  411. m.cacheMu.Unlock()
  412. }
  413. // dataKey returns the key for a row.
  414. func (m *TableManager) dataKey(table, pk string) string {
  415. return fmt.Sprintf("%s:_data:%s:%s", m.database, strings.ToLower(table), pk)
  416. }
  417. // dataPrefix returns the prefix for all rows in a table.
  418. func (m *TableManager) dataPrefix(table string) string {
  419. return fmt.Sprintf("%s:_data:%s:", m.database, strings.ToLower(table))
  420. }
  421. // prepareInsert validates and normalizes an insert row, generating the ROWID
  422. // when required. It returns the normalized row and the full data key without
  423. // writing anything, so both the autocommit path (compare-and-swap) and the
  424. // buffered transaction path (staging) can share it. The input row has its
  425. // primary key populated as a side effect.
  426. func (m *TableManager) prepareInsert(table string, row Row) (Row, string, error) {
  427. schema, err := m.schema.GetSchema(table)
  428. if err != nil {
  429. return nil, "", err
  430. }
  431. // Get primary key value
  432. pkValue, ok := row[schema.PrimaryKey]
  433. if !ok {
  434. for k, v := range row {
  435. if strings.EqualFold(k, schema.PrimaryKey) {
  436. pkValue = v
  437. ok = true
  438. break
  439. }
  440. }
  441. }
  442. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  443. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  444. var rowid int64
  445. if !ok || pkValue == nil {
  446. if isIntegerPK || !ok {
  447. rowid, err = m.schema.GetNextRowID(table)
  448. if err != nil {
  449. return nil, "", err
  450. }
  451. pkValue = rowid
  452. row[schema.PrimaryKey] = rowid
  453. ok = true
  454. } else {
  455. return nil, "", fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  456. }
  457. } else if isIntegerPK {
  458. switch v := pkValue.(type) {
  459. case int64:
  460. rowid = v
  461. case float64:
  462. if math.Trunc(v) != v {
  463. return nil, "", fmt.Errorf("invalid integer primary key: %v", v)
  464. }
  465. rowid = int64(v)
  466. case int:
  467. rowid = int64(v)
  468. default:
  469. rowid = 0
  470. }
  471. if rowid > 0 {
  472. if err := m.schema.UpdateMaxRowID(table, rowid); err != nil {
  473. return nil, "", err
  474. }
  475. }
  476. }
  477. pk := fmt.Sprintf("%v", pkValue)
  478. normalizedRow := make(Row)
  479. for _, col := range schema.Columns {
  480. for k, v := range row {
  481. if strings.EqualFold(k, col.Name) {
  482. normalizedRow[col.Name] = v
  483. break
  484. }
  485. }
  486. }
  487. for _, col := range schema.Columns {
  488. if _, ok := normalizedRow[col.Name]; !ok && col.Default != nil {
  489. normalizedRow[col.Name] = col.Default
  490. }
  491. }
  492. if rowid > 0 {
  493. normalizedRow["_rowid_"] = rowid
  494. } else {
  495. newRowID, err := m.schema.GetNextRowID(table)
  496. if err != nil {
  497. return nil, "", err
  498. }
  499. normalizedRow["_rowid_"] = newRowID
  500. }
  501. if err := validateNotNull(schema, normalizedRow); err != nil {
  502. return nil, "", err
  503. }
  504. return normalizedRow, m.dataKey(table, pk), nil
  505. }
  506. // Insert inserts a new row. The duplicate check and write are one atomic
  507. // compare-and-swap so concurrent inserts of the same primary key cannot both
  508. // persist. The shared table gate is acquired before the key's striped lock so
  509. // a queued scan writer cannot invert the lock order with point operations.
  510. func (m *TableManager) Insert(table string, row Row) error {
  511. _, err := m.InsertWithRowID(table, row)
  512. return err
  513. }
  514. // InsertWithRowID returns the actual row identifier, never a concurrent counter.
  515. func (m *TableManager) InsertWithRowID(table string, row Row) (int64, error) {
  516. nr, key, err := m.prepareInsert(table, row)
  517. if err != nil {
  518. return 0, err
  519. }
  520. rowid, _ := rowIDFromRow(nr)
  521. data, err := encodeRow(nr)
  522. if err != nil {
  523. return 0, fmt.Errorf("failed to serialize row: %w", err)
  524. }
  525. wasInit := m.countInitialized(table)
  526. // Point writers share this gate with each other. Transaction commits and
  527. // scan-based writes take it exclusively, so generation validation and cache
  528. // publication are ordered without serializing writes to different keys.
  529. // A UNIQUE index forces the exclusive gate so the validating scan below
  530. // cannot race another writer (or a concurrent CREATE UNIQUE INDEX).
  531. unlock, uniq, err := m.lockForWrite(table)
  532. if err != nil {
  533. return 0, err
  534. }
  535. defer unlock()
  536. var uniquePlan uniqueCachePlan
  537. if uniq {
  538. uniquePlan, err = m.prepareUniqueRows(table, map[string]*overlayEntry{key: {row: nr}})
  539. if err != nil {
  540. return 0, err
  541. }
  542. }
  543. st := m.stripeKey(key)
  544. st.Lock()
  545. defer st.Unlock()
  546. committed, err := m.compareWritePoint(
  547. []CompareCheck{{Key: []byte(key), LSN: 0}},
  548. []BatchOp{{Op: batchPut, Key: []byte(key), Value: data}},
  549. )
  550. if err != nil {
  551. m.invalidateWriteFailure(table)
  552. return 0, err
  553. }
  554. if !committed {
  555. return 0, &PrimaryKeyViolationError{Table: table, Key: strings.TrimPrefix(key, m.dataPrefix(table))}
  556. }
  557. m.applyUniquePlan(uniquePlan, nil)
  558. if err := m.updateIndexesForRow(table, nr, true, key); err != nil {
  559. return 0, err
  560. }
  561. // Publish derived index state before its generations. A reader that races
  562. // with publication either sees the old generation and aborts or sees the
  563. // complete new state.
  564. if err := m.bumpIndexPredicates(table, nr); err != nil {
  565. return 0, err
  566. }
  567. m.bumpGeneration(table)
  568. if schema, serr := m.schema.GetSchema(table); serr == nil {
  569. m.incrCount(table, schema.CreatedAt, 1, wasInit)
  570. }
  571. return rowid, nil
  572. }
  573. // bulkBatchByteBudget bounds a single atomic BATCH_WRITE payload below the
  574. // PKBFI frame limit so a bulk insert never emits a frame the server rejects.
  575. // Each op contributes 12 header bytes plus its key and value.
  576. const bulkBatchByteBudget = 60 * 1024 * 1024
  577. // chunkBatchOps splits ops into atomic BATCH_WRITE chunks bounded by both the
  578. // PKBFI operation-count limit and the frame-size limit. Each chunk is a slice
  579. // of the backing array, valid until the next append to ops.
  580. func chunkBatchOps(ops []BatchOp) [][]BatchOp {
  581. var chunks [][]BatchOp
  582. for i := 0; i < len(ops); {
  583. end := i + maxOperations
  584. if end > len(ops) {
  585. end = len(ops)
  586. }
  587. bytes := 0
  588. j := i
  589. for j < end {
  590. sz := 12 + len(ops[j].Key) + len(ops[j].Value)
  591. if j > i && bytes+sz > bulkBatchByteBudget {
  592. break
  593. }
  594. bytes += sz
  595. j++
  596. }
  597. if j == i {
  598. j = i + 1
  599. }
  600. chunks = append(chunks, ops[i:j])
  601. i = j
  602. }
  603. return chunks
  604. }
  605. // InsertBulk inserts multiple rows efficiently using atomic BATCH_WRITE chunks
  606. // bounded by the PKBFI operation-count and frame-size limits. Skips per-row
  607. // duplicate checks (caller must ensure uniqueness). Used by INSERT ... SELECT.
  608. func (m *TableManager) InsertBulk(table string, rows []Row) (int, error) {
  609. count, _, err := m.InsertBulkWithLastRowID(table, rows)
  610. return count, err
  611. }
  612. // InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the
  613. // last persisted row (0 when nothing persisted).
  614. func (m *TableManager) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) {
  615. if len(rows) == 0 {
  616. return 0, 0, nil
  617. }
  618. tl := m.tableLock(table)
  619. tl.Lock()
  620. defer tl.Unlock()
  621. schema, err := m.schema.GetSchema(table)
  622. if err != nil {
  623. return 0, 0, err
  624. }
  625. wasInit := m.countInitialized(table)
  626. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  627. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  628. // Normalize rows and assign _rowid_.
  629. normalized := make([]Row, 0, len(rows))
  630. var maxRowID int64
  631. for _, row := range rows {
  632. nr := make(Row)
  633. for _, col := range schema.Columns {
  634. for k, v := range row {
  635. if strings.EqualFold(k, col.Name) {
  636. nr[col.Name] = v
  637. break
  638. }
  639. }
  640. }
  641. for _, col := range schema.Columns {
  642. if _, ok := nr[col.Name]; !ok && col.Default != nil {
  643. nr[col.Name] = col.Default
  644. }
  645. }
  646. var rowid int64
  647. var hasRowid bool
  648. if isIntegerPK {
  649. switch v := nr[schema.PrimaryKey].(type) {
  650. case float64:
  651. if math.Trunc(v) != v {
  652. return 0, 0, fmt.Errorf("invalid integer primary key: %v", v)
  653. }
  654. rowid = int64(v)
  655. hasRowid = true
  656. case int64:
  657. rowid = v
  658. hasRowid = true
  659. case int:
  660. rowid = int64(v)
  661. hasRowid = true
  662. }
  663. }
  664. if !hasRowid {
  665. // Auto-generate for an INTEGER primary key or the synthetic _rowid_,
  666. // mirroring prepareInsert so INSERT ... SELECT works on AUTOINCREMENT
  667. // tables.
  668. if isIntegerPK || schema.PrimaryKey == "_rowid_" {
  669. rowid, err = m.schema.GetNextRowID(table)
  670. if err != nil {
  671. return 0, 0, err
  672. }
  673. if isIntegerPK {
  674. nr[schema.PrimaryKey] = rowid
  675. } else {
  676. nr["_rowid_"] = rowid
  677. }
  678. } else {
  679. pk, ok := nr[schema.PrimaryKey]
  680. if !ok || pk == nil {
  681. return 0, 0, fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  682. }
  683. rowid, err = m.schema.GetNextRowID(table)
  684. if err != nil {
  685. return 0, 0, err
  686. }
  687. }
  688. }
  689. nr["_rowid_"] = rowid
  690. if err := validateNotNull(schema, nr); err != nil {
  691. return 0, 0, err
  692. }
  693. if rowid > maxRowID {
  694. maxRowID = rowid
  695. }
  696. normalized = append(normalized, nr)
  697. }
  698. if maxRowID > 0 {
  699. m.schema.UpdateMaxRowID(table, maxRowID)
  700. }
  701. // Serialize all rows into batch operations. A parallel slice keeps the
  702. // normalized Row for each op for index maintenance after the write.
  703. ops := make([]BatchOp, 0, len(normalized))
  704. encoded := make([]Row, 0, len(normalized))
  705. for _, nr := range normalized {
  706. pk := fmt.Sprintf("%v", nr[schema.PrimaryKey])
  707. data, err := encodeRow(nr)
  708. if err != nil {
  709. return 0, 0, err
  710. }
  711. ops = append(ops, BatchOp{Op: batchPut, Key: []byte(m.dataKey(table, pk)), Value: data})
  712. encoded = append(encoded, nr)
  713. }
  714. keys := make([][]byte, len(ops))
  715. seen := make(map[string]struct{}, len(ops))
  716. for i, op := range ops {
  717. key := string(op.Key)
  718. if _, duplicate := seen[key]; duplicate {
  719. return 0, 0, &PrimaryKeyViolationError{Table: table, Key: key}
  720. }
  721. seen[key] = struct{}{}
  722. keys[i] = op.Key
  723. }
  724. existing := make([]bool, len(keys))
  725. if err := m.pool.WithClient(func(client *KVClient) error {
  726. for start := 0; start < len(keys); start += maxOperations {
  727. end := start + maxOperations
  728. if end > len(keys) {
  729. end = len(keys)
  730. }
  731. found, err := client.ExistsMany(keys[start:end])
  732. if err != nil {
  733. return err
  734. }
  735. copy(existing[start:end], found)
  736. }
  737. return nil
  738. }); err != nil {
  739. return 0, 0, err
  740. }
  741. for i, found := range existing {
  742. if found {
  743. return 0, 0, &PrimaryKeyViolationError{Table: table, Key: string(ops[i].Key)}
  744. }
  745. }
  746. var uniquePlan uniqueCachePlan
  747. unique, err := m.hasUniqueIndex(table)
  748. if err != nil {
  749. return 0, 0, err
  750. }
  751. if unique {
  752. uniqueEntries := make(map[string]*overlayEntry, len(encoded))
  753. for i, row := range encoded {
  754. uniqueEntries[string(ops[i].Key)] = &overlayEntry{row: row}
  755. }
  756. uniquePlan, err = m.prepareUniqueRows(table, uniqueEntries)
  757. if err != nil {
  758. return 0, 0, err
  759. }
  760. }
  761. // Write rows in atomic BATCH_WRITE chunks. Maintain in-memory indexes only
  762. // for rows that actually persisted, so a partial failure cannot leave an
  763. // already-built index stale.
  764. var firstErr error
  765. numOK := 0
  766. for _, chunk := range chunkBatchOps(ops) {
  767. err := m.pool.WithClient(func(c *KVClient) error {
  768. _, err := c.BatchWrite(chunk, nil)
  769. return err
  770. })
  771. if err != nil {
  772. if firstErr == nil {
  773. firstErr = err
  774. }
  775. break
  776. }
  777. numOK += len(chunk)
  778. }
  779. // Maintain indexes for the rows that persisted (the first numOK ops).
  780. if len(uniquePlan) != 0 {
  781. accepted := make(map[string]bool, numOK)
  782. for i := 0; i < numOK; i++ {
  783. accepted[string(ops[i].Key)] = true
  784. }
  785. m.applyUniquePlan(uniquePlan, accepted)
  786. }
  787. for i := 0; i < numOK; i++ {
  788. if err := m.updateIndexesForRow(table, encoded[i], true, string(ops[i].Key)); err != nil {
  789. return numOK, 0, err
  790. }
  791. }
  792. if numOK > 0 {
  793. m.bumpGeneration(table)
  794. if err := m.bumpIndexPredicates(table, encoded[:numOK]...); err != nil {
  795. return numOK, 0, err
  796. }
  797. }
  798. m.incrCount(table, schema.CreatedAt, numOK, wasInit)
  799. if firstErr != nil {
  800. m.invalidateWriteFailure(table)
  801. }
  802. var lastRowID int64
  803. if numOK > 0 {
  804. lastRowID, _ = rowIDFromRow(encoded[numOK-1])
  805. }
  806. return numOK, lastRowID, firstErr
  807. }
  808. // updateIndexesForRow adds or removes entries from already-built in-memory
  809. // indexes. Index entries are rebuildable from durable row data, so this method
  810. // intentionally does not write idx:* keys to KV.
  811. func (m *TableManager) updateIndexesForRow(table string, row Row, add bool, actualKey string) error {
  812. indexes, err := m.schema.ListTableIndexes(table)
  813. if err != nil || len(indexes) == 0 {
  814. return err
  815. }
  816. rowid, ok := rowIDFromRow(row)
  817. if !ok {
  818. return nil
  819. }
  820. tableKey := strings.ToLower(table)
  821. m.cacheMu.Lock()
  822. if keys, initialized := m.rowKeyCache[tableKey]; initialized {
  823. pk := strings.TrimPrefix(actualKey, m.dataPrefix(table))
  824. reverse := m.keyRowCache[tableKey]
  825. if add {
  826. keys[rowid] = pk
  827. reverse[pk] = rowid
  828. } else {
  829. delete(keys, rowid)
  830. delete(reverse, pk)
  831. }
  832. }
  833. m.cacheMu.Unlock()
  834. for _, idx := range indexes {
  835. indexName := strings.ToLower(idx.Name)
  836. m.cacheMu.RLock()
  837. _, initialized := m.indexCache[indexName]
  838. m.cacheMu.RUnlock()
  839. if !initialized {
  840. continue
  841. }
  842. colValue, err := m.buildIndexValue(row, idx.Columns)
  843. if err != nil {
  844. return err
  845. }
  846. if add {
  847. m.AddIndexEntry(idx.Name, colValue, rowid)
  848. } else {
  849. m.RemoveIndexEntry(idx.Name, colValue, rowid)
  850. }
  851. }
  852. return nil
  853. }
  854. // Select retrieves rows from a table by scanning durable rows and collecting
  855. // only matching rows.
  856. func (m *TableManager) Select(table string, filter func(Row) bool) ([]Row, error) {
  857. tl := m.tableLock(table)
  858. tl.RLock()
  859. defer tl.RUnlock()
  860. return m.selectRows(table, filter)
  861. }
  862. // selectRows scans a table while the caller holds its shared or exclusive gate.
  863. func (m *TableManager) selectRows(table string, filter func(Row) bool) ([]Row, error) {
  864. if !m.schema.TableExists(table) {
  865. return nil, fmt.Errorf("table not found: %s", table)
  866. }
  867. rows := make([]Row, 0)
  868. err := m.scanRows(table, func(row Row) (bool, error) {
  869. if filter == nil || filter(row) {
  870. rows = append(rows, row)
  871. }
  872. return false, nil
  873. })
  874. if err != nil {
  875. return nil, err
  876. }
  877. return rows, nil
  878. }
  879. func cloneRow(row Row) Row {
  880. if row == nil {
  881. return nil
  882. }
  883. cloned := make(Row, len(row))
  884. for k, v := range row {
  885. cloned[k] = v
  886. }
  887. return cloned
  888. }
  889. // SelectWithLimit retrieves rows with limit and offset, applying filter/offset
  890. // while scanning and closing the scan early once the limit is reached.
  891. func (m *TableManager) SelectWithLimit(table string, filter func(Row) bool, limit, offset int) ([]Row, error) {
  892. tl := m.tableLock(table)
  893. tl.RLock()
  894. defer tl.RUnlock()
  895. if !m.schema.TableExists(table) {
  896. return nil, fmt.Errorf("table not found: %s", table)
  897. }
  898. rows := make([]Row, 0)
  899. skipped := 0
  900. pageSize := scanPageSize
  901. if limit > 0 {
  902. desired := limit
  903. if offset > 0 {
  904. if offset >= scanPageSize-desired {
  905. desired = scanPageSize
  906. } else {
  907. desired += offset
  908. }
  909. }
  910. if desired < 1 {
  911. desired = 1
  912. }
  913. if desired < pageSize {
  914. pageSize = desired
  915. }
  916. }
  917. err := m.scanRowsWithPageSize(table, uint32(pageSize), func(row Row) (bool, error) {
  918. if filter != nil && !filter(row) {
  919. return false, nil
  920. }
  921. if skipped < offset {
  922. skipped++
  923. return false, nil
  924. }
  925. rows = append(rows, row)
  926. return limit > 0 && len(rows) >= limit, nil
  927. })
  928. if err != nil {
  929. return nil, err
  930. }
  931. return rows, nil
  932. }
  933. // Update updates rows matching the filter.
  934. func (m *TableManager) Update(table string, updates Row, filter func(Row) bool) (int, error) {
  935. return m.updateLocked(table, func(Row) (Row, error) { return updates, nil }, filter)
  936. }
  937. // UpdateFunc updates rows matching the filter using a function to compute new values.
  938. // The updateFn receives the current row and returns the updates to apply.
  939. func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
  940. return m.updateLocked(table, updateFn, filter)
  941. }
  942. // updateLocked stages the complete update before one compare-batch publication.
  943. // Evaluation, encoding, validation or I/O errors must never leave earlier rows
  944. // committed, including when a primary key moves to a different durable key.
  945. func (m *TableManager) updateLocked(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
  946. session := NewSession(m.schema, m)
  947. if err := session.Begin(); err != nil {
  948. return 0, err
  949. }
  950. defer session.Rollback()
  951. count, err := session.UpdateFunc(table, updateFn, filter)
  952. if err != nil {
  953. return 0, err
  954. }
  955. if err := session.Commit(); err != nil {
  956. return 0, err
  957. }
  958. return count, nil
  959. }
  960. // UpdateByPK updates one row without scanning the table. It uses the key's
  961. // striped lock and a compare-and-swap write so a concurrent modification of the
  962. // same row fails with a serialization error instead of being silently lost.
  963. func (m *TableManager) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) {
  964. schema, err := m.schema.GetSchema(table)
  965. if err != nil {
  966. return nil, false, err
  967. }
  968. key := m.dataKey(table, pk)
  969. unlock, uniq, err := m.lockForWrite(table)
  970. if err != nil {
  971. return nil, false, err
  972. }
  973. defer unlock()
  974. st := m.stripeKey(key)
  975. st.Lock()
  976. defer st.Unlock()
  977. row, lsn, err := m.getByPKWithLSN(table, pk)
  978. if err == ErrKeyNotFound {
  979. return nil, false, nil
  980. }
  981. if err != nil {
  982. return nil, false, err
  983. }
  984. oldRow := cloneRow(row)
  985. updates, err := updateFn(row)
  986. if err != nil {
  987. return nil, false, err
  988. }
  989. for name, value := range updates {
  990. for _, column := range schema.Columns {
  991. if strings.EqualFold(name, column.Name) {
  992. row[column.Name] = value
  993. break
  994. }
  995. }
  996. }
  997. if err := validateNotNull(schema, row); err != nil {
  998. return nil, false, err
  999. }
  1000. data, err := encodeRow(row)
  1001. if err != nil {
  1002. return nil, false, err
  1003. }
  1004. var uniquePlan uniqueCachePlan
  1005. if uniq {
  1006. uniquePlan, err = m.prepareUniqueRows(table, map[string]*overlayEntry{key: {row: row}})
  1007. if err != nil {
  1008. return nil, false, err
  1009. }
  1010. }
  1011. committed, err := m.compareWritePoint(
  1012. []CompareCheck{{Key: []byte(key), LSN: lsn}},
  1013. []BatchOp{{Op: batchPut, Key: []byte(key), Value: data}},
  1014. )
  1015. if err != nil {
  1016. m.invalidateWriteFailure(table)
  1017. return nil, false, err
  1018. }
  1019. if !committed {
  1020. return nil, false, ErrSerialization
  1021. }
  1022. m.applyUniquePlan(uniquePlan, nil)
  1023. if err := m.updateIndexesForRow(table, oldRow, false, key); err != nil {
  1024. return nil, false, err
  1025. }
  1026. if err := m.updateIndexesForRow(table, row, true, key); err != nil {
  1027. return nil, false, err
  1028. }
  1029. if err := m.bumpIndexPredicates(table, oldRow, row); err != nil {
  1030. return nil, false, err
  1031. }
  1032. m.bumpGeneration(table)
  1033. return oldRow, true, nil
  1034. }
  1035. // Delete deletes rows matching the filter.
  1036. func (m *TableManager) Delete(table string, filter func(Row) bool) (int, error) {
  1037. session := NewSession(m.schema, m)
  1038. if err := session.Begin(); err != nil {
  1039. return 0, err
  1040. }
  1041. defer session.Rollback()
  1042. count, err := session.Delete(table, filter)
  1043. if err != nil {
  1044. return 0, err
  1045. }
  1046. if err := session.Commit(); err != nil {
  1047. return 0, err
  1048. }
  1049. return count, nil
  1050. }
  1051. // DeleteByPK deletes one row without scanning the table, using the key's
  1052. // striped lock and a compare-and-swap delete.
  1053. func (m *TableManager) DeleteByPK(table, pk string) (Row, bool, error) {
  1054. schema, err := m.schema.GetSchema(table)
  1055. if err != nil {
  1056. return nil, false, err
  1057. }
  1058. key := m.dataKey(table, pk)
  1059. tl := m.tableLock(table)
  1060. tl.RLock()
  1061. defer tl.RUnlock()
  1062. st := m.stripeKey(key)
  1063. st.Lock()
  1064. defer st.Unlock()
  1065. wasInit := m.countInitialized(table)
  1066. row, lsn, err := m.getByPKWithLSN(table, pk)
  1067. if err == ErrKeyNotFound {
  1068. return nil, false, nil
  1069. }
  1070. if err != nil {
  1071. return nil, false, err
  1072. }
  1073. committed, err := m.compareWritePoint(
  1074. []CompareCheck{{Key: []byte(key), LSN: lsn}},
  1075. []BatchOp{{Op: batchDelete, Key: []byte(key)}},
  1076. )
  1077. if err != nil {
  1078. m.invalidateWriteFailure(table)
  1079. return nil, false, err
  1080. }
  1081. if !committed {
  1082. return nil, false, ErrSerialization
  1083. }
  1084. m.removeUniqueKey(table, key)
  1085. if err := m.updateIndexesForRow(table, row, false, key); err != nil {
  1086. return nil, false, err
  1087. }
  1088. if err := m.bumpIndexPredicates(table, row); err != nil {
  1089. return nil, false, err
  1090. }
  1091. m.bumpGeneration(table)
  1092. m.incrCount(table, schema.CreatedAt, -1, wasInit)
  1093. return row, true, nil
  1094. }
  1095. // GetByPK retrieves a row by primary key. Point reads take only the key's
  1096. // striped lock so reads of different keys progress concurrently.
  1097. func (m *TableManager) GetByPK(table string, pk string) (Row, error) {
  1098. key := m.dataKey(table, pk)
  1099. st := m.stripeKey(key)
  1100. st.Lock()
  1101. defer st.Unlock()
  1102. if !m.schema.TableExists(table) {
  1103. return nil, fmt.Errorf("table not found: %s", table)
  1104. }
  1105. row, _, err := m.getByPKWithLSN(table, pk)
  1106. return row, err
  1107. }
  1108. // getByPKWithLSN reads a row by primary key and returns its KV LSN (0 when
  1109. // absent). It performs no locking; callers must hold the appropriate striped
  1110. // or table lock.
  1111. func (m *TableManager) getByPKWithLSN(table, pk string) (Row, uint64, error) {
  1112. key := m.dataKey(table, pk)
  1113. var value []byte
  1114. var lsn uint64
  1115. err := m.pool.WithClient(func(c *KVClient) error {
  1116. res, err := c.Get([]byte(key))
  1117. if err != nil {
  1118. return err
  1119. }
  1120. value = res.Value
  1121. lsn = res.LSN
  1122. return nil
  1123. })
  1124. if err != nil {
  1125. if err == ErrKeyNotFound {
  1126. return nil, 0, ErrKeyNotFound
  1127. }
  1128. return nil, 0, err
  1129. }
  1130. row, err := decodeRow(value)
  1131. if err != nil {
  1132. return nil, 0, fmt.Errorf("failed to parse row: %w", err)
  1133. }
  1134. return row, lsn, nil
  1135. }
  1136. // Count returns the number of rows in a table matching the filter.
  1137. func (m *TableManager) Count(table string, filter func(Row) bool) (int, error) {
  1138. tl := m.tableLock(table)
  1139. tl.RLock()
  1140. defer tl.RUnlock()
  1141. if !m.schema.TableExists(table) {
  1142. return 0, fmt.Errorf("table not found: %s", table)
  1143. }
  1144. count := 0
  1145. err := m.scanRows(table, func(row Row) (bool, error) {
  1146. if filter == nil || filter(row) {
  1147. count++
  1148. }
  1149. return false, nil
  1150. })
  1151. if err != nil {
  1152. return 0, err
  1153. }
  1154. return count, nil
  1155. }
  1156. // Truncate removes all rows from a table.
  1157. func (m *TableManager) Truncate(table string) (int, error) {
  1158. return m.Delete(table, nil)
  1159. }
  1160. // isIntegerType checks if a type name is an integer type.
  1161. func isIntegerType(typeName string) bool {
  1162. t := strings.ToUpper(typeName)
  1163. switch t {
  1164. case "INTEGER", "INT", "SMALLINT", "BIGINT", "TINYINT", "MEDIUMINT":
  1165. return true
  1166. }
  1167. return false
  1168. }
  1169. // IsRowIDColumn checks if a column name is a ROWID alias.
  1170. func IsRowIDColumn(name string) bool {
  1171. n := strings.ToLower(name)
  1172. return n == "rowid" || n == "oid" || n == "_rowid_"
  1173. }
  1174. // Index entry methods - leveraging radix trie for prefix-based lookups
  1175. // Format: {database}:idx:{index_name}:{column_value} → JSON array of rowids
  1176. // indexEntryKey returns the key for an index entry.
  1177. func (m *TableManager) indexEntryKey(indexName string, colValue interface{}) string {
  1178. return fmt.Sprintf("%s:idx:%s:%s", m.database, strings.ToLower(indexName), formatIndexValue(colValue))
  1179. }
  1180. // indexPrefix returns the prefix for all entries of an index.
  1181. func (m *TableManager) indexPrefix(indexName string) string {
  1182. return fmt.Sprintf("%s:idx:%s:", m.database, strings.ToLower(indexName))
  1183. }
  1184. func formatIndexValue(value interface{}) string {
  1185. // Equality predicates compare numeric values (including numeric text) as
  1186. // float64. Use the same equivalence classes without decimal rounding.
  1187. // Index hits are candidates: the executor still evaluates the predicate.
  1188. var number float64
  1189. switch v := value.(type) {
  1190. case float64:
  1191. number = v
  1192. case int64:
  1193. number = float64(v)
  1194. case int:
  1195. number = float64(v)
  1196. case bool:
  1197. if v {
  1198. return "1"
  1199. }
  1200. return "0"
  1201. case []byte:
  1202. return formatIndexValue(string(v))
  1203. case string:
  1204. parsed, err := strconv.ParseFloat(v, 64)
  1205. if err != nil {
  1206. return v
  1207. }
  1208. number = parsed
  1209. default:
  1210. return fmt.Sprintf("%v", v)
  1211. }
  1212. if number == 0 {
  1213. return "0"
  1214. }
  1215. return strconv.FormatFloat(number, 'g', -1, 64)
  1216. }
  1217. func rowIDFromRow(row Row) (int64, bool) {
  1218. switch v := row["_rowid_"].(type) {
  1219. case int64:
  1220. return v, true
  1221. case int:
  1222. return int64(v), true
  1223. case float64:
  1224. return int64(v), true
  1225. default:
  1226. return 0, false
  1227. }
  1228. }
  1229. func (m *TableManager) ensureIndex(index *Index) error {
  1230. indexKey := strings.ToLower(index.Name)
  1231. m.cacheMu.RLock()
  1232. disabled := m.disabledIndexes[indexKey]
  1233. _, initialized := m.indexCache[indexKey]
  1234. m.cacheMu.RUnlock()
  1235. if disabled {
  1236. return nil
  1237. }
  1238. if initialized {
  1239. return nil
  1240. }
  1241. // Serialize index build against writes to the same table so the derived
  1242. // entries cannot miss a concurrently-inserted row.
  1243. table := index.Table
  1244. key := strings.ToLower(table)
  1245. tl := m.tableLock(key)
  1246. tl.Lock()
  1247. defer tl.Unlock()
  1248. m.cacheMu.RLock()
  1249. disabled = m.disabledIndexes[indexKey]
  1250. _, initialized = m.indexCache[indexKey]
  1251. m.cacheMu.RUnlock()
  1252. if disabled {
  1253. return nil
  1254. }
  1255. if initialized {
  1256. return nil
  1257. }
  1258. tableSchema, err := m.schema.GetSchema(table)
  1259. if err != nil {
  1260. return err
  1261. }
  1262. values := make(map[string][]int64)
  1263. rowKeys := make(map[int64]string)
  1264. if err := m.scanRows(table, func(row Row) (bool, error) {
  1265. rowid, ok := rowIDFromRow(row)
  1266. if !ok {
  1267. return false, nil
  1268. }
  1269. colValue, err := m.buildIndexValue(row, index.Columns)
  1270. if err != nil {
  1271. return true, err
  1272. }
  1273. valueKey := formatIndexValue(colValue)
  1274. values[valueKey] = append(values[valueKey], rowid)
  1275. rowKeys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  1276. return false, nil
  1277. }); err != nil {
  1278. return err
  1279. }
  1280. m.cacheMu.Lock()
  1281. if _, initialized := m.indexCache[indexKey]; !initialized {
  1282. m.installIndexCacheLocked(indexKey, key, values, rowKeys)
  1283. }
  1284. m.cacheMu.Unlock()
  1285. return nil
  1286. }
  1287. // AddIndexEntry adds a rowid to an in-memory index entry.
  1288. func (m *TableManager) AddIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1289. indexKey := strings.ToLower(indexName)
  1290. valueKey := formatIndexValue(colValue)
  1291. m.cacheMu.Lock()
  1292. defer m.cacheMu.Unlock()
  1293. values, ok := m.indexCache[indexKey]
  1294. if !ok {
  1295. return nil
  1296. }
  1297. rowids := values[valueKey]
  1298. for _, r := range rowids {
  1299. if r == rowid {
  1300. return nil
  1301. }
  1302. }
  1303. values[valueKey] = append(rowids, rowid)
  1304. if m.indexRowValues[indexKey] == nil {
  1305. m.indexRowValues[indexKey] = make(map[int64]string)
  1306. }
  1307. m.indexRowValues[indexKey][rowid] = valueKey
  1308. return nil
  1309. }
  1310. // RemoveIndexEntry removes a rowid from an in-memory index entry.
  1311. func (m *TableManager) RemoveIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1312. indexKey := strings.ToLower(indexName)
  1313. valueKey := formatIndexValue(colValue)
  1314. m.cacheMu.Lock()
  1315. defer m.cacheMu.Unlock()
  1316. m.removeIndexEntryLocked(indexKey, valueKey, rowid)
  1317. return nil
  1318. }
  1319. // LookupIndex returns rowids matching a column value using the index.
  1320. func (m *TableManager) LookupIndex(indexName string, colValue interface{}) ([]int64, error) {
  1321. index, err := m.schema.GetIndex(indexName)
  1322. if err != nil {
  1323. return nil, err
  1324. }
  1325. indexKey := strings.ToLower(indexName)
  1326. valueKey := formatIndexValue(colValue)
  1327. for {
  1328. if err := m.ensureIndex(index); err != nil {
  1329. return nil, err
  1330. }
  1331. m.cacheMu.RLock()
  1332. values, initialized := m.indexCache[indexKey]
  1333. disabled := m.disabledIndexes[indexKey]
  1334. rowids := append([]int64(nil), values[valueKey]...)
  1335. m.cacheMu.RUnlock()
  1336. if disabled {
  1337. return nil, fmt.Errorf("%w: %s is disabled", ErrIndexNotFound, indexName)
  1338. }
  1339. if initialized {
  1340. return rowids, nil
  1341. }
  1342. // A shared-gate commit invalidated the cache after ensureIndex.
  1343. }
  1344. }
  1345. // ClearIndex removes an index's in-memory entries and marks it disabled so a
  1346. // concurrent lookup cannot rebuild it after DROP but before the schema entry
  1347. // is removed. Index entries are derived from durable rows, so there is nothing
  1348. // durable to delete here.
  1349. func (m *TableManager) ClearIndex(indexName, tableName string, columns []string) error {
  1350. indexKey := strings.ToLower(indexName)
  1351. tableKey := strings.ToLower(tableName)
  1352. m.cacheMu.Lock()
  1353. delete(m.indexCache, indexKey)
  1354. delete(m.indexRowValues, indexKey)
  1355. delete(m.indexTable, indexKey)
  1356. delete(m.uniqueCache, indexKey)
  1357. m.disabledIndexes[indexKey] = true
  1358. rowKeysNeeded := false
  1359. for _, indexedTable := range m.indexTable {
  1360. if indexedTable == tableKey {
  1361. rowKeysNeeded = true
  1362. break
  1363. }
  1364. }
  1365. if !rowKeysNeeded {
  1366. delete(m.rowKeyCache, tableKey)
  1367. delete(m.keyRowCache, tableKey)
  1368. }
  1369. m.cacheMu.Unlock()
  1370. return nil
  1371. }
  1372. // BuildIndex builds index entries for all existing rows in a table.
  1373. func (m *TableManager) BuildIndex(indexName, tableName string, columns []string) error {
  1374. indexKey := strings.ToLower(indexName)
  1375. m.cacheMu.Lock()
  1376. delete(m.disabledIndexes, indexKey)
  1377. delete(m.indexCache, indexKey)
  1378. delete(m.indexRowValues, indexKey)
  1379. delete(m.indexTable, indexKey)
  1380. m.cacheMu.Unlock()
  1381. index, err := m.schema.GetIndex(indexName)
  1382. if err == nil {
  1383. return m.ensureIndex(index)
  1384. }
  1385. if !errors.Is(err, ErrIndexNotFound) {
  1386. return err
  1387. }
  1388. tableSchema, schemaErr := m.schema.GetSchema(tableName)
  1389. if schemaErr != nil {
  1390. return schemaErr
  1391. }
  1392. indexCols := make([]IndexColumn, len(columns))
  1393. for i, name := range columns {
  1394. indexCols[i] = IndexColumn{Name: name}
  1395. }
  1396. values := make(map[string][]int64)
  1397. rowKeys := make(map[int64]string)
  1398. if err := m.scanRows(tableName, func(row Row) (bool, error) {
  1399. rowid, ok := rowIDFromRow(row)
  1400. if !ok {
  1401. return false, nil
  1402. }
  1403. colValue, err := m.buildIndexValue(row, indexCols)
  1404. if err != nil {
  1405. return true, err
  1406. }
  1407. values[formatIndexValue(colValue)] = append(values[formatIndexValue(colValue)], rowid)
  1408. rowKeys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  1409. return false, nil
  1410. }); err != nil {
  1411. return err
  1412. }
  1413. m.cacheMu.Lock()
  1414. m.installIndexCacheLocked(indexKey, strings.ToLower(tableName), values, rowKeys)
  1415. m.cacheMu.Unlock()
  1416. return nil
  1417. }
  1418. // indexColumnValue resolves a single index column's value for a row. A plain
  1419. // column is looked up case-insensitively; an expression column is evaluated
  1420. // through the registered expression evaluator.
  1421. func (m *TableManager) indexColumnValue(row Row, col IndexColumn) (interface{}, error) {
  1422. if col.Expression != "" {
  1423. eval := m.expressionEvaluator()
  1424. if eval == nil {
  1425. return nil, fmt.Errorf("expression index column %q requires an expression evaluator", col.Expression)
  1426. }
  1427. return eval(col.Expression, row)
  1428. }
  1429. if v, ok := row[col.Name]; ok {
  1430. return v, nil
  1431. }
  1432. for k, v := range row {
  1433. if strings.EqualFold(k, col.Name) {
  1434. return v, nil
  1435. }
  1436. }
  1437. return nil, nil
  1438. }
  1439. // buildIndexValue creates the index key value from an index's columns. Single
  1440. // column values are formatted directly; composite values are joined with a NUL
  1441. // separator.
  1442. func (m *TableManager) buildIndexValue(row Row, columns []IndexColumn) (string, error) {
  1443. if len(columns) == 1 {
  1444. v, err := m.indexColumnValue(row, columns[0])
  1445. if err != nil {
  1446. return "", err
  1447. }
  1448. return formatIndexValue(v), nil
  1449. }
  1450. parts := make([]string, 0, len(columns))
  1451. for _, col := range columns {
  1452. v, err := m.indexColumnValue(row, col)
  1453. if err != nil {
  1454. return "", err
  1455. }
  1456. parts = append(parts, formatIndexValue(v))
  1457. }
  1458. return strings.Join(parts, "\x00"), nil
  1459. }
  1460. // ── UNIQUE index enforcement ────────────────────────────────────────────────
  1461. //
  1462. // Uniqueness is enforced by scanning durable rows rather than by persisting a
  1463. // separate claim key. There is therefore no storage migration and no new durable
  1464. // state: a table with a UNIQUE index takes its exclusive table gate for writes so
  1465. // a validating scan can never race a concurrent writer. NULL values are exempt
  1466. // (SQLite permits multiple NULLs in a unique index), and composite values are
  1467. // encoded with length prefixes so "a\x00b" in one column can never collide with
  1468. // "a", "b" across two columns.
  1469. // uniqueIndexes returns the UNIQUE indexes defined on a table.
  1470. func (m *TableManager) uniqueIndexes(table string) ([]*Index, error) {
  1471. indexes, err := m.schema.ListTableIndexes(table)
  1472. if err != nil {
  1473. return nil, err
  1474. }
  1475. var uniq []*Index
  1476. for _, idx := range indexes {
  1477. if idx.Unique {
  1478. uniq = append(uniq, idx)
  1479. }
  1480. }
  1481. return uniq, nil
  1482. }
  1483. // hasUniqueIndex reports whether the table has any UNIQUE index. An IO error is
  1484. // surfaced so callers never silently skip uniqueness enforcement.
  1485. func (m *TableManager) hasUniqueIndex(table string) (bool, error) {
  1486. indexes, err := m.uniqueIndexes(table)
  1487. if err != nil {
  1488. return false, err
  1489. }
  1490. return len(indexes) > 0, nil
  1491. }
  1492. // HasUniqueIndex reports whether the table has any UNIQUE index, surfacing IO
  1493. // errors. It is the exported form used by the executor to decide whether to run a
  1494. // DML statement inside an implicit transaction for statement-level atomicity.
  1495. func (m *TableManager) HasUniqueIndex(table string) (bool, error) {
  1496. return m.hasUniqueIndex(table)
  1497. }
  1498. // encodeUniqueValue encodes the indexed value of a row with per-column type tags
  1499. // and length prefixes. It returns isNull=true when any indexed column is NULL, in
  1500. // which case the row is exempt from uniqueness. The length prefix makes composite
  1501. // encodings collision-free. Expression columns are evaluated through the
  1502. // registered evaluator.
  1503. func (m *TableManager) encodeUniqueValue(row Row, columns []IndexColumn) (string, bool, error) {
  1504. var sb strings.Builder
  1505. for _, col := range columns {
  1506. v, err := m.indexColumnValue(row, col)
  1507. if err != nil {
  1508. return "", false, err
  1509. }
  1510. if v == nil {
  1511. return "", true, nil
  1512. }
  1513. s := encodeUniqueScalar(v)
  1514. sb.WriteString(strconv.Itoa(len(s)))
  1515. sb.WriteByte(':')
  1516. sb.WriteString(s)
  1517. }
  1518. return sb.String(), false, nil
  1519. }
  1520. // encodeUniqueScalar encodes a scalar for uniqueness comparison. Integral
  1521. // numerics are canonicalized across their integer/float/unsigned Go
  1522. // representations so a computed value such as 1.5-0.5 (float64) collides with
  1523. // the literal 1 (int64) the same way SQLite's numeric affinity does. Non-integral
  1524. // reals keep a full-precision tag so no distinct value is lost. TEXT and BLOB
  1525. // remain distinct from numbers.
  1526. func encodeUniqueScalar(v interface{}) string {
  1527. switch t := v.(type) {
  1528. case bool:
  1529. if t {
  1530. return "b1"
  1531. }
  1532. return "b0"
  1533. case int:
  1534. return "i" + strconv.FormatInt(int64(t), 10)
  1535. case int64:
  1536. return "i" + strconv.FormatInt(t, 10)
  1537. case uint:
  1538. return "i" + strconv.FormatUint(uint64(t), 10)
  1539. case uint8:
  1540. return "i" + strconv.FormatUint(uint64(t), 10)
  1541. case uint16:
  1542. return "i" + strconv.FormatUint(uint64(t), 10)
  1543. case uint32:
  1544. return "i" + strconv.FormatUint(uint64(t), 10)
  1545. case uint64:
  1546. return "i" + strconv.FormatUint(t, 10)
  1547. case uintptr:
  1548. return "i" + strconv.FormatUint(uint64(t), 10)
  1549. case float64:
  1550. if s, ok := encodeIntegralFloat(t); ok {
  1551. return s
  1552. }
  1553. return "r" + strconv.FormatFloat(t, 'g', -1, 64)
  1554. case string:
  1555. return "s" + t
  1556. case []byte:
  1557. return "x" + string(t)
  1558. default:
  1559. return "s" + fmt.Sprintf("%v", t)
  1560. }
  1561. }
  1562. // encodeIntegralFloat canonicalizes a whole float64 to the same "i" encoding used
  1563. // by integer values, without precision loss, so integral reals and integers are
  1564. // not treated as distinct under a unique index. It reports false (and returns
  1565. // nothing) for non-integral values or values outside the int64 range.
  1566. func encodeIntegralFloat(v float64) (string, bool) {
  1567. if v != math.Trunc(v) {
  1568. return "", false
  1569. }
  1570. if v < -9223372036854775808.0 || v >= 9223372036854775808.0 {
  1571. return "", false
  1572. }
  1573. return "i" + strconv.FormatInt(int64(v), 10), true
  1574. }
  1575. // lockForWrite acquires the table gate appropriate for a point write: exclusive
  1576. // when the table has a UNIQUE index (so the validating scan cannot race another
  1577. // writer), shared otherwise. The uniqueness decision is re-checked under the lock
  1578. // so a concurrent CREATE UNIQUE INDEX cannot slip in between the unlocked probe
  1579. // and the lock acquisition; the returned unlock closes the gate.
  1580. func (m *TableManager) lockForWrite(table string) (unlock func(), unique bool, err error) {
  1581. tl := m.tableLock(table)
  1582. for {
  1583. uniq, err := m.hasUniqueIndex(table)
  1584. if err != nil {
  1585. return nil, false, err
  1586. }
  1587. if uniq {
  1588. tl.Lock()
  1589. // Re-check under the exclusive gate (authoritative): CreateUniqueIndex
  1590. // also takes the exclusive gate, so it cannot run while we hold it.
  1591. uniq, err = m.hasUniqueIndex(table)
  1592. if err != nil {
  1593. tl.Unlock()
  1594. return nil, false, err
  1595. }
  1596. if uniq {
  1597. return tl.Unlock, true, nil
  1598. }
  1599. tl.Unlock()
  1600. continue
  1601. }
  1602. tl.RLock()
  1603. // A unique index could have appeared before we acquired the shared gate;
  1604. // re-check and escalate if so.
  1605. uniq, err = m.hasUniqueIndex(table)
  1606. if err != nil {
  1607. tl.RUnlock()
  1608. return nil, false, err
  1609. }
  1610. if uniq {
  1611. tl.RUnlock()
  1612. continue
  1613. }
  1614. return tl.RUnlock, false, nil
  1615. }
  1616. }
  1617. // ValidateUniqueIndex rejects a UNIQUE index definition if existing rows already
  1618. // contain duplicate non-NULL values for the indexed columns.
  1619. func (m *TableManager) ValidateUniqueIndex(index *Index) error {
  1620. seen := make(map[string]bool)
  1621. return m.scanRows(index.Table, func(row Row) (bool, error) {
  1622. encoded, isNull, err := m.encodeUniqueValue(row, index.Columns)
  1623. if err != nil {
  1624. return true, err
  1625. }
  1626. if isNull {
  1627. return false, nil
  1628. }
  1629. if seen[encoded] {
  1630. return true, &UniqueViolationError{Constraint: index.Name}
  1631. }
  1632. seen[encoded] = true
  1633. return false, nil
  1634. })
  1635. }
  1636. // CreateUniqueIndex validates existing rows and registers the index while
  1637. // holding the table's exclusive gate, so a concurrent writer cannot insert a
  1638. // conflicting value between validation and registration.
  1639. func (m *TableManager) CreateUniqueIndex(index *Index) error {
  1640. tl := m.tableLock(index.Table)
  1641. tl.Lock()
  1642. defer tl.Unlock()
  1643. if err := m.ValidateUniqueIndex(index); err != nil {
  1644. return err
  1645. }
  1646. return m.schema.CreateIndex(index)
  1647. }
  1648. type indexedRowVersion struct {
  1649. row Row
  1650. key string
  1651. lsn uint64
  1652. }
  1653. type indexPredicateSnapshot struct {
  1654. valueKey string
  1655. valueGen uint64
  1656. wildcardKey string
  1657. wildcardGen uint64
  1658. }
  1659. // selectByIndexWithLSN retrieves indexed rows and their durable versions. The
  1660. // transaction layer uses the versions for optimistic commit validation.
  1661. func (m *TableManager) selectByIndexWithLSN(table, indexName string, colValue interface{}) ([]indexedRowVersion, indexPredicateSnapshot, error) {
  1662. return m.selectByIndexKeyWithLSN(table, indexName, formatIndexValue(colValue))
  1663. }
  1664. // IndexRowKey computes the formatted in-memory index key a row maps to for an
  1665. // index. Expression columns are evaluated through the same registered evaluator
  1666. // used by uniqueness enforcement, so the key matches both durable index entries
  1667. // and staged overlay rows.
  1668. func (m *TableManager) IndexRowKey(idx *Index, row Row) (string, error) {
  1669. value, err := m.buildIndexValue(row, idx.Columns)
  1670. if err != nil {
  1671. return "", err
  1672. }
  1673. return formatIndexValue(value), nil
  1674. }
  1675. // IndexValueContainsNull reports whether any column contributing to an index
  1676. // key evaluates to NULL. UNIQUE indexes exempt such rows from conflicts under
  1677. // SQLite semantics.
  1678. func (m *TableManager) IndexValueContainsNull(idx *Index, row Row) (bool, error) {
  1679. _, isNull, err := m.encodeUniqueValue(row, idx.Columns)
  1680. return isNull, err
  1681. }
  1682. // selectByIndexKeyWithLSN is selectByIndexWithLSN with an already-computed
  1683. // formatted index key, so callers can look up composite and expression indexes
  1684. // without re-deriving the key from a single column value.
  1685. func (m *TableManager) selectByIndexKeyWithLSN(table, indexName, valueKey string) ([]indexedRowVersion, indexPredicateSnapshot, error) {
  1686. index, err := m.schema.GetIndex(indexName)
  1687. if err != nil {
  1688. return nil, indexPredicateSnapshot{}, err
  1689. }
  1690. tableKey := strings.ToLower(table)
  1691. tl := m.tableLock(tableKey)
  1692. indexKey := strings.ToLower(indexName)
  1693. for {
  1694. if err := m.ensureIndex(index); err != nil {
  1695. return nil, indexPredicateSnapshot{}, err
  1696. }
  1697. tl.RLock()
  1698. if !m.schema.TableExists(table) {
  1699. tl.RUnlock()
  1700. return nil, indexPredicateSnapshot{}, fmt.Errorf("table not found: %s", table)
  1701. }
  1702. m.cacheMu.RLock()
  1703. if m.disabledIndexes[indexKey] {
  1704. m.cacheMu.RUnlock()
  1705. tl.RUnlock()
  1706. return nil, indexPredicateSnapshot{}, fmt.Errorf("%w: %s is disabled", ErrIndexNotFound, indexName)
  1707. }
  1708. if _, initialized := m.indexCache[indexKey]; initialized {
  1709. break // Keep cacheMu through the complete candidate/key snapshot.
  1710. }
  1711. m.cacheMu.RUnlock()
  1712. tl.RUnlock()
  1713. // Commits can invalidate under a shared table gate. Retry the build,
  1714. // never interpret a missing cache as a legitimately empty index.
  1715. }
  1716. defer tl.RUnlock()
  1717. predicateKey, predicateGen, wildcardKey, wildcardGen := m.predicateSnapshot(table, indexName, valueKey)
  1718. snapshot := indexPredicateSnapshot{
  1719. valueKey: predicateKey, valueGen: predicateGen,
  1720. wildcardKey: wildcardKey, wildcardGen: wildcardGen,
  1721. }
  1722. rowids := append([]int64(nil), m.indexCache[indexKey][valueKey]...)
  1723. primaryKeys := make([]string, 0, len(rowids))
  1724. missingRowID := int64(0)
  1725. missingRowKey := false
  1726. for _, rowid := range rowids {
  1727. if primaryKey, ok := m.rowKeyCache[tableKey][rowid]; ok {
  1728. primaryKeys = append(primaryKeys, primaryKey)
  1729. } else {
  1730. missingRowID = rowid
  1731. missingRowKey = true
  1732. break
  1733. }
  1734. }
  1735. m.cacheMu.RUnlock()
  1736. if missingRowKey {
  1737. return nil, indexPredicateSnapshot{}, fmt.Errorf("index %s is missing rowid %d", indexName, missingRowID)
  1738. }
  1739. // If no rowids found, return empty result
  1740. if len(rowids) == 0 {
  1741. return []indexedRowVersion{}, snapshot, nil
  1742. }
  1743. rows := make([]indexedRowVersion, 0, len(primaryKeys))
  1744. err = m.pool.WithClient(func(client *KVClient) error {
  1745. keys := make([][]byte, len(primaryKeys))
  1746. for i, primaryKey := range primaryKeys {
  1747. keys[i] = []byte(m.dataKey(table, primaryKey))
  1748. }
  1749. results, err := client.MultiGet(keys)
  1750. if err != nil {
  1751. return err
  1752. }
  1753. for i, result := range results {
  1754. if !result.Found {
  1755. primaryKey := primaryKeys[i]
  1756. return fmt.Errorf("index %s references missing primary key %s", indexName, primaryKey)
  1757. }
  1758. row, err := decodeRow(result.Value)
  1759. if err != nil {
  1760. return err
  1761. }
  1762. rows = append(rows, indexedRowVersion{
  1763. row: row,
  1764. key: m.dataKey(table, primaryKeys[i]),
  1765. lsn: result.LSN,
  1766. })
  1767. }
  1768. return nil
  1769. })
  1770. if err != nil {
  1771. return nil, indexPredicateSnapshot{}, err
  1772. }
  1773. return rows, snapshot, nil
  1774. }
  1775. // SelectByIndex retrieves rows using an in-memory equality index followed by a
  1776. // single MultiGet for the matching primary keys.
  1777. func (m *TableManager) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
  1778. versions, _, err := m.selectByIndexWithLSN(table, indexName, colValue)
  1779. if err != nil {
  1780. return nil, err
  1781. }
  1782. rows := make([]Row, len(versions))
  1783. for i := range versions {
  1784. rows[i] = versions[i].row
  1785. }
  1786. return rows, nil
  1787. }