table.go 54 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938
  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. indexCache map[string]map[string][]int64 // index name → indexed value → rowids
  31. indexTable map[string]string // index name → table name
  32. rowKeyCache map[string]map[int64]string // table name → rowid → primary key
  33. // disabledIndexes prevents a concurrent lookup from rebuilding an index
  34. // after DROP has cleared it but before the schema entry is removed.
  35. disabledIndexes map[string]bool
  36. // counts holds exact per-table row counts for the COUNT(*) fast path.
  37. // It is derived lazily from durable rows on first use and maintained
  38. // incrementally by Insert/InsertBulk/Delete thereafter.
  39. counts map[string]int
  40. countsInit map[string]bool
  41. countGeneration map[string]time.Time
  42. // stripes are deterministic per-key locks used by point operations
  43. // (GetByPK/Insert/UpdateByPK/DeleteByPK). They replace the table-wide lock
  44. // so point operations on different keys of the same table proceed
  45. // concurrently. Index is derived from the full data key (database+table+pk).
  46. stripes [64]sync.Mutex
  47. // generations tracks full-table scans. predicateGenerations narrows indexed
  48. // equality reads to one index value so unrelated writes do not conflict.
  49. genMu sync.Mutex
  50. generations map[string]uint64
  51. predicateGenerations map[string]uint64
  52. // exprMu guards exprEval, the pluggable evaluator for expression-index
  53. // columns. It is set by the SQL executor so storage stays independent of
  54. // the SQL front end.
  55. exprMu sync.RWMutex
  56. exprEval ExpressionEvaluator
  57. }
  58. // ExpressionEvaluator computes the value of a persisted index expression for a
  59. // row. The SQL executor registers one so the storage layer can enforce
  60. // expression indexes without importing the parser/executor.
  61. type ExpressionEvaluator func(expression string, row Row) (interface{}, error)
  62. // SetExpressionEvaluator installs the evaluator used to compute expression-index
  63. // values during uniqueness validation and index maintenance.
  64. func (m *TableManager) SetExpressionEvaluator(eval ExpressionEvaluator) {
  65. m.exprMu.Lock()
  66. m.exprEval = eval
  67. m.exprMu.Unlock()
  68. }
  69. func (m *TableManager) expressionEvaluator() ExpressionEvaluator {
  70. m.exprMu.RLock()
  71. defer m.exprMu.RUnlock()
  72. return m.exprEval
  73. }
  74. // NewTableManager creates a new table manager.
  75. func NewTableManager(pool *KVPool, schema *SchemaManager, database string) *TableManager {
  76. return &TableManager{
  77. pool: pool,
  78. schema: schema,
  79. database: database,
  80. uniqueCache: make(map[string]*uniqueIndexCache),
  81. indexCache: make(map[string]map[string][]int64),
  82. indexTable: make(map[string]string),
  83. rowKeyCache: make(map[string]map[int64]string),
  84. disabledIndexes: make(map[string]bool),
  85. counts: make(map[string]int),
  86. countsInit: make(map[string]bool),
  87. countGeneration: make(map[string]time.Time),
  88. generations: make(map[string]uint64),
  89. predicateGenerations: make(map[string]uint64),
  90. }
  91. }
  92. func (m *TableManager) tableLock(key string) *sync.RWMutex {
  93. return m.schema.tableLock(key)
  94. }
  95. // stripeKey returns the deterministic striped lock for a point operation on the
  96. // given full data key. Point operations on different keys therefore serialize
  97. // independently, while operations on the same key are mutually exclusive.
  98. func (m *TableManager) stripeKey(key string) *sync.Mutex {
  99. h := fnv.New32a()
  100. h.Write([]byte(key))
  101. return &m.stripes[h.Sum32()%uint32(len(m.stripes))]
  102. }
  103. // generation returns the current in-process generation for a table. It is
  104. // bumped on every committed write and captured by transaction scans.
  105. func (m *TableManager) generation(table string) uint64 {
  106. key := strings.ToLower(table)
  107. m.genMu.Lock()
  108. g := m.generations[key]
  109. m.genMu.Unlock()
  110. return g
  111. }
  112. // bumpGeneration advances a table's generation. Callers hold the table gate in
  113. // shared mode for point writes or exclusive mode for scan-based writes.
  114. func (m *TableManager) bumpGeneration(table string) {
  115. key := strings.ToLower(table)
  116. m.genMu.Lock()
  117. m.generations[key]++
  118. m.genMu.Unlock()
  119. }
  120. func indexPredicateKey(table, indexName, value string) string {
  121. return strings.ToLower(table) + "\x00" + strings.ToLower(indexName) + "\x00" + value
  122. }
  123. func indexPredicateWildcardKey(table string) string {
  124. return strings.ToLower(table) + "\x00*"
  125. }
  126. func (m *TableManager) predicateSnapshot(table, indexName, value string) (string, uint64, string, uint64) {
  127. valueKey := indexPredicateKey(table, indexName, value)
  128. wildcardKey := indexPredicateWildcardKey(table)
  129. m.genMu.Lock()
  130. valueGen := m.predicateGenerations[valueKey]
  131. wildcardGen := m.predicateGenerations[wildcardKey]
  132. m.genMu.Unlock()
  133. return valueKey, valueGen, wildcardKey, wildcardGen
  134. }
  135. func (m *TableManager) predicateGeneration(key string) uint64 {
  136. m.genMu.Lock()
  137. gen := m.predicateGenerations[key]
  138. m.genMu.Unlock()
  139. return gen
  140. }
  141. func (m *TableManager) bumpIndexPredicates(table string, rows ...Row) error {
  142. indexes, err := m.schema.ListTableIndexes(table)
  143. if err != nil || len(indexes) == 0 {
  144. return err
  145. }
  146. m.genMu.Lock()
  147. defer m.genMu.Unlock()
  148. for _, row := range rows {
  149. if row == nil {
  150. continue
  151. }
  152. for _, index := range indexes {
  153. value, err := m.buildIndexValue(row, index.Columns)
  154. if err != nil {
  155. return err
  156. }
  157. m.predicateGenerations[indexPredicateKey(table, index.Name, value)]++
  158. }
  159. }
  160. return nil
  161. }
  162. func (m *TableManager) bumpIndexPredicateWildcard(table string) {
  163. m.genMu.Lock()
  164. m.predicateGenerations[indexPredicateWildcardKey(table)]++
  165. m.genMu.Unlock()
  166. }
  167. // compareWritePoint issues a single CompareBatchWrite against the pooled KV.
  168. // It returns whether the compare checks held and the ops committed.
  169. func (m *TableManager) compareWritePoint(checks []CompareCheck, ops []BatchOp) (bool, error) {
  170. var committed bool
  171. err := m.pool.WithClient(func(c *KVClient) error {
  172. _, ok, err := c.CompareBatchWrite(checks, ops, nil)
  173. committed = ok
  174. return err
  175. })
  176. return committed, err
  177. }
  178. // invalidateCache removes a table's derived in-memory indexes.
  179. func (m *TableManager) invalidateCache(table string) {
  180. m.cacheMu.Lock()
  181. key := strings.ToLower(table)
  182. delete(m.rowKeyCache, key)
  183. delete(m.counts, key)
  184. delete(m.countsInit, key)
  185. delete(m.countGeneration, key)
  186. for indexName, tableName := range m.indexTable {
  187. if tableName == key {
  188. delete(m.indexCache, indexName)
  189. delete(m.indexTable, indexName)
  190. }
  191. }
  192. m.cacheMu.Unlock()
  193. }
  194. // InvalidateCache is the exported version for use by the executor.
  195. func (m *TableManager) InvalidateCache(table string) {
  196. m.invalidateCache(table)
  197. m.cacheMu.Lock()
  198. for name, cache := range m.uniqueCache {
  199. if cache.table == strings.ToLower(table) {
  200. delete(m.uniqueCache, name)
  201. }
  202. }
  203. m.cacheMu.Unlock()
  204. }
  205. // rowVisitFunc is invoked for each decoded row in a streaming scan. Return
  206. // stop=true to end the scan early; a non-nil error aborts the scan.
  207. type rowVisitFunc func(Row) (stop bool, err error)
  208. // scanRows streams the rows of table by scanning durable KV rows one page at a
  209. // time under a single pooled client. Each page is decoded as it arrives and
  210. // passed to fn, which may stop the scan early. The cursor is always closed and
  211. // the client always returned to the pool, even on error.
  212. func (m *TableManager) scanRows(table string, fn rowVisitFunc) error {
  213. return m.scanRowsWithPageSize(table, scanPageSize, fn)
  214. }
  215. func (m *TableManager) scanRowsWithPageSize(table string, pageSize uint32, fn rowVisitFunc) error {
  216. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  217. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), pageSize)
  218. if err != nil {
  219. return err
  220. }
  221. defer func() {
  222. if err := cursor.Close(); retErr == nil {
  223. retErr = err
  224. }
  225. }()
  226. for {
  227. entries, done, err := cursor.Next()
  228. if err != nil {
  229. return err
  230. }
  231. for _, e := range entries {
  232. row, err := decodeRow(e.Value)
  233. if err != nil {
  234. return err
  235. }
  236. stop, err := fn(row)
  237. if err != nil {
  238. return err
  239. }
  240. if stop {
  241. return nil
  242. }
  243. }
  244. if done {
  245. return nil
  246. }
  247. }
  248. })
  249. }
  250. // rowWithLSNVisitFunc is like rowVisitFunc but also passes the durable row LSN.
  251. type rowWithLSNVisitFunc func(row Row, lsn uint64) (stop bool, err error)
  252. // scanRowsWithLSN streams a table's rows together with their durable LSNs. It
  253. // is used by buffered transactions to capture a per-row read set for
  254. // compare-and-swap validation at commit.
  255. func (m *TableManager) scanRowsWithLSN(table string, fn rowWithLSNVisitFunc) error {
  256. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  257. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), scanPageSize)
  258. if err != nil {
  259. return err
  260. }
  261. defer func() {
  262. if err := cursor.Close(); retErr == nil {
  263. retErr = err
  264. }
  265. }()
  266. for {
  267. entries, done, err := cursor.Next()
  268. if err != nil {
  269. return err
  270. }
  271. for _, e := range entries {
  272. row, err := decodeRow(e.Value)
  273. if err != nil {
  274. return err
  275. }
  276. stop, err := fn(row, e.LSN)
  277. if err != nil {
  278. return err
  279. }
  280. if stop {
  281. return nil
  282. }
  283. }
  284. if done {
  285. return nil
  286. }
  287. }
  288. })
  289. }
  290. // scanCountKeys counts the durable rows of table using a key-only scan so row
  291. // values are never pulled across the wire. It is used for first-time COUNT(*)
  292. // derivation.
  293. func (m *TableManager) scanCountKeys(table string) (int, error) {
  294. count := 0
  295. err := m.pool.WithClient(func(client *KVClient) (retErr error) {
  296. cursor, err := client.ScanKeys([]byte(m.dataPrefix(table)))
  297. if err != nil {
  298. return err
  299. }
  300. defer func() {
  301. if err := cursor.Close(); retErr == nil {
  302. retErr = err
  303. }
  304. }()
  305. for {
  306. entries, done, err := cursor.Next()
  307. if err != nil {
  308. return err
  309. }
  310. count += len(entries)
  311. if done {
  312. return nil
  313. }
  314. }
  315. })
  316. return count, err
  317. }
  318. // CountFast returns the exact number of rows in a table. The count is derived
  319. // from durable rows on first use (recovering across restarts) and then
  320. // maintained incrementally by the write paths, so repeated COUNT(*) queries
  321. // avoid a full table scan. It intentionally does not persist a counter to KV:
  322. // the KV layer has no atomic increment primitive, and a durable counter that
  323. // could diverge from the rows on crash would be worse than a lazily-derived,
  324. // always-exact value. The cost is one key-only scan the first time COUNT(*) is
  325. // issued after startup.
  326. func (m *TableManager) CountFast(table string) (int, error) {
  327. key := strings.ToLower(table)
  328. tableSchema, err := m.schema.GetSchema(table)
  329. if err != nil {
  330. return 0, err
  331. }
  332. // Cached read path: if the count is initialized for the current table
  333. // generation, return it without taking any table lock.
  334. m.cacheMu.Lock()
  335. if m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  336. if m.countsInit[key] {
  337. n := m.counts[key]
  338. m.cacheMu.Unlock()
  339. return n, nil
  340. }
  341. } else {
  342. delete(m.counts, key)
  343. delete(m.countsInit, key)
  344. m.countGeneration[key] = tableSchema.CreatedAt
  345. }
  346. m.cacheMu.Unlock()
  347. // First derivation: hold the table write lock so the key-only scan is
  348. // exact against concurrent writes.
  349. tl := m.tableLock(key)
  350. tl.Lock()
  351. defer tl.Unlock()
  352. m.cacheMu.Lock()
  353. if !m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  354. delete(m.counts, key)
  355. delete(m.countsInit, key)
  356. m.countGeneration[key] = tableSchema.CreatedAt
  357. }
  358. if m.countsInit[key] {
  359. n := m.counts[key]
  360. m.cacheMu.Unlock()
  361. return n, nil
  362. }
  363. m.cacheMu.Unlock()
  364. count, err := m.scanCountKeys(table)
  365. if err != nil {
  366. return 0, err
  367. }
  368. m.cacheMu.Lock()
  369. m.counts[key] = count
  370. m.countsInit[key] = true
  371. m.cacheMu.Unlock()
  372. return count, nil
  373. }
  374. // countInitialized reports whether the derived count cache is initialized for
  375. // the table at this instant. Write paths capture it before their KV write so a
  376. // concurrent first derivation does not double-count the just-written row.
  377. func (m *TableManager) countInitialized(table string) bool {
  378. key := strings.ToLower(table)
  379. m.cacheMu.Lock()
  380. init := m.countsInit[key]
  381. m.cacheMu.Unlock()
  382. return init
  383. }
  384. // incrCount adjusts the derived per-table row count. wasInit reports whether
  385. // the count was already initialized before the corresponding write, so a count
  386. // that was not yet initialized is left to be re-derived from durable rows
  387. // (which already reflect the write) on next use.
  388. func (m *TableManager) incrCount(table string, generation time.Time, delta int, wasInit bool) {
  389. key := strings.ToLower(table)
  390. m.cacheMu.Lock()
  391. if !m.countGeneration[key].Equal(generation) {
  392. delete(m.counts, key)
  393. delete(m.countsInit, key)
  394. m.countGeneration[key] = generation
  395. }
  396. if wasInit && m.countsInit[key] {
  397. m.counts[key] += delta
  398. }
  399. if !wasInit {
  400. // The count was not initialized before the write, so a concurrent first
  401. // derivation may have missed the just-written row. Invalidate to force
  402. // an exact re-derivation on the next COUNT(*).
  403. delete(m.countsInit, key)
  404. }
  405. m.cacheMu.Unlock()
  406. }
  407. // dataKey returns the key for a row.
  408. func (m *TableManager) dataKey(table, pk string) string {
  409. return fmt.Sprintf("%s:_data:%s:%s", m.database, strings.ToLower(table), pk)
  410. }
  411. // dataPrefix returns the prefix for all rows in a table.
  412. func (m *TableManager) dataPrefix(table string) string {
  413. return fmt.Sprintf("%s:_data:%s:", m.database, strings.ToLower(table))
  414. }
  415. // prepareInsert validates and normalizes an insert row, generating the ROWID
  416. // when required. It returns the normalized row and the full data key without
  417. // writing anything, so both the autocommit path (compare-and-swap) and the
  418. // buffered transaction path (staging) can share it. The input row has its
  419. // primary key populated as a side effect.
  420. func (m *TableManager) prepareInsert(table string, row Row) (Row, string, error) {
  421. schema, err := m.schema.GetSchema(table)
  422. if err != nil {
  423. return nil, "", err
  424. }
  425. // Get primary key value
  426. pkValue, ok := row[schema.PrimaryKey]
  427. if !ok {
  428. for k, v := range row {
  429. if strings.EqualFold(k, schema.PrimaryKey) {
  430. pkValue = v
  431. ok = true
  432. break
  433. }
  434. }
  435. }
  436. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  437. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  438. var rowid int64
  439. if !ok || pkValue == nil {
  440. if isIntegerPK || !ok {
  441. rowid, err = m.schema.GetNextRowID(table)
  442. if err != nil {
  443. return nil, "", err
  444. }
  445. pkValue = rowid
  446. row[schema.PrimaryKey] = rowid
  447. ok = true
  448. } else {
  449. return nil, "", fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  450. }
  451. } else if isIntegerPK {
  452. switch v := pkValue.(type) {
  453. case int64:
  454. rowid = v
  455. case float64:
  456. if math.Trunc(v) != v {
  457. return nil, "", fmt.Errorf("invalid integer primary key: %v", v)
  458. }
  459. rowid = int64(v)
  460. case int:
  461. rowid = int64(v)
  462. default:
  463. rowid = 0
  464. }
  465. if rowid > 0 {
  466. if err := m.schema.UpdateMaxRowID(table, rowid); err != nil {
  467. return nil, "", err
  468. }
  469. }
  470. }
  471. pk := fmt.Sprintf("%v", pkValue)
  472. normalizedRow := make(Row)
  473. for _, col := range schema.Columns {
  474. for k, v := range row {
  475. if strings.EqualFold(k, col.Name) {
  476. normalizedRow[col.Name] = v
  477. break
  478. }
  479. }
  480. }
  481. for _, col := range schema.Columns {
  482. if _, ok := normalizedRow[col.Name]; !ok && col.Default != nil {
  483. normalizedRow[col.Name] = col.Default
  484. }
  485. }
  486. if rowid > 0 {
  487. normalizedRow["_rowid_"] = rowid
  488. } else {
  489. newRowID, err := m.schema.GetNextRowID(table)
  490. if err != nil {
  491. return nil, "", err
  492. }
  493. normalizedRow["_rowid_"] = newRowID
  494. }
  495. if err := validateNotNull(schema, normalizedRow); err != nil {
  496. return nil, "", err
  497. }
  498. return normalizedRow, m.dataKey(table, pk), nil
  499. }
  500. // Insert inserts a new row. The duplicate check and write are one atomic
  501. // compare-and-swap so concurrent inserts of the same primary key cannot both
  502. // persist. The shared table gate is acquired before the key's striped lock so
  503. // a queued scan writer cannot invert the lock order with point operations.
  504. func (m *TableManager) Insert(table string, row Row) error {
  505. _, err := m.InsertWithRowID(table, row)
  506. return err
  507. }
  508. // InsertWithRowID returns the actual row identifier, never a concurrent counter.
  509. func (m *TableManager) InsertWithRowID(table string, row Row) (int64, error) {
  510. nr, key, err := m.prepareInsert(table, row)
  511. if err != nil {
  512. return 0, err
  513. }
  514. rowid, _ := rowIDFromRow(nr)
  515. data, err := encodeRow(nr)
  516. if err != nil {
  517. return 0, fmt.Errorf("failed to serialize row: %w", err)
  518. }
  519. wasInit := m.countInitialized(table)
  520. // Point writers share this gate with each other. Transaction commits and
  521. // scan-based writes take it exclusively, so generation validation and cache
  522. // publication are ordered without serializing writes to different keys.
  523. // A UNIQUE index forces the exclusive gate so the validating scan below
  524. // cannot race another writer (or a concurrent CREATE UNIQUE INDEX).
  525. unlock, uniq, err := m.lockForWrite(table)
  526. if err != nil {
  527. return 0, err
  528. }
  529. defer unlock()
  530. var uniquePlan uniqueCachePlan
  531. if uniq {
  532. uniquePlan, err = m.prepareUniqueRows(table, map[string]*overlayEntry{key: {row: nr}})
  533. if err != nil {
  534. return 0, err
  535. }
  536. }
  537. st := m.stripeKey(key)
  538. st.Lock()
  539. defer st.Unlock()
  540. committed, err := m.compareWritePoint(
  541. []CompareCheck{{Key: []byte(key), LSN: 0}},
  542. []BatchOp{{Op: batchPut, Key: []byte(key), Value: data}},
  543. )
  544. if err != nil {
  545. m.invalidateWriteFailure(table)
  546. return 0, err
  547. }
  548. if !committed {
  549. return 0, fmt.Errorf("duplicate primary key: %v", row[schemaPrimaryKey(m.schema, table)])
  550. }
  551. m.applyUniquePlan(uniquePlan, nil)
  552. if err := m.updateIndexesForRow(table, nr, true); err != nil {
  553. return 0, err
  554. }
  555. // Publish derived index state before its generations. A reader that races
  556. // with publication either sees the old generation and aborts or sees the
  557. // complete new state.
  558. if err := m.bumpIndexPredicates(table, nr); err != nil {
  559. return 0, err
  560. }
  561. m.bumpGeneration(table)
  562. if schema, serr := m.schema.GetSchema(table); serr == nil {
  563. m.incrCount(table, schema.CreatedAt, 1, wasInit)
  564. }
  565. return rowid, nil
  566. }
  567. func schemaPrimaryKey(s *SchemaManager, table string) string {
  568. schema, err := s.GetSchema(table)
  569. if err != nil {
  570. return "_rowid_"
  571. }
  572. return schema.PrimaryKey
  573. }
  574. // bulkBatchByteBudget bounds a single atomic BATCH_WRITE payload below the
  575. // PKBFI frame limit so a bulk insert never emits a frame the server rejects.
  576. // Each op contributes 12 header bytes plus its key and value.
  577. const bulkBatchByteBudget = 60 * 1024 * 1024
  578. // chunkBatchOps splits ops into atomic BATCH_WRITE chunks bounded by both the
  579. // PKBFI operation-count limit and the frame-size limit. Each chunk is a slice
  580. // of the backing array, valid until the next append to ops.
  581. func chunkBatchOps(ops []BatchOp) [][]BatchOp {
  582. var chunks [][]BatchOp
  583. for i := 0; i < len(ops); {
  584. end := i + maxOperations
  585. if end > len(ops) {
  586. end = len(ops)
  587. }
  588. bytes := 0
  589. j := i
  590. for j < end {
  591. sz := 12 + len(ops[j].Key) + len(ops[j].Value)
  592. if j > i && bytes+sz > bulkBatchByteBudget {
  593. break
  594. }
  595. bytes += sz
  596. j++
  597. }
  598. if j == i {
  599. j = i + 1
  600. }
  601. chunks = append(chunks, ops[i:j])
  602. i = j
  603. }
  604. return chunks
  605. }
  606. // InsertBulk inserts multiple rows efficiently using atomic BATCH_WRITE chunks
  607. // bounded by the PKBFI operation-count and frame-size limits. Skips per-row
  608. // duplicate checks (caller must ensure uniqueness). Used by INSERT ... SELECT.
  609. func (m *TableManager) InsertBulk(table string, rows []Row) (int, error) {
  610. count, _, err := m.InsertBulkWithLastRowID(table, rows)
  611. return count, err
  612. }
  613. // InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the
  614. // last persisted row (0 when nothing persisted).
  615. func (m *TableManager) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) {
  616. if len(rows) == 0 {
  617. return 0, 0, nil
  618. }
  619. tl := m.tableLock(table)
  620. tl.Lock()
  621. defer tl.Unlock()
  622. schema, err := m.schema.GetSchema(table)
  623. if err != nil {
  624. return 0, 0, err
  625. }
  626. wasInit := m.countInitialized(table)
  627. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  628. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  629. // Normalize rows and assign _rowid_.
  630. normalized := make([]Row, 0, len(rows))
  631. var maxRowID int64
  632. for _, row := range rows {
  633. nr := make(Row)
  634. for _, col := range schema.Columns {
  635. for k, v := range row {
  636. if strings.EqualFold(k, col.Name) {
  637. nr[col.Name] = v
  638. break
  639. }
  640. }
  641. }
  642. for _, col := range schema.Columns {
  643. if _, ok := nr[col.Name]; !ok && col.Default != nil {
  644. nr[col.Name] = col.Default
  645. }
  646. }
  647. var rowid int64
  648. var hasRowid bool
  649. if isIntegerPK {
  650. switch v := nr[schema.PrimaryKey].(type) {
  651. case float64:
  652. if math.Trunc(v) != v {
  653. return 0, 0, fmt.Errorf("invalid integer primary key: %v", v)
  654. }
  655. rowid = int64(v)
  656. hasRowid = true
  657. case int64:
  658. rowid = v
  659. hasRowid = true
  660. case int:
  661. rowid = int64(v)
  662. hasRowid = true
  663. }
  664. }
  665. if !hasRowid {
  666. // Auto-generate for an INTEGER primary key or the synthetic _rowid_,
  667. // mirroring prepareInsert so INSERT ... SELECT works on AUTOINCREMENT
  668. // tables.
  669. if isIntegerPK || schema.PrimaryKey == "_rowid_" {
  670. rowid, err = m.schema.GetNextRowID(table)
  671. if err != nil {
  672. return 0, 0, err
  673. }
  674. if isIntegerPK {
  675. nr[schema.PrimaryKey] = rowid
  676. } else {
  677. nr["_rowid_"] = rowid
  678. }
  679. } else {
  680. pk, ok := nr[schema.PrimaryKey]
  681. if !ok || pk == nil {
  682. return 0, 0, fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  683. }
  684. rowid, err = m.schema.GetNextRowID(table)
  685. if err != nil {
  686. return 0, 0, err
  687. }
  688. }
  689. }
  690. nr["_rowid_"] = rowid
  691. if err := validateNotNull(schema, nr); err != nil {
  692. return 0, 0, err
  693. }
  694. if rowid > maxRowID {
  695. maxRowID = rowid
  696. }
  697. normalized = append(normalized, nr)
  698. }
  699. if maxRowID > 0 {
  700. m.schema.UpdateMaxRowID(table, maxRowID)
  701. }
  702. // Serialize all rows into batch operations. A parallel slice keeps the
  703. // normalized Row for each op for index maintenance after the write.
  704. ops := make([]BatchOp, 0, len(normalized))
  705. encoded := make([]Row, 0, len(normalized))
  706. for _, nr := range normalized {
  707. pk := fmt.Sprintf("%v", nr[schema.PrimaryKey])
  708. data, err := encodeRow(nr)
  709. if err != nil {
  710. return 0, 0, err
  711. }
  712. ops = append(ops, BatchOp{Op: batchPut, Key: []byte(m.dataKey(table, pk)), Value: data})
  713. encoded = append(encoded, nr)
  714. }
  715. keys := make([][]byte, len(ops))
  716. seen := make(map[string]struct{}, len(ops))
  717. for i, op := range ops {
  718. key := string(op.Key)
  719. if _, duplicate := seen[key]; duplicate {
  720. return 0, 0, fmt.Errorf("duplicate primary key: %s", key)
  721. }
  722. seen[key] = struct{}{}
  723. keys[i] = op.Key
  724. }
  725. existing := make([]bool, len(keys))
  726. if err := m.pool.WithClient(func(client *KVClient) error {
  727. for start := 0; start < len(keys); start += maxOperations {
  728. end := start + maxOperations
  729. if end > len(keys) {
  730. end = len(keys)
  731. }
  732. found, err := client.ExistsMany(keys[start:end])
  733. if err != nil {
  734. return err
  735. }
  736. copy(existing[start:end], found)
  737. }
  738. return nil
  739. }); err != nil {
  740. return 0, 0, err
  741. }
  742. for i, found := range existing {
  743. if found {
  744. return 0, 0, fmt.Errorf("duplicate primary key: %s", ops[i].Key)
  745. }
  746. }
  747. var uniquePlan uniqueCachePlan
  748. unique, err := m.hasUniqueIndex(table)
  749. if err != nil {
  750. return 0, 0, err
  751. }
  752. if unique {
  753. uniqueEntries := make(map[string]*overlayEntry, len(encoded))
  754. for i, row := range encoded {
  755. uniqueEntries[string(ops[i].Key)] = &overlayEntry{row: row}
  756. }
  757. uniquePlan, err = m.prepareUniqueRows(table, uniqueEntries)
  758. if err != nil {
  759. return 0, 0, err
  760. }
  761. }
  762. // Write rows in atomic BATCH_WRITE chunks. Maintain in-memory indexes only
  763. // for rows that actually persisted, so a partial failure cannot leave an
  764. // already-built index stale.
  765. var firstErr error
  766. numOK := 0
  767. for _, chunk := range chunkBatchOps(ops) {
  768. err := m.pool.WithClient(func(c *KVClient) error {
  769. _, err := c.BatchWrite(chunk, nil)
  770. return err
  771. })
  772. if err != nil {
  773. if firstErr == nil {
  774. firstErr = err
  775. }
  776. break
  777. }
  778. numOK += len(chunk)
  779. }
  780. // Maintain indexes for the rows that persisted (the first numOK ops).
  781. if len(uniquePlan) != 0 {
  782. accepted := make(map[string]bool, numOK)
  783. for i := 0; i < numOK; i++ {
  784. accepted[string(ops[i].Key)] = true
  785. }
  786. m.applyUniquePlan(uniquePlan, accepted)
  787. }
  788. for i := 0; i < numOK; i++ {
  789. if err := m.updateIndexesForRow(table, encoded[i], true); err != nil {
  790. return numOK, 0, err
  791. }
  792. }
  793. if numOK > 0 {
  794. m.bumpGeneration(table)
  795. if err := m.bumpIndexPredicates(table, encoded[:numOK]...); err != nil {
  796. return numOK, 0, err
  797. }
  798. }
  799. m.incrCount(table, schema.CreatedAt, numOK, wasInit)
  800. if firstErr != nil {
  801. m.invalidateWriteFailure(table)
  802. }
  803. var lastRowID int64
  804. if numOK > 0 {
  805. lastRowID, _ = rowIDFromRow(encoded[numOK-1])
  806. }
  807. return numOK, lastRowID, firstErr
  808. }
  809. // updateIndexesForRow adds or removes entries from already-built in-memory
  810. // indexes. Index entries are rebuildable from durable row data, so this method
  811. // intentionally does not write idx:* keys to KV.
  812. func (m *TableManager) updateIndexesForRow(table string, row Row, add bool) error {
  813. indexes, err := m.schema.ListTableIndexes(table)
  814. if err != nil || len(indexes) == 0 {
  815. return err
  816. }
  817. rowid, ok := rowIDFromRow(row)
  818. if !ok {
  819. return nil
  820. }
  821. tableKey := strings.ToLower(table)
  822. tableSchema, schemaErr := m.schema.GetSchema(table)
  823. if schemaErr == nil {
  824. m.cacheMu.Lock()
  825. if keys, initialized := m.rowKeyCache[tableKey]; initialized {
  826. if add {
  827. keys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  828. } else {
  829. delete(keys, rowid)
  830. }
  831. }
  832. m.cacheMu.Unlock()
  833. }
  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); err != nil {
  1024. return nil, false, err
  1025. }
  1026. if err := m.updateIndexesForRow(table, row, true); 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); 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.indexCache[indexKey] = values
  1283. m.indexTable[indexKey] = key
  1284. m.rowKeyCache[key] = rowKeys
  1285. }
  1286. m.cacheMu.Unlock()
  1287. return nil
  1288. }
  1289. // AddIndexEntry adds a rowid to an in-memory index entry.
  1290. func (m *TableManager) AddIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1291. indexKey := strings.ToLower(indexName)
  1292. valueKey := formatIndexValue(colValue)
  1293. m.cacheMu.Lock()
  1294. defer m.cacheMu.Unlock()
  1295. values, ok := m.indexCache[indexKey]
  1296. if !ok {
  1297. return nil
  1298. }
  1299. rowids := values[valueKey]
  1300. for _, r := range rowids {
  1301. if r == rowid {
  1302. return nil
  1303. }
  1304. }
  1305. values[valueKey] = append(rowids, rowid)
  1306. return nil
  1307. }
  1308. // RemoveIndexEntry removes a rowid from an in-memory index entry.
  1309. func (m *TableManager) RemoveIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1310. indexKey := strings.ToLower(indexName)
  1311. valueKey := formatIndexValue(colValue)
  1312. m.cacheMu.Lock()
  1313. defer m.cacheMu.Unlock()
  1314. values, ok := m.indexCache[indexKey]
  1315. if !ok {
  1316. return nil
  1317. }
  1318. rowids := values[valueKey]
  1319. newRowids := make([]int64, 0, len(rowids))
  1320. for _, r := range rowids {
  1321. if r != rowid {
  1322. newRowids = append(newRowids, r)
  1323. }
  1324. }
  1325. if len(newRowids) == 0 {
  1326. delete(values, valueKey)
  1327. return nil
  1328. }
  1329. values[valueKey] = newRowids
  1330. return nil
  1331. }
  1332. // LookupIndex returns rowids matching a column value using the index.
  1333. func (m *TableManager) LookupIndex(indexName string, colValue interface{}) ([]int64, error) {
  1334. index, err := m.schema.GetIndex(indexName)
  1335. if err != nil {
  1336. return nil, err
  1337. }
  1338. if err := m.ensureIndex(index); err != nil {
  1339. return nil, err
  1340. }
  1341. indexKey := strings.ToLower(indexName)
  1342. valueKey := formatIndexValue(colValue)
  1343. m.cacheMu.RLock()
  1344. rowids := append([]int64(nil), m.indexCache[indexKey][valueKey]...)
  1345. m.cacheMu.RUnlock()
  1346. return rowids, nil
  1347. }
  1348. // ClearIndex removes an index's in-memory entries and marks it disabled so a
  1349. // concurrent lookup cannot rebuild it after DROP but before the schema entry
  1350. // is removed. Index entries are derived from durable rows, so there is nothing
  1351. // durable to delete here.
  1352. func (m *TableManager) ClearIndex(indexName, tableName string, columns []string) error {
  1353. indexKey := strings.ToLower(indexName)
  1354. tableKey := strings.ToLower(tableName)
  1355. m.cacheMu.Lock()
  1356. delete(m.indexCache, indexKey)
  1357. delete(m.indexTable, indexKey)
  1358. delete(m.uniqueCache, indexKey)
  1359. m.disabledIndexes[indexKey] = true
  1360. rowKeysNeeded := false
  1361. for _, indexedTable := range m.indexTable {
  1362. if indexedTable == tableKey {
  1363. rowKeysNeeded = true
  1364. break
  1365. }
  1366. }
  1367. if !rowKeysNeeded {
  1368. delete(m.rowKeyCache, tableKey)
  1369. }
  1370. m.cacheMu.Unlock()
  1371. return nil
  1372. }
  1373. // BuildIndex builds index entries for all existing rows in a table.
  1374. func (m *TableManager) BuildIndex(indexName, tableName string, columns []string) error {
  1375. indexKey := strings.ToLower(indexName)
  1376. m.cacheMu.Lock()
  1377. delete(m.disabledIndexes, indexKey)
  1378. delete(m.indexCache, 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.indexCache[indexKey] = values
  1415. m.indexTable[indexKey] = strings.ToLower(tableName)
  1416. m.rowKeyCache[strings.ToLower(tableName)] = rowKeys
  1417. m.cacheMu.Unlock()
  1418. return nil
  1419. }
  1420. // indexColumnValue resolves a single index column's value for a row. A plain
  1421. // column is looked up case-insensitively; an expression column is evaluated
  1422. // through the registered expression evaluator.
  1423. func (m *TableManager) indexColumnValue(row Row, col IndexColumn) (interface{}, error) {
  1424. if col.Expression != "" {
  1425. eval := m.expressionEvaluator()
  1426. if eval == nil {
  1427. return nil, fmt.Errorf("expression index column %q requires an expression evaluator", col.Expression)
  1428. }
  1429. return eval(col.Expression, row)
  1430. }
  1431. if v, ok := row[col.Name]; ok {
  1432. return v, nil
  1433. }
  1434. for k, v := range row {
  1435. if strings.EqualFold(k, col.Name) {
  1436. return v, nil
  1437. }
  1438. }
  1439. return nil, nil
  1440. }
  1441. // buildIndexValue creates the index key value from an index's columns. Single
  1442. // column values are formatted directly; composite values are joined with a NUL
  1443. // separator.
  1444. func (m *TableManager) buildIndexValue(row Row, columns []IndexColumn) (string, error) {
  1445. if len(columns) == 1 {
  1446. v, err := m.indexColumnValue(row, columns[0])
  1447. if err != nil {
  1448. return "", err
  1449. }
  1450. return formatIndexValue(v), nil
  1451. }
  1452. parts := make([]string, 0, len(columns))
  1453. for _, col := range columns {
  1454. v, err := m.indexColumnValue(row, col)
  1455. if err != nil {
  1456. return "", err
  1457. }
  1458. parts = append(parts, formatIndexValue(v))
  1459. }
  1460. return strings.Join(parts, "\x00"), nil
  1461. }
  1462. // ── UNIQUE index enforcement ────────────────────────────────────────────────
  1463. //
  1464. // Uniqueness is enforced by scanning durable rows rather than by persisting a
  1465. // separate claim key. There is therefore no storage migration and no new durable
  1466. // state: a table with a UNIQUE index takes its exclusive table gate for writes so
  1467. // a validating scan can never race a concurrent writer. NULL values are exempt
  1468. // (SQLite permits multiple NULLs in a unique index), and composite values are
  1469. // encoded with length prefixes so "a\x00b" in one column can never collide with
  1470. // "a", "b" across two columns.
  1471. // uniqueIndexes returns the UNIQUE indexes defined on a table.
  1472. func (m *TableManager) uniqueIndexes(table string) ([]*Index, error) {
  1473. indexes, err := m.schema.ListTableIndexes(table)
  1474. if err != nil {
  1475. return nil, err
  1476. }
  1477. var uniq []*Index
  1478. for _, idx := range indexes {
  1479. if idx.Unique {
  1480. uniq = append(uniq, idx)
  1481. }
  1482. }
  1483. return uniq, nil
  1484. }
  1485. // hasUniqueIndex reports whether the table has any UNIQUE index. An IO error is
  1486. // surfaced so callers never silently skip uniqueness enforcement.
  1487. func (m *TableManager) hasUniqueIndex(table string) (bool, error) {
  1488. indexes, err := m.uniqueIndexes(table)
  1489. if err != nil {
  1490. return false, err
  1491. }
  1492. return len(indexes) > 0, nil
  1493. }
  1494. // HasUniqueIndex reports whether the table has any UNIQUE index, surfacing IO
  1495. // errors. It is the exported form used by the executor to decide whether to run a
  1496. // DML statement inside an implicit transaction for statement-level atomicity.
  1497. func (m *TableManager) HasUniqueIndex(table string) (bool, error) {
  1498. return m.hasUniqueIndex(table)
  1499. }
  1500. // encodeUniqueValue encodes the indexed value of a row with per-column type tags
  1501. // and length prefixes. It returns isNull=true when any indexed column is NULL, in
  1502. // which case the row is exempt from uniqueness. The length prefix makes composite
  1503. // encodings collision-free. Expression columns are evaluated through the
  1504. // registered evaluator.
  1505. func (m *TableManager) encodeUniqueValue(row Row, columns []IndexColumn) (string, bool, error) {
  1506. var sb strings.Builder
  1507. for _, col := range columns {
  1508. v, err := m.indexColumnValue(row, col)
  1509. if err != nil {
  1510. return "", false, err
  1511. }
  1512. if v == nil {
  1513. return "", true, nil
  1514. }
  1515. s := encodeUniqueScalar(v)
  1516. sb.WriteString(strconv.Itoa(len(s)))
  1517. sb.WriteByte(':')
  1518. sb.WriteString(s)
  1519. }
  1520. return sb.String(), false, nil
  1521. }
  1522. // encodeUniqueScalar encodes a scalar for uniqueness comparison. Integral
  1523. // numerics are canonicalized across their integer/float/unsigned Go
  1524. // representations so a computed value such as 1.5-0.5 (float64) collides with
  1525. // the literal 1 (int64) the same way SQLite's numeric affinity does. Non-integral
  1526. // reals keep a full-precision tag so no distinct value is lost. TEXT and BLOB
  1527. // remain distinct from numbers.
  1528. func encodeUniqueScalar(v interface{}) string {
  1529. switch t := v.(type) {
  1530. case bool:
  1531. if t {
  1532. return "b1"
  1533. }
  1534. return "b0"
  1535. case int:
  1536. return "i" + strconv.FormatInt(int64(t), 10)
  1537. case int64:
  1538. return "i" + strconv.FormatInt(t, 10)
  1539. case uint:
  1540. return "i" + strconv.FormatUint(uint64(t), 10)
  1541. case uint8:
  1542. return "i" + strconv.FormatUint(uint64(t), 10)
  1543. case uint16:
  1544. return "i" + strconv.FormatUint(uint64(t), 10)
  1545. case uint32:
  1546. return "i" + strconv.FormatUint(uint64(t), 10)
  1547. case uint64:
  1548. return "i" + strconv.FormatUint(t, 10)
  1549. case uintptr:
  1550. return "i" + strconv.FormatUint(uint64(t), 10)
  1551. case float64:
  1552. if s, ok := encodeIntegralFloat(t); ok {
  1553. return s
  1554. }
  1555. return "r" + strconv.FormatFloat(t, 'g', -1, 64)
  1556. case string:
  1557. return "s" + t
  1558. case []byte:
  1559. return "x" + string(t)
  1560. default:
  1561. return "s" + fmt.Sprintf("%v", t)
  1562. }
  1563. }
  1564. // encodeIntegralFloat canonicalizes a whole float64 to the same "i" encoding used
  1565. // by integer values, without precision loss, so integral reals and integers are
  1566. // not treated as distinct under a unique index. It reports false (and returns
  1567. // nothing) for non-integral values or values outside the int64 range.
  1568. func encodeIntegralFloat(v float64) (string, bool) {
  1569. if v != math.Trunc(v) {
  1570. return "", false
  1571. }
  1572. if v < -9223372036854775808.0 || v >= 9223372036854775808.0 {
  1573. return "", false
  1574. }
  1575. return "i" + strconv.FormatInt(int64(v), 10), true
  1576. }
  1577. // lockForWrite acquires the table gate appropriate for a point write: exclusive
  1578. // when the table has a UNIQUE index (so the validating scan cannot race another
  1579. // writer), shared otherwise. The uniqueness decision is re-checked under the lock
  1580. // so a concurrent CREATE UNIQUE INDEX cannot slip in between the unlocked probe
  1581. // and the lock acquisition; the returned unlock closes the gate.
  1582. func (m *TableManager) lockForWrite(table string) (unlock func(), unique bool, err error) {
  1583. tl := m.tableLock(table)
  1584. for {
  1585. uniq, err := m.hasUniqueIndex(table)
  1586. if err != nil {
  1587. return nil, false, err
  1588. }
  1589. if uniq {
  1590. tl.Lock()
  1591. // Re-check under the exclusive gate (authoritative): CreateUniqueIndex
  1592. // also takes the exclusive gate, so it cannot run while we hold it.
  1593. uniq, err = m.hasUniqueIndex(table)
  1594. if err != nil {
  1595. tl.Unlock()
  1596. return nil, false, err
  1597. }
  1598. if uniq {
  1599. return tl.Unlock, true, nil
  1600. }
  1601. tl.Unlock()
  1602. continue
  1603. }
  1604. tl.RLock()
  1605. // A unique index could have appeared before we acquired the shared gate;
  1606. // re-check and escalate if so.
  1607. uniq, err = m.hasUniqueIndex(table)
  1608. if err != nil {
  1609. tl.RUnlock()
  1610. return nil, false, err
  1611. }
  1612. if uniq {
  1613. tl.RUnlock()
  1614. continue
  1615. }
  1616. return tl.RUnlock, false, nil
  1617. }
  1618. }
  1619. // ValidateUniqueIndex rejects a UNIQUE index definition if existing rows already
  1620. // contain duplicate non-NULL values for the indexed columns.
  1621. func (m *TableManager) ValidateUniqueIndex(index *Index) error {
  1622. seen := make(map[string]bool)
  1623. return m.scanRows(index.Table, func(row Row) (bool, error) {
  1624. encoded, isNull, err := m.encodeUniqueValue(row, index.Columns)
  1625. if err != nil {
  1626. return true, err
  1627. }
  1628. if isNull {
  1629. return false, nil
  1630. }
  1631. if seen[encoded] {
  1632. return true, &UniqueViolationError{Constraint: index.Name}
  1633. }
  1634. seen[encoded] = true
  1635. return false, nil
  1636. })
  1637. }
  1638. // CreateUniqueIndex validates existing rows and registers the index while
  1639. // holding the table's exclusive gate, so a concurrent writer cannot insert a
  1640. // conflicting value between validation and registration.
  1641. func (m *TableManager) CreateUniqueIndex(index *Index) error {
  1642. tl := m.tableLock(index.Table)
  1643. tl.Lock()
  1644. defer tl.Unlock()
  1645. if err := m.ValidateUniqueIndex(index); err != nil {
  1646. return err
  1647. }
  1648. return m.schema.CreateIndex(index)
  1649. }
  1650. type indexedRowVersion struct {
  1651. row Row
  1652. key string
  1653. lsn uint64
  1654. }
  1655. type indexPredicateSnapshot struct {
  1656. valueKey string
  1657. valueGen uint64
  1658. wildcardKey string
  1659. wildcardGen uint64
  1660. }
  1661. // selectByIndexWithLSN retrieves indexed rows and their durable versions. The
  1662. // transaction layer uses the versions for optimistic commit validation.
  1663. func (m *TableManager) selectByIndexWithLSN(table, indexName string, colValue interface{}) ([]indexedRowVersion, indexPredicateSnapshot, error) {
  1664. return m.selectByIndexKeyWithLSN(table, indexName, formatIndexValue(colValue))
  1665. }
  1666. // IndexRowKey computes the formatted in-memory index key a row maps to for an
  1667. // index. Expression columns are evaluated through the same registered evaluator
  1668. // used by uniqueness enforcement, so the key matches both durable index entries
  1669. // and staged overlay rows.
  1670. func (m *TableManager) IndexRowKey(idx *Index, row Row) (string, error) {
  1671. value, err := m.buildIndexValue(row, idx.Columns)
  1672. if err != nil {
  1673. return "", err
  1674. }
  1675. return formatIndexValue(value), nil
  1676. }
  1677. // IndexValueContainsNull reports whether any column contributing to an index
  1678. // key evaluates to NULL. UNIQUE indexes exempt such rows from conflicts under
  1679. // SQLite semantics.
  1680. func (m *TableManager) IndexValueContainsNull(idx *Index, row Row) (bool, error) {
  1681. _, isNull, err := m.encodeUniqueValue(row, idx.Columns)
  1682. return isNull, err
  1683. }
  1684. // selectByIndexKeyWithLSN is selectByIndexWithLSN with an already-computed
  1685. // formatted index key, so callers can look up composite and expression indexes
  1686. // without re-deriving the key from a single column value.
  1687. func (m *TableManager) selectByIndexKeyWithLSN(table, indexName, valueKey string) ([]indexedRowVersion, indexPredicateSnapshot, error) {
  1688. index, err := m.schema.GetIndex(indexName)
  1689. if err != nil {
  1690. return nil, indexPredicateSnapshot{}, err
  1691. }
  1692. if err := m.ensureIndex(index); err != nil {
  1693. return nil, indexPredicateSnapshot{}, err
  1694. }
  1695. tableKey := strings.ToLower(table)
  1696. tl := m.tableLock(tableKey)
  1697. tl.RLock()
  1698. defer tl.RUnlock()
  1699. if !m.schema.TableExists(table) {
  1700. return nil, indexPredicateSnapshot{}, fmt.Errorf("table not found: %s", table)
  1701. }
  1702. indexKey := strings.ToLower(indexName)
  1703. predicateKey, predicateGen, wildcardKey, wildcardGen := m.predicateSnapshot(table, indexName, valueKey)
  1704. snapshot := indexPredicateSnapshot{
  1705. valueKey: predicateKey, valueGen: predicateGen,
  1706. wildcardKey: wildcardKey, wildcardGen: wildcardGen,
  1707. }
  1708. m.cacheMu.RLock()
  1709. rowids := append([]int64(nil), m.indexCache[indexKey][valueKey]...)
  1710. primaryKeys := make([]string, 0, len(rowids))
  1711. missingRowID := int64(0)
  1712. missingRowKey := false
  1713. for _, rowid := range rowids {
  1714. if primaryKey, ok := m.rowKeyCache[tableKey][rowid]; ok {
  1715. primaryKeys = append(primaryKeys, primaryKey)
  1716. } else {
  1717. missingRowID = rowid
  1718. missingRowKey = true
  1719. break
  1720. }
  1721. }
  1722. m.cacheMu.RUnlock()
  1723. if missingRowKey {
  1724. return nil, indexPredicateSnapshot{}, fmt.Errorf("index %s is missing rowid %d", indexName, missingRowID)
  1725. }
  1726. // If no rowids found, return empty result
  1727. if len(rowids) == 0 {
  1728. return []indexedRowVersion{}, snapshot, nil
  1729. }
  1730. rows := make([]indexedRowVersion, 0, len(primaryKeys))
  1731. err = m.pool.WithClient(func(client *KVClient) error {
  1732. keys := make([][]byte, len(primaryKeys))
  1733. for i, primaryKey := range primaryKeys {
  1734. keys[i] = []byte(m.dataKey(table, primaryKey))
  1735. }
  1736. results, err := client.MultiGet(keys)
  1737. if err != nil {
  1738. return err
  1739. }
  1740. for i, result := range results {
  1741. if !result.Found {
  1742. primaryKey := primaryKeys[i]
  1743. return fmt.Errorf("index %s references missing primary key %s", indexName, primaryKey)
  1744. }
  1745. row, err := decodeRow(result.Value)
  1746. if err != nil {
  1747. return err
  1748. }
  1749. rows = append(rows, indexedRowVersion{
  1750. row: row,
  1751. key: m.dataKey(table, primaryKeys[i]),
  1752. lsn: result.LSN,
  1753. })
  1754. }
  1755. return nil
  1756. })
  1757. if err != nil {
  1758. return nil, indexPredicateSnapshot{}, err
  1759. }
  1760. return rows, snapshot, nil
  1761. }
  1762. // SelectByIndex retrieves rows using an in-memory equality index followed by a
  1763. // single MultiGet for the matching primary keys.
  1764. func (m *TableManager) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
  1765. versions, _, err := m.selectByIndexWithLSN(table, indexName, colValue)
  1766. if err != nil {
  1767. return nil, err
  1768. }
  1769. rows := make([]Row, len(versions))
  1770. for i := range versions {
  1771. rows[i] = versions[i].row
  1772. }
  1773. return rows, nil
  1774. }