2
0

table.go 43 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669
  1. package storage
  2. import (
  3. "fmt"
  4. "hash/fnv"
  5. "math"
  6. "strings"
  7. "sync"
  8. "time"
  9. )
  10. // Row represents a database row.
  11. type Row map[string]interface{}
  12. // TableManager manages table data operations.
  13. type TableManager struct {
  14. pool *KVPool
  15. schema *SchemaManager
  16. database string
  17. cacheMu sync.RWMutex
  18. indexCache map[string]map[string][]int64 // index name → indexed value → rowids
  19. indexTable map[string]string // index name → table name
  20. rowKeyCache map[string]map[int64]string // table name → rowid → primary key
  21. // disabledIndexes prevents a concurrent lookup from rebuilding an index
  22. // after DROP has cleared it but before the schema entry is removed.
  23. disabledIndexes map[string]bool
  24. // counts holds exact per-table row counts for the COUNT(*) fast path.
  25. // It is derived lazily from durable rows on first use and maintained
  26. // incrementally by Insert/InsertBulk/Delete thereafter.
  27. counts map[string]int
  28. countsInit map[string]bool
  29. countGeneration map[string]time.Time
  30. // stripes are deterministic per-key locks used by point operations
  31. // (GetByPK/Insert/UpdateByPK/DeleteByPK). They replace the table-wide lock
  32. // so point operations on different keys of the same table proceed
  33. // concurrently. Index is derived from the full data key (database+table+pk).
  34. stripes [64]sync.Mutex
  35. // generations tracks full-table scans. predicateGenerations narrows indexed
  36. // equality reads to one index value so unrelated writes do not conflict.
  37. genMu sync.Mutex
  38. generations map[string]uint64
  39. predicateGenerations map[string]uint64
  40. }
  41. // NewTableManager creates a new table manager.
  42. func NewTableManager(pool *KVPool, schema *SchemaManager, database string) *TableManager {
  43. return &TableManager{
  44. pool: pool,
  45. schema: schema,
  46. database: database,
  47. indexCache: make(map[string]map[string][]int64),
  48. indexTable: make(map[string]string),
  49. rowKeyCache: make(map[string]map[int64]string),
  50. disabledIndexes: make(map[string]bool),
  51. counts: make(map[string]int),
  52. countsInit: make(map[string]bool),
  53. countGeneration: make(map[string]time.Time),
  54. generations: make(map[string]uint64),
  55. predicateGenerations: make(map[string]uint64),
  56. }
  57. }
  58. func (m *TableManager) tableLock(key string) *sync.RWMutex {
  59. return m.schema.tableLock(key)
  60. }
  61. // stripeKey returns the deterministic striped lock for a point operation on the
  62. // given full data key. Point operations on different keys therefore serialize
  63. // independently, while operations on the same key are mutually exclusive.
  64. func (m *TableManager) stripeKey(key string) *sync.Mutex {
  65. h := fnv.New32a()
  66. h.Write([]byte(key))
  67. return &m.stripes[h.Sum32()%uint32(len(m.stripes))]
  68. }
  69. // generation returns the current in-process generation for a table. It is
  70. // bumped on every committed write and captured by transaction scans.
  71. func (m *TableManager) generation(table string) uint64 {
  72. key := strings.ToLower(table)
  73. m.genMu.Lock()
  74. g := m.generations[key]
  75. m.genMu.Unlock()
  76. return g
  77. }
  78. // bumpGeneration advances a table's generation. Callers hold the table gate in
  79. // shared mode for point writes or exclusive mode for scan-based writes.
  80. func (m *TableManager) bumpGeneration(table string) {
  81. key := strings.ToLower(table)
  82. m.genMu.Lock()
  83. m.generations[key]++
  84. m.genMu.Unlock()
  85. }
  86. func indexPredicateKey(table, indexName, value string) string {
  87. return strings.ToLower(table) + "\x00" + strings.ToLower(indexName) + "\x00" + value
  88. }
  89. func indexPredicateWildcardKey(table string) string {
  90. return strings.ToLower(table) + "\x00*"
  91. }
  92. func (m *TableManager) predicateSnapshot(table, indexName, value string) (string, uint64, string, uint64) {
  93. valueKey := indexPredicateKey(table, indexName, value)
  94. wildcardKey := indexPredicateWildcardKey(table)
  95. m.genMu.Lock()
  96. valueGen := m.predicateGenerations[valueKey]
  97. wildcardGen := m.predicateGenerations[wildcardKey]
  98. m.genMu.Unlock()
  99. return valueKey, valueGen, wildcardKey, wildcardGen
  100. }
  101. func (m *TableManager) predicateGeneration(key string) uint64 {
  102. m.genMu.Lock()
  103. gen := m.predicateGenerations[key]
  104. m.genMu.Unlock()
  105. return gen
  106. }
  107. func (m *TableManager) bumpIndexPredicates(table string, rows ...Row) {
  108. indexes, err := m.schema.ListTableIndexes(table)
  109. if err != nil || len(indexes) == 0 {
  110. return
  111. }
  112. m.genMu.Lock()
  113. defer m.genMu.Unlock()
  114. for _, row := range rows {
  115. if row == nil {
  116. continue
  117. }
  118. for _, index := range indexes {
  119. columns := make([]string, len(index.Columns))
  120. for i, column := range index.Columns {
  121. columns[i] = column.Name
  122. }
  123. value := formatIndexValue(m.buildIndexValue(row, columns))
  124. m.predicateGenerations[indexPredicateKey(table, index.Name, value)]++
  125. }
  126. }
  127. }
  128. func (m *TableManager) bumpIndexPredicateWildcard(table string) {
  129. m.genMu.Lock()
  130. m.predicateGenerations[indexPredicateWildcardKey(table)]++
  131. m.genMu.Unlock()
  132. }
  133. // compareWritePoint issues a single CompareBatchWrite against the pooled KV.
  134. // It returns whether the compare checks held and the ops committed.
  135. func (m *TableManager) compareWritePoint(checks []CompareCheck, ops []BatchOp) (bool, error) {
  136. var committed bool
  137. err := m.pool.WithClient(func(c *KVClient) error {
  138. _, ok, err := c.CompareBatchWrite(checks, ops, nil)
  139. committed = ok
  140. return err
  141. })
  142. return committed, err
  143. }
  144. // invalidateCache removes a table's derived in-memory indexes.
  145. func (m *TableManager) invalidateCache(table string) {
  146. m.cacheMu.Lock()
  147. key := strings.ToLower(table)
  148. delete(m.rowKeyCache, key)
  149. delete(m.counts, key)
  150. delete(m.countsInit, key)
  151. delete(m.countGeneration, key)
  152. for indexName, tableName := range m.indexTable {
  153. if tableName == key {
  154. delete(m.indexCache, indexName)
  155. delete(m.indexTable, indexName)
  156. }
  157. }
  158. m.cacheMu.Unlock()
  159. }
  160. // InvalidateCache is the exported version for use by the executor.
  161. func (m *TableManager) InvalidateCache(table string) {
  162. m.invalidateCache(table)
  163. }
  164. // rowVisitFunc is invoked for each decoded row in a streaming scan. Return
  165. // stop=true to end the scan early; a non-nil error aborts the scan.
  166. type rowVisitFunc func(Row) (stop bool, err error)
  167. // scanRows streams the rows of table by scanning durable KV rows one page at a
  168. // time under a single pooled client. Each page is decoded as it arrives and
  169. // passed to fn, which may stop the scan early. The cursor is always closed and
  170. // the client always returned to the pool, even on error.
  171. func (m *TableManager) scanRows(table string, fn rowVisitFunc) error {
  172. return m.scanRowsWithPageSize(table, scanPageSize, fn)
  173. }
  174. func (m *TableManager) scanRowsWithPageSize(table string, pageSize uint32, fn rowVisitFunc) error {
  175. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  176. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), pageSize)
  177. if err != nil {
  178. return err
  179. }
  180. defer func() {
  181. if err := cursor.Close(); retErr == nil {
  182. retErr = err
  183. }
  184. }()
  185. for {
  186. entries, done, err := cursor.Next()
  187. if err != nil {
  188. return err
  189. }
  190. for _, e := range entries {
  191. row, err := decodeRow(e.Value)
  192. if err != nil {
  193. return err
  194. }
  195. stop, err := fn(row)
  196. if err != nil {
  197. return err
  198. }
  199. if stop {
  200. return nil
  201. }
  202. }
  203. if done {
  204. return nil
  205. }
  206. }
  207. })
  208. }
  209. // rowWithLSNVisitFunc is like rowVisitFunc but also passes the durable row LSN.
  210. type rowWithLSNVisitFunc func(row Row, lsn uint64) (stop bool, err error)
  211. // scanRowsWithLSN streams a table's rows together with their durable LSNs. It
  212. // is used by buffered transactions to capture a per-row read set for
  213. // compare-and-swap validation at commit.
  214. func (m *TableManager) scanRowsWithLSN(table string, fn rowWithLSNVisitFunc) error {
  215. return m.pool.WithClient(func(client *KVClient) (retErr error) {
  216. cursor, err := client.ScanWithLimit([]byte(m.dataPrefix(table)), scanPageSize)
  217. if err != nil {
  218. return err
  219. }
  220. defer func() {
  221. if err := cursor.Close(); retErr == nil {
  222. retErr = err
  223. }
  224. }()
  225. for {
  226. entries, done, err := cursor.Next()
  227. if err != nil {
  228. return err
  229. }
  230. for _, e := range entries {
  231. row, err := decodeRow(e.Value)
  232. if err != nil {
  233. return err
  234. }
  235. stop, err := fn(row, e.LSN)
  236. if err != nil {
  237. return err
  238. }
  239. if stop {
  240. return nil
  241. }
  242. }
  243. if done {
  244. return nil
  245. }
  246. }
  247. })
  248. }
  249. // scanCountKeys counts the durable rows of table using a key-only scan so row
  250. // values are never pulled across the wire. It is used for first-time COUNT(*)
  251. // derivation.
  252. func (m *TableManager) scanCountKeys(table string) (int, error) {
  253. count := 0
  254. err := m.pool.WithClient(func(client *KVClient) (retErr error) {
  255. cursor, err := client.ScanKeys([]byte(m.dataPrefix(table)))
  256. if err != nil {
  257. return err
  258. }
  259. defer func() {
  260. if err := cursor.Close(); retErr == nil {
  261. retErr = err
  262. }
  263. }()
  264. for {
  265. entries, done, err := cursor.Next()
  266. if err != nil {
  267. return err
  268. }
  269. count += len(entries)
  270. if done {
  271. return nil
  272. }
  273. }
  274. })
  275. return count, err
  276. }
  277. // CountFast returns the exact number of rows in a table. The count is derived
  278. // from durable rows on first use (recovering across restarts) and then
  279. // maintained incrementally by the write paths, so repeated COUNT(*) queries
  280. // avoid a full table scan. It intentionally does not persist a counter to KV:
  281. // the KV layer has no atomic increment primitive, and a durable counter that
  282. // could diverge from the rows on crash would be worse than a lazily-derived,
  283. // always-exact value. The cost is one key-only scan the first time COUNT(*) is
  284. // issued after startup.
  285. func (m *TableManager) CountFast(table string) (int, error) {
  286. key := strings.ToLower(table)
  287. tableSchema, err := m.schema.GetSchema(table)
  288. if err != nil {
  289. return 0, err
  290. }
  291. // Cached read path: if the count is initialized for the current table
  292. // generation, return it without taking any table lock.
  293. m.cacheMu.Lock()
  294. if m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  295. if m.countsInit[key] {
  296. n := m.counts[key]
  297. m.cacheMu.Unlock()
  298. return n, nil
  299. }
  300. } else {
  301. delete(m.counts, key)
  302. delete(m.countsInit, key)
  303. m.countGeneration[key] = tableSchema.CreatedAt
  304. }
  305. m.cacheMu.Unlock()
  306. // First derivation: hold the table write lock so the key-only scan is
  307. // exact against concurrent writes.
  308. tl := m.tableLock(key)
  309. tl.Lock()
  310. defer tl.Unlock()
  311. m.cacheMu.Lock()
  312. if !m.countGeneration[key].Equal(tableSchema.CreatedAt) {
  313. delete(m.counts, key)
  314. delete(m.countsInit, key)
  315. m.countGeneration[key] = tableSchema.CreatedAt
  316. }
  317. if m.countsInit[key] {
  318. n := m.counts[key]
  319. m.cacheMu.Unlock()
  320. return n, nil
  321. }
  322. m.cacheMu.Unlock()
  323. count, err := m.scanCountKeys(table)
  324. if err != nil {
  325. return 0, err
  326. }
  327. m.cacheMu.Lock()
  328. m.counts[key] = count
  329. m.countsInit[key] = true
  330. m.cacheMu.Unlock()
  331. return count, nil
  332. }
  333. // countInitialized reports whether the derived count cache is initialized for
  334. // the table at this instant. Write paths capture it before their KV write so a
  335. // concurrent first derivation does not double-count the just-written row.
  336. func (m *TableManager) countInitialized(table string) bool {
  337. key := strings.ToLower(table)
  338. m.cacheMu.Lock()
  339. init := m.countsInit[key]
  340. m.cacheMu.Unlock()
  341. return init
  342. }
  343. // incrCount adjusts the derived per-table row count. wasInit reports whether
  344. // the count was already initialized before the corresponding write, so a count
  345. // that was not yet initialized is left to be re-derived from durable rows
  346. // (which already reflect the write) on next use.
  347. func (m *TableManager) incrCount(table string, generation time.Time, delta int, wasInit bool) {
  348. key := strings.ToLower(table)
  349. m.cacheMu.Lock()
  350. if !m.countGeneration[key].Equal(generation) {
  351. delete(m.counts, key)
  352. delete(m.countsInit, key)
  353. m.countGeneration[key] = generation
  354. }
  355. if wasInit && m.countsInit[key] {
  356. m.counts[key] += delta
  357. }
  358. if !wasInit {
  359. // The count was not initialized before the write, so a concurrent first
  360. // derivation may have missed the just-written row. Invalidate to force
  361. // an exact re-derivation on the next COUNT(*).
  362. delete(m.countsInit, key)
  363. }
  364. m.cacheMu.Unlock()
  365. }
  366. // dataKey returns the key for a row.
  367. func (m *TableManager) dataKey(table, pk string) string {
  368. return fmt.Sprintf("%s:_data:%s:%s", m.database, strings.ToLower(table), pk)
  369. }
  370. // dataPrefix returns the prefix for all rows in a table.
  371. func (m *TableManager) dataPrefix(table string) string {
  372. return fmt.Sprintf("%s:_data:%s:", m.database, strings.ToLower(table))
  373. }
  374. // prepareInsert validates and normalizes an insert row, generating the ROWID
  375. // when required. It returns the normalized row and the full data key without
  376. // writing anything, so both the autocommit path (compare-and-swap) and the
  377. // buffered transaction path (staging) can share it. The input row has its
  378. // primary key populated as a side effect.
  379. func (m *TableManager) prepareInsert(table string, row Row) (Row, string, error) {
  380. schema, err := m.schema.GetSchema(table)
  381. if err != nil {
  382. return nil, "", err
  383. }
  384. // Get primary key value
  385. pkValue, ok := row[schema.PrimaryKey]
  386. if !ok {
  387. for k, v := range row {
  388. if strings.EqualFold(k, schema.PrimaryKey) {
  389. pkValue = v
  390. ok = true
  391. break
  392. }
  393. }
  394. }
  395. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  396. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  397. var rowid int64
  398. if !ok || pkValue == nil {
  399. if isIntegerPK || !ok {
  400. rowid, err = m.schema.GetNextRowID(table)
  401. if err != nil {
  402. return nil, "", err
  403. }
  404. pkValue = rowid
  405. row[schema.PrimaryKey] = rowid
  406. ok = true
  407. } else {
  408. return nil, "", fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  409. }
  410. } else if isIntegerPK {
  411. switch v := pkValue.(type) {
  412. case int64:
  413. rowid = v
  414. case float64:
  415. if math.Trunc(v) != v {
  416. return nil, "", fmt.Errorf("invalid integer primary key: %v", v)
  417. }
  418. rowid = int64(v)
  419. case int:
  420. rowid = int64(v)
  421. default:
  422. rowid = 0
  423. }
  424. if rowid > 0 {
  425. if err := m.schema.UpdateMaxRowID(table, rowid); err != nil {
  426. return nil, "", err
  427. }
  428. }
  429. }
  430. pk := fmt.Sprintf("%v", pkValue)
  431. for _, col := range schema.Columns {
  432. if !col.Nullable && col.Default == nil {
  433. val, hasVal := row[col.Name]
  434. if !hasVal {
  435. for k, v := range row {
  436. if strings.EqualFold(k, col.Name) {
  437. val = v
  438. hasVal = true
  439. break
  440. }
  441. }
  442. }
  443. if !hasVal || val == nil {
  444. return nil, "", fmt.Errorf("missing required column: %s", col.Name)
  445. }
  446. }
  447. }
  448. normalizedRow := make(Row)
  449. for _, col := range schema.Columns {
  450. for k, v := range row {
  451. if strings.EqualFold(k, col.Name) {
  452. normalizedRow[col.Name] = v
  453. break
  454. }
  455. }
  456. }
  457. for _, col := range schema.Columns {
  458. if _, ok := normalizedRow[col.Name]; !ok && col.Default != nil {
  459. normalizedRow[col.Name] = col.Default
  460. }
  461. }
  462. if rowid > 0 {
  463. normalizedRow["_rowid_"] = rowid
  464. } else {
  465. newRowID, err := m.schema.GetNextRowID(table)
  466. if err != nil {
  467. return nil, "", err
  468. }
  469. normalizedRow["_rowid_"] = newRowID
  470. }
  471. return normalizedRow, m.dataKey(table, pk), nil
  472. }
  473. // Insert inserts a new row. The duplicate check and write are one atomic
  474. // compare-and-swap so concurrent inserts of the same primary key cannot both
  475. // persist. The shared table gate is acquired before the key's striped lock so
  476. // a queued scan writer cannot invert the lock order with point operations.
  477. func (m *TableManager) Insert(table string, row Row) error {
  478. nr, key, err := m.prepareInsert(table, row)
  479. if err != nil {
  480. return err
  481. }
  482. data, err := encodeRow(nr)
  483. if err != nil {
  484. return fmt.Errorf("failed to serialize row: %w", err)
  485. }
  486. wasInit := m.countInitialized(table)
  487. // Point writers share this gate with each other. Transaction commits and
  488. // scan-based writes take it exclusively, so generation validation and cache
  489. // publication are ordered without serializing writes to different keys.
  490. tl := m.tableLock(table)
  491. tl.RLock()
  492. defer tl.RUnlock()
  493. st := m.stripeKey(key)
  494. st.Lock()
  495. defer st.Unlock()
  496. committed, err := m.compareWritePoint(
  497. []CompareCheck{{Key: []byte(key), LSN: 0}},
  498. []BatchOp{{Op: batchPut, Key: []byte(key), Value: data}},
  499. )
  500. if err != nil {
  501. return err
  502. }
  503. if !committed {
  504. return fmt.Errorf("duplicate primary key: %v", row[schemaPrimaryKey(m.schema, table)])
  505. }
  506. m.updateIndexesForRow(table, nr, true)
  507. // Publish derived index state before its generations. A reader that races
  508. // with publication either sees the old generation and aborts or sees the
  509. // complete new state.
  510. m.bumpIndexPredicates(table, nr)
  511. m.bumpGeneration(table)
  512. if schema, serr := m.schema.GetSchema(table); serr == nil {
  513. m.incrCount(table, schema.CreatedAt, 1, wasInit)
  514. }
  515. return nil
  516. }
  517. func schemaPrimaryKey(s *SchemaManager, table string) string {
  518. schema, err := s.GetSchema(table)
  519. if err != nil {
  520. return "_rowid_"
  521. }
  522. return schema.PrimaryKey
  523. }
  524. // bulkBatchByteBudget bounds a single atomic BATCH_WRITE payload below the
  525. // PKBFI frame limit so a bulk insert never emits a frame the server rejects.
  526. // Each op contributes 12 header bytes plus its key and value.
  527. const bulkBatchByteBudget = 60 * 1024 * 1024
  528. // chunkBatchOps splits ops into atomic BATCH_WRITE chunks bounded by both the
  529. // PKBFI operation-count limit and the frame-size limit. Each chunk is a slice
  530. // of the backing array, valid until the next append to ops.
  531. func chunkBatchOps(ops []BatchOp) [][]BatchOp {
  532. var chunks [][]BatchOp
  533. for i := 0; i < len(ops); {
  534. end := i + maxOperations
  535. if end > len(ops) {
  536. end = len(ops)
  537. }
  538. bytes := 0
  539. j := i
  540. for j < end {
  541. sz := 12 + len(ops[j].Key) + len(ops[j].Value)
  542. if j > i && bytes+sz > bulkBatchByteBudget {
  543. break
  544. }
  545. bytes += sz
  546. j++
  547. }
  548. if j == i {
  549. j = i + 1
  550. }
  551. chunks = append(chunks, ops[i:j])
  552. i = j
  553. }
  554. return chunks
  555. }
  556. // InsertBulk inserts multiple rows efficiently using atomic BATCH_WRITE chunks
  557. // bounded by the PKBFI operation-count and frame-size limits. Skips per-row
  558. // duplicate checks (caller must ensure uniqueness). Used by INSERT ... SELECT.
  559. func (m *TableManager) InsertBulk(table string, rows []Row) (int, error) {
  560. if len(rows) == 0 {
  561. return 0, nil
  562. }
  563. tl := m.tableLock(table)
  564. tl.Lock()
  565. defer tl.Unlock()
  566. schema, err := m.schema.GetSchema(table)
  567. if err != nil {
  568. return 0, err
  569. }
  570. wasInit := m.countInitialized(table)
  571. pkCol, _ := schema.GetColumn(schema.PrimaryKey)
  572. isIntegerPK := pkCol != nil && isIntegerType(pkCol.Type)
  573. // Normalize rows and assign _rowid_.
  574. normalized := make([]Row, 0, len(rows))
  575. var maxRowID int64
  576. for _, row := range rows {
  577. nr := make(Row)
  578. for _, col := range schema.Columns {
  579. for k, v := range row {
  580. if strings.EqualFold(k, col.Name) {
  581. nr[col.Name] = v
  582. break
  583. }
  584. }
  585. }
  586. for _, col := range schema.Columns {
  587. if _, ok := nr[col.Name]; !ok && col.Default != nil {
  588. nr[col.Name] = col.Default
  589. }
  590. }
  591. var rowid int64
  592. var hasRowid bool
  593. if isIntegerPK {
  594. switch v := nr[schema.PrimaryKey].(type) {
  595. case float64:
  596. if math.Trunc(v) != v {
  597. return 0, fmt.Errorf("invalid integer primary key: %v", v)
  598. }
  599. rowid = int64(v)
  600. hasRowid = true
  601. case int64:
  602. rowid = v
  603. hasRowid = true
  604. case int:
  605. rowid = int64(v)
  606. hasRowid = true
  607. }
  608. }
  609. if !hasRowid {
  610. if schema.PrimaryKey != "_rowid_" {
  611. pk, ok := nr[schema.PrimaryKey]
  612. if !ok || pk == nil {
  613. return 0, fmt.Errorf("missing primary key: %s", schema.PrimaryKey)
  614. }
  615. }
  616. rowid, err = m.schema.GetNextRowID(table)
  617. if err != nil {
  618. return 0, err
  619. }
  620. if schema.PrimaryKey == "_rowid_" {
  621. nr[schema.PrimaryKey] = rowid
  622. }
  623. }
  624. nr["_rowid_"] = rowid
  625. if rowid > maxRowID {
  626. maxRowID = rowid
  627. }
  628. normalized = append(normalized, nr)
  629. }
  630. if maxRowID > 0 {
  631. m.schema.UpdateMaxRowID(table, maxRowID)
  632. }
  633. // Serialize all rows into batch operations. A parallel slice keeps the
  634. // normalized Row for each op for index maintenance after the write.
  635. ops := make([]BatchOp, 0, len(normalized))
  636. encoded := make([]Row, 0, len(normalized))
  637. for _, nr := range normalized {
  638. pk := fmt.Sprintf("%v", nr[schema.PrimaryKey])
  639. data, err := encodeRow(nr)
  640. if err != nil {
  641. return 0, err
  642. }
  643. ops = append(ops, BatchOp{Op: batchPut, Key: []byte(m.dataKey(table, pk)), Value: data})
  644. encoded = append(encoded, nr)
  645. }
  646. keys := make([][]byte, len(ops))
  647. seen := make(map[string]struct{}, len(ops))
  648. for i, op := range ops {
  649. key := string(op.Key)
  650. if _, duplicate := seen[key]; duplicate {
  651. return 0, fmt.Errorf("duplicate primary key: %s", key)
  652. }
  653. seen[key] = struct{}{}
  654. keys[i] = op.Key
  655. }
  656. existing := make([]bool, len(keys))
  657. if err := m.pool.WithClient(func(client *KVClient) error {
  658. for start := 0; start < len(keys); start += maxOperations {
  659. end := start + maxOperations
  660. if end > len(keys) {
  661. end = len(keys)
  662. }
  663. found, err := client.ExistsMany(keys[start:end])
  664. if err != nil {
  665. return err
  666. }
  667. copy(existing[start:end], found)
  668. }
  669. return nil
  670. }); err != nil {
  671. return 0, err
  672. }
  673. for i, found := range existing {
  674. if found {
  675. return 0, fmt.Errorf("duplicate primary key: %s", ops[i].Key)
  676. }
  677. }
  678. // Write rows in atomic BATCH_WRITE chunks. Maintain in-memory indexes only
  679. // for rows that actually persisted, so a partial failure cannot leave an
  680. // already-built index stale.
  681. var firstErr error
  682. numOK := 0
  683. for _, chunk := range chunkBatchOps(ops) {
  684. err := m.pool.WithClient(func(c *KVClient) error {
  685. _, err := c.BatchWrite(chunk, nil)
  686. return err
  687. })
  688. if err != nil {
  689. if firstErr == nil {
  690. firstErr = err
  691. }
  692. break
  693. }
  694. numOK += len(chunk)
  695. }
  696. // Maintain indexes for the rows that persisted (the first numOK ops).
  697. for i := 0; i < numOK; i++ {
  698. m.updateIndexesForRow(table, encoded[i], true)
  699. }
  700. if numOK > 0 {
  701. m.bumpGeneration(table)
  702. m.bumpIndexPredicates(table, encoded[:numOK]...)
  703. }
  704. m.incrCount(table, schema.CreatedAt, numOK, wasInit)
  705. return numOK, firstErr
  706. }
  707. // updateIndexesForRow adds or removes entries from already-built in-memory
  708. // indexes. Index entries are rebuildable from durable row data, so this method
  709. // intentionally does not write idx:* keys to KV.
  710. func (m *TableManager) updateIndexesForRow(table string, row Row, add bool) {
  711. indexes, err := m.schema.ListTableIndexes(table)
  712. if err != nil || len(indexes) == 0 {
  713. return
  714. }
  715. rowid, ok := rowIDFromRow(row)
  716. if !ok {
  717. return
  718. }
  719. tableKey := strings.ToLower(table)
  720. tableSchema, schemaErr := m.schema.GetSchema(table)
  721. if schemaErr == nil {
  722. m.cacheMu.Lock()
  723. if keys, initialized := m.rowKeyCache[tableKey]; initialized {
  724. if add {
  725. keys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  726. } else {
  727. delete(keys, rowid)
  728. }
  729. }
  730. m.cacheMu.Unlock()
  731. }
  732. for _, idx := range indexes {
  733. indexName := strings.ToLower(idx.Name)
  734. m.cacheMu.RLock()
  735. _, initialized := m.indexCache[indexName]
  736. m.cacheMu.RUnlock()
  737. if !initialized {
  738. continue
  739. }
  740. columns := make([]string, len(idx.Columns))
  741. for i, col := range idx.Columns {
  742. columns[i] = col.Name
  743. }
  744. colValue := m.buildIndexValue(row, columns)
  745. if add {
  746. m.AddIndexEntry(idx.Name, colValue, rowid)
  747. } else {
  748. m.RemoveIndexEntry(idx.Name, colValue, rowid)
  749. }
  750. }
  751. }
  752. // Select retrieves rows from a table by scanning durable rows and collecting
  753. // only matching rows.
  754. func (m *TableManager) Select(table string, filter func(Row) bool) ([]Row, error) {
  755. tl := m.tableLock(table)
  756. tl.RLock()
  757. defer tl.RUnlock()
  758. return m.selectRows(table, filter)
  759. }
  760. // selectRows scans a table while the caller holds its shared or exclusive gate.
  761. func (m *TableManager) selectRows(table string, filter func(Row) bool) ([]Row, error) {
  762. if !m.schema.TableExists(table) {
  763. return nil, fmt.Errorf("table not found: %s", table)
  764. }
  765. rows := make([]Row, 0)
  766. err := m.scanRows(table, func(row Row) (bool, error) {
  767. if filter == nil || filter(row) {
  768. rows = append(rows, row)
  769. }
  770. return false, nil
  771. })
  772. if err != nil {
  773. return nil, err
  774. }
  775. return rows, nil
  776. }
  777. func cloneRow(row Row) Row {
  778. if row == nil {
  779. return nil
  780. }
  781. cloned := make(Row, len(row))
  782. for k, v := range row {
  783. cloned[k] = v
  784. }
  785. return cloned
  786. }
  787. // SelectWithLimit retrieves rows with limit and offset, applying filter/offset
  788. // while scanning and closing the scan early once the limit is reached.
  789. func (m *TableManager) SelectWithLimit(table string, filter func(Row) bool, limit, offset int) ([]Row, error) {
  790. tl := m.tableLock(table)
  791. tl.RLock()
  792. defer tl.RUnlock()
  793. if !m.schema.TableExists(table) {
  794. return nil, fmt.Errorf("table not found: %s", table)
  795. }
  796. rows := make([]Row, 0)
  797. skipped := 0
  798. pageSize := scanPageSize
  799. if limit > 0 {
  800. desired := limit
  801. if offset > 0 {
  802. if offset >= scanPageSize-desired {
  803. desired = scanPageSize
  804. } else {
  805. desired += offset
  806. }
  807. }
  808. if desired < 1 {
  809. desired = 1
  810. }
  811. if desired < pageSize {
  812. pageSize = desired
  813. }
  814. }
  815. err := m.scanRowsWithPageSize(table, uint32(pageSize), func(row Row) (bool, error) {
  816. if filter != nil && !filter(row) {
  817. return false, nil
  818. }
  819. if skipped < offset {
  820. skipped++
  821. return false, nil
  822. }
  823. rows = append(rows, row)
  824. return limit > 0 && len(rows) >= limit, nil
  825. })
  826. if err != nil {
  827. return nil, err
  828. }
  829. return rows, nil
  830. }
  831. // Update updates rows matching the filter.
  832. func (m *TableManager) Update(table string, updates Row, filter func(Row) bool) (int, error) {
  833. tl := m.tableLock(strings.ToLower(table))
  834. tl.Lock()
  835. defer tl.Unlock()
  836. rows, err := m.selectRows(table, filter)
  837. if err != nil {
  838. return 0, err
  839. }
  840. schema, err := m.schema.GetSchema(table)
  841. if err != nil {
  842. return 0, err
  843. }
  844. count := 0
  845. for _, row := range rows {
  846. // Snapshot the pre-update row so removed index entries can be restored
  847. // if persistence fails.
  848. oldRow := cloneRow(row)
  849. m.updateIndexesForRow(table, row, false)
  850. // Apply updates
  851. for k, v := range updates {
  852. // Normalize column name
  853. for _, col := range schema.Columns {
  854. if strings.EqualFold(k, col.Name) {
  855. row[col.Name] = v
  856. break
  857. }
  858. }
  859. }
  860. // Get primary key
  861. pkValue := row[schema.PrimaryKey]
  862. pk := fmt.Sprintf("%v", pkValue)
  863. // Serialize row
  864. data, err := encodeRow(row)
  865. if err != nil {
  866. m.updateIndexesForRow(table, oldRow, true)
  867. continue
  868. }
  869. // Write back
  870. key := m.dataKey(table, pk)
  871. err = m.pool.WithClient(func(c *KVClient) error {
  872. _, err := c.Put([]byte(key), data)
  873. return err
  874. })
  875. if err == nil {
  876. // Add new index entries after update
  877. m.updateIndexesForRow(table, row, true)
  878. m.bumpIndexPredicates(table, oldRow, row)
  879. count++
  880. } else {
  881. m.updateIndexesForRow(table, oldRow, true)
  882. }
  883. }
  884. if count > 0 {
  885. m.bumpGeneration(table)
  886. }
  887. return count, nil
  888. }
  889. // UpdateFunc updates rows matching the filter using a function to compute new values.
  890. // The updateFn receives the current row and returns the updates to apply.
  891. func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
  892. tl := m.tableLock(strings.ToLower(table))
  893. tl.Lock()
  894. defer tl.Unlock()
  895. rows, err := m.selectRows(table, filter)
  896. if err != nil {
  897. return 0, err
  898. }
  899. schema, err := m.schema.GetSchema(table)
  900. if err != nil {
  901. return 0, err
  902. }
  903. count := 0
  904. for _, row := range rows {
  905. oldRow := cloneRow(row)
  906. m.updateIndexesForRow(table, row, false)
  907. // Compute updates using the provided function
  908. updates, err := updateFn(row)
  909. if err != nil {
  910. m.updateIndexesForRow(table, oldRow, true)
  911. return count, err
  912. }
  913. // Apply updates
  914. for k, v := range updates {
  915. // Normalize column name
  916. for _, col := range schema.Columns {
  917. if strings.EqualFold(k, col.Name) {
  918. row[col.Name] = v
  919. break
  920. }
  921. }
  922. }
  923. // Get primary key
  924. pkValue := row[schema.PrimaryKey]
  925. pk := fmt.Sprintf("%v", pkValue)
  926. // Serialize row
  927. data, err := encodeRow(row)
  928. if err != nil {
  929. m.updateIndexesForRow(table, oldRow, true)
  930. continue
  931. }
  932. // Write back
  933. key := m.dataKey(table, pk)
  934. err = m.pool.WithClient(func(c *KVClient) error {
  935. _, err := c.Put([]byte(key), data)
  936. return err
  937. })
  938. if err == nil {
  939. // Add new index entries after update
  940. m.updateIndexesForRow(table, row, true)
  941. m.bumpIndexPredicates(table, oldRow, row)
  942. count++
  943. } else {
  944. m.updateIndexesForRow(table, oldRow, true)
  945. }
  946. }
  947. if count > 0 {
  948. m.bumpGeneration(table)
  949. }
  950. return count, nil
  951. }
  952. // UpdateByPK updates one row without scanning the table. It uses the key's
  953. // striped lock and a compare-and-swap write so a concurrent modification of the
  954. // same row fails with a serialization error instead of being silently lost.
  955. func (m *TableManager) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) {
  956. schema, err := m.schema.GetSchema(table)
  957. if err != nil {
  958. return nil, false, err
  959. }
  960. key := m.dataKey(table, pk)
  961. tl := m.tableLock(table)
  962. tl.RLock()
  963. defer tl.RUnlock()
  964. st := m.stripeKey(key)
  965. st.Lock()
  966. defer st.Unlock()
  967. row, lsn, err := m.getByPKWithLSN(table, pk)
  968. if err == ErrKeyNotFound {
  969. return nil, false, nil
  970. }
  971. if err != nil {
  972. return nil, false, err
  973. }
  974. oldRow := cloneRow(row)
  975. updates, err := updateFn(row)
  976. if err != nil {
  977. return nil, false, err
  978. }
  979. for name, value := range updates {
  980. for _, column := range schema.Columns {
  981. if strings.EqualFold(name, column.Name) {
  982. row[column.Name] = value
  983. break
  984. }
  985. }
  986. }
  987. data, err := encodeRow(row)
  988. if err != nil {
  989. return nil, false, err
  990. }
  991. committed, err := m.compareWritePoint(
  992. []CompareCheck{{Key: []byte(key), LSN: lsn}},
  993. []BatchOp{{Op: batchPut, Key: []byte(key), Value: data}},
  994. )
  995. if err != nil {
  996. return nil, false, err
  997. }
  998. if !committed {
  999. return nil, false, ErrSerialization
  1000. }
  1001. m.updateIndexesForRow(table, oldRow, false)
  1002. m.updateIndexesForRow(table, row, true)
  1003. m.bumpIndexPredicates(table, oldRow, row)
  1004. m.bumpGeneration(table)
  1005. return oldRow, true, nil
  1006. }
  1007. // Delete deletes rows matching the filter.
  1008. func (m *TableManager) Delete(table string, filter func(Row) bool) (int, error) {
  1009. tl := m.tableLock(strings.ToLower(table))
  1010. tl.Lock()
  1011. defer tl.Unlock()
  1012. rows, err := m.selectRows(table, filter)
  1013. if err != nil {
  1014. return 0, err
  1015. }
  1016. schema, err := m.schema.GetSchema(table)
  1017. if err != nil {
  1018. return 0, err
  1019. }
  1020. wasInit := m.countInitialized(table)
  1021. count := 0
  1022. for _, row := range rows {
  1023. // Remove index entries before deleting row
  1024. m.updateIndexesForRow(table, row, false)
  1025. pkValue := row[schema.PrimaryKey]
  1026. pk := fmt.Sprintf("%v", pkValue)
  1027. key := m.dataKey(table, pk)
  1028. err = m.pool.WithClient(func(c *KVClient) error {
  1029. _, err := c.Del([]byte(key))
  1030. return err
  1031. })
  1032. if err == nil {
  1033. m.bumpIndexPredicates(table, row)
  1034. count++
  1035. } else {
  1036. // Restore the index entries removed above.
  1037. m.updateIndexesForRow(table, row, true)
  1038. }
  1039. }
  1040. if count > 0 {
  1041. m.bumpGeneration(table)
  1042. }
  1043. m.incrCount(table, schema.CreatedAt, -count, wasInit)
  1044. return count, nil
  1045. }
  1046. // DeleteByPK deletes one row without scanning the table, using the key's
  1047. // striped lock and a compare-and-swap delete.
  1048. func (m *TableManager) DeleteByPK(table, pk string) (Row, bool, error) {
  1049. schema, err := m.schema.GetSchema(table)
  1050. if err != nil {
  1051. return nil, false, err
  1052. }
  1053. key := m.dataKey(table, pk)
  1054. tl := m.tableLock(table)
  1055. tl.RLock()
  1056. defer tl.RUnlock()
  1057. st := m.stripeKey(key)
  1058. st.Lock()
  1059. defer st.Unlock()
  1060. wasInit := m.countInitialized(table)
  1061. row, lsn, err := m.getByPKWithLSN(table, pk)
  1062. if err == ErrKeyNotFound {
  1063. return nil, false, nil
  1064. }
  1065. if err != nil {
  1066. return nil, false, err
  1067. }
  1068. committed, err := m.compareWritePoint(
  1069. []CompareCheck{{Key: []byte(key), LSN: lsn}},
  1070. []BatchOp{{Op: batchDelete, Key: []byte(key)}},
  1071. )
  1072. if err != nil {
  1073. return nil, false, err
  1074. }
  1075. if !committed {
  1076. return nil, false, ErrSerialization
  1077. }
  1078. m.updateIndexesForRow(table, row, false)
  1079. m.bumpIndexPredicates(table, row)
  1080. m.bumpGeneration(table)
  1081. m.incrCount(table, schema.CreatedAt, -1, wasInit)
  1082. return row, true, nil
  1083. }
  1084. // GetByPK retrieves a row by primary key. Point reads take only the key's
  1085. // striped lock so reads of different keys progress concurrently.
  1086. func (m *TableManager) GetByPK(table string, pk string) (Row, error) {
  1087. key := m.dataKey(table, pk)
  1088. st := m.stripeKey(key)
  1089. st.Lock()
  1090. defer st.Unlock()
  1091. if !m.schema.TableExists(table) {
  1092. return nil, fmt.Errorf("table not found: %s", table)
  1093. }
  1094. row, _, err := m.getByPKWithLSN(table, pk)
  1095. return row, err
  1096. }
  1097. // getByPKWithLSN reads a row by primary key and returns its KV LSN (0 when
  1098. // absent). It performs no locking; callers must hold the appropriate striped
  1099. // or table lock.
  1100. func (m *TableManager) getByPKWithLSN(table, pk string) (Row, uint64, error) {
  1101. key := m.dataKey(table, pk)
  1102. var value []byte
  1103. var lsn uint64
  1104. err := m.pool.WithClient(func(c *KVClient) error {
  1105. res, err := c.Get([]byte(key))
  1106. if err != nil {
  1107. return err
  1108. }
  1109. value = res.Value
  1110. lsn = res.LSN
  1111. return nil
  1112. })
  1113. if err != nil {
  1114. if err == ErrKeyNotFound {
  1115. return nil, 0, ErrKeyNotFound
  1116. }
  1117. return nil, 0, err
  1118. }
  1119. row, err := decodeRow(value)
  1120. if err != nil {
  1121. return nil, 0, fmt.Errorf("failed to parse row: %w", err)
  1122. }
  1123. return row, lsn, nil
  1124. }
  1125. // Count returns the number of rows in a table matching the filter.
  1126. func (m *TableManager) Count(table string, filter func(Row) bool) (int, error) {
  1127. tl := m.tableLock(table)
  1128. tl.RLock()
  1129. defer tl.RUnlock()
  1130. if !m.schema.TableExists(table) {
  1131. return 0, fmt.Errorf("table not found: %s", table)
  1132. }
  1133. count := 0
  1134. err := m.scanRows(table, func(row Row) (bool, error) {
  1135. if filter == nil || filter(row) {
  1136. count++
  1137. }
  1138. return false, nil
  1139. })
  1140. if err != nil {
  1141. return 0, err
  1142. }
  1143. return count, nil
  1144. }
  1145. // Truncate removes all rows from a table.
  1146. func (m *TableManager) Truncate(table string) (int, error) {
  1147. return m.Delete(table, nil)
  1148. }
  1149. // isIntegerType checks if a type name is an integer type.
  1150. func isIntegerType(typeName string) bool {
  1151. t := strings.ToUpper(typeName)
  1152. switch t {
  1153. case "INTEGER", "INT", "SMALLINT", "BIGINT", "TINYINT", "MEDIUMINT":
  1154. return true
  1155. }
  1156. return false
  1157. }
  1158. // IsRowIDColumn checks if a column name is a ROWID alias.
  1159. func IsRowIDColumn(name string) bool {
  1160. n := strings.ToLower(name)
  1161. return n == "rowid" || n == "oid" || n == "_rowid_"
  1162. }
  1163. // Index entry methods - leveraging radix trie for prefix-based lookups
  1164. // Format: {database}:idx:{index_name}:{column_value} → JSON array of rowids
  1165. // indexEntryKey returns the key for an index entry.
  1166. func (m *TableManager) indexEntryKey(indexName string, colValue interface{}) string {
  1167. return fmt.Sprintf("%s:idx:%s:%s", m.database, strings.ToLower(indexName), formatIndexValue(colValue))
  1168. }
  1169. // indexPrefix returns the prefix for all entries of an index.
  1170. func (m *TableManager) indexPrefix(indexName string) string {
  1171. return fmt.Sprintf("%s:idx:%s:", m.database, strings.ToLower(indexName))
  1172. }
  1173. func formatIndexValue(value interface{}) string {
  1174. switch v := value.(type) {
  1175. case float64:
  1176. if v == float64(int64(v)) {
  1177. return fmt.Sprintf("%d", int64(v))
  1178. }
  1179. return fmt.Sprintf("%f", v)
  1180. case int64:
  1181. return fmt.Sprintf("%d", v)
  1182. case int:
  1183. return fmt.Sprintf("%d", v)
  1184. default:
  1185. return fmt.Sprintf("%v", v)
  1186. }
  1187. }
  1188. func rowIDFromRow(row Row) (int64, bool) {
  1189. switch v := row["_rowid_"].(type) {
  1190. case int64:
  1191. return v, true
  1192. case int:
  1193. return int64(v), true
  1194. case float64:
  1195. return int64(v), true
  1196. default:
  1197. return 0, false
  1198. }
  1199. }
  1200. func (m *TableManager) ensureIndex(index *Index) error {
  1201. indexKey := strings.ToLower(index.Name)
  1202. m.cacheMu.RLock()
  1203. disabled := m.disabledIndexes[indexKey]
  1204. _, initialized := m.indexCache[indexKey]
  1205. m.cacheMu.RUnlock()
  1206. if disabled {
  1207. return nil
  1208. }
  1209. if initialized {
  1210. return nil
  1211. }
  1212. // Serialize index build against writes to the same table so the derived
  1213. // entries cannot miss a concurrently-inserted row.
  1214. table := index.Table
  1215. key := strings.ToLower(table)
  1216. tl := m.tableLock(key)
  1217. tl.Lock()
  1218. defer tl.Unlock()
  1219. m.cacheMu.RLock()
  1220. disabled = m.disabledIndexes[indexKey]
  1221. _, initialized = m.indexCache[indexKey]
  1222. m.cacheMu.RUnlock()
  1223. if disabled {
  1224. return nil
  1225. }
  1226. if initialized {
  1227. return nil
  1228. }
  1229. columns := make([]string, len(index.Columns))
  1230. for i, col := range index.Columns {
  1231. columns[i] = col.Name
  1232. }
  1233. tableSchema, err := m.schema.GetSchema(table)
  1234. if err != nil {
  1235. return err
  1236. }
  1237. values := make(map[string][]int64)
  1238. rowKeys := make(map[int64]string)
  1239. if err := m.scanRows(table, func(row Row) (bool, error) {
  1240. rowid, ok := rowIDFromRow(row)
  1241. if !ok {
  1242. return false, nil
  1243. }
  1244. colValue := m.buildIndexValue(row, columns)
  1245. valueKey := formatIndexValue(colValue)
  1246. values[valueKey] = append(values[valueKey], rowid)
  1247. rowKeys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  1248. return false, nil
  1249. }); err != nil {
  1250. return err
  1251. }
  1252. m.cacheMu.Lock()
  1253. if _, initialized := m.indexCache[indexKey]; !initialized {
  1254. m.indexCache[indexKey] = values
  1255. m.indexTable[indexKey] = key
  1256. m.rowKeyCache[key] = rowKeys
  1257. }
  1258. m.cacheMu.Unlock()
  1259. return nil
  1260. }
  1261. // AddIndexEntry adds a rowid to an in-memory index entry.
  1262. func (m *TableManager) AddIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1263. indexKey := strings.ToLower(indexName)
  1264. valueKey := formatIndexValue(colValue)
  1265. m.cacheMu.Lock()
  1266. defer m.cacheMu.Unlock()
  1267. values, ok := m.indexCache[indexKey]
  1268. if !ok {
  1269. return nil
  1270. }
  1271. rowids := values[valueKey]
  1272. for _, r := range rowids {
  1273. if r == rowid {
  1274. return nil
  1275. }
  1276. }
  1277. values[valueKey] = append(rowids, rowid)
  1278. return nil
  1279. }
  1280. // RemoveIndexEntry removes a rowid from an in-memory index entry.
  1281. func (m *TableManager) RemoveIndexEntry(indexName string, colValue interface{}, rowid int64) error {
  1282. indexKey := strings.ToLower(indexName)
  1283. valueKey := formatIndexValue(colValue)
  1284. m.cacheMu.Lock()
  1285. defer m.cacheMu.Unlock()
  1286. values, ok := m.indexCache[indexKey]
  1287. if !ok {
  1288. return nil
  1289. }
  1290. rowids := values[valueKey]
  1291. newRowids := make([]int64, 0, len(rowids))
  1292. for _, r := range rowids {
  1293. if r != rowid {
  1294. newRowids = append(newRowids, r)
  1295. }
  1296. }
  1297. if len(newRowids) == 0 {
  1298. delete(values, valueKey)
  1299. return nil
  1300. }
  1301. values[valueKey] = newRowids
  1302. return nil
  1303. }
  1304. // LookupIndex returns rowids matching a column value using the index.
  1305. func (m *TableManager) LookupIndex(indexName string, colValue interface{}) ([]int64, error) {
  1306. index, err := m.schema.GetIndex(indexName)
  1307. if err != nil {
  1308. return nil, err
  1309. }
  1310. if err := m.ensureIndex(index); err != nil {
  1311. return nil, err
  1312. }
  1313. indexKey := strings.ToLower(indexName)
  1314. valueKey := formatIndexValue(colValue)
  1315. m.cacheMu.RLock()
  1316. rowids := append([]int64(nil), m.indexCache[indexKey][valueKey]...)
  1317. m.cacheMu.RUnlock()
  1318. return rowids, nil
  1319. }
  1320. // ClearIndex removes an index's in-memory entries and marks it disabled so a
  1321. // concurrent lookup cannot rebuild it after DROP but before the schema entry
  1322. // is removed. Index entries are derived from durable rows, so there is nothing
  1323. // durable to delete here.
  1324. func (m *TableManager) ClearIndex(indexName, tableName string, columns []string) error {
  1325. indexKey := strings.ToLower(indexName)
  1326. tableKey := strings.ToLower(tableName)
  1327. m.cacheMu.Lock()
  1328. delete(m.indexCache, indexKey)
  1329. delete(m.indexTable, indexKey)
  1330. m.disabledIndexes[indexKey] = true
  1331. rowKeysNeeded := false
  1332. for _, indexedTable := range m.indexTable {
  1333. if indexedTable == tableKey {
  1334. rowKeysNeeded = true
  1335. break
  1336. }
  1337. }
  1338. if !rowKeysNeeded {
  1339. delete(m.rowKeyCache, tableKey)
  1340. }
  1341. m.cacheMu.Unlock()
  1342. return nil
  1343. }
  1344. // BuildIndex builds index entries for all existing rows in a table.
  1345. func (m *TableManager) BuildIndex(indexName, tableName string, columns []string) error {
  1346. indexKey := strings.ToLower(indexName)
  1347. m.cacheMu.Lock()
  1348. delete(m.disabledIndexes, indexKey)
  1349. delete(m.indexCache, indexKey)
  1350. delete(m.indexTable, indexKey)
  1351. m.cacheMu.Unlock()
  1352. index, err := m.schema.GetIndex(indexName)
  1353. if err == nil {
  1354. return m.ensureIndex(index)
  1355. }
  1356. tableSchema, schemaErr := m.schema.GetSchema(tableName)
  1357. if schemaErr != nil {
  1358. return schemaErr
  1359. }
  1360. values := make(map[string][]int64)
  1361. rowKeys := make(map[int64]string)
  1362. if err := m.scanRows(tableName, func(row Row) (bool, error) {
  1363. rowid, ok := rowIDFromRow(row)
  1364. if !ok {
  1365. return false, nil
  1366. }
  1367. colValue := m.buildIndexValue(row, columns)
  1368. values[formatIndexValue(colValue)] = append(values[formatIndexValue(colValue)], rowid)
  1369. rowKeys[rowid] = fmt.Sprintf("%v", row[tableSchema.PrimaryKey])
  1370. return false, nil
  1371. }); err != nil {
  1372. return err
  1373. }
  1374. m.cacheMu.Lock()
  1375. m.indexCache[indexKey] = values
  1376. m.indexTable[indexKey] = strings.ToLower(tableName)
  1377. m.rowKeyCache[strings.ToLower(tableName)] = rowKeys
  1378. m.cacheMu.Unlock()
  1379. return nil
  1380. }
  1381. // buildIndexValue creates the index key value from row columns.
  1382. func (m *TableManager) buildIndexValue(row Row, columns []string) string {
  1383. formatValue := func(v interface{}) string {
  1384. switch val := v.(type) {
  1385. case float64:
  1386. // Check if it's actually an integer value
  1387. if val == float64(int64(val)) {
  1388. return fmt.Sprintf("%d", int64(val))
  1389. }
  1390. return fmt.Sprintf("%f", val)
  1391. case int64:
  1392. return fmt.Sprintf("%d", val)
  1393. case int:
  1394. return fmt.Sprintf("%d", val)
  1395. default:
  1396. return fmt.Sprintf("%v", val)
  1397. }
  1398. }
  1399. if len(columns) == 1 {
  1400. return formatValue(row[columns[0]])
  1401. }
  1402. // Multi-column index: concatenate values with separator
  1403. var parts []string
  1404. for _, col := range columns {
  1405. parts = append(parts, formatValue(row[col]))
  1406. }
  1407. return strings.Join(parts, "\x00")
  1408. }
  1409. type indexedRowVersion struct {
  1410. row Row
  1411. key string
  1412. lsn uint64
  1413. }
  1414. type indexPredicateSnapshot struct {
  1415. valueKey string
  1416. valueGen uint64
  1417. wildcardKey string
  1418. wildcardGen uint64
  1419. }
  1420. // selectByIndexWithLSN retrieves indexed rows and their durable versions. The
  1421. // transaction layer uses the versions for optimistic commit validation.
  1422. func (m *TableManager) selectByIndexWithLSN(table, indexName string, colValue interface{}) ([]indexedRowVersion, indexPredicateSnapshot, error) {
  1423. index, err := m.schema.GetIndex(indexName)
  1424. if err != nil {
  1425. return nil, indexPredicateSnapshot{}, err
  1426. }
  1427. if err := m.ensureIndex(index); err != nil {
  1428. return nil, indexPredicateSnapshot{}, err
  1429. }
  1430. tableKey := strings.ToLower(table)
  1431. tl := m.tableLock(tableKey)
  1432. tl.RLock()
  1433. defer tl.RUnlock()
  1434. if !m.schema.TableExists(table) {
  1435. return nil, indexPredicateSnapshot{}, fmt.Errorf("table not found: %s", table)
  1436. }
  1437. indexKey := strings.ToLower(indexName)
  1438. valueKey := formatIndexValue(colValue)
  1439. predicateKey, predicateGen, wildcardKey, wildcardGen := m.predicateSnapshot(table, indexName, valueKey)
  1440. snapshot := indexPredicateSnapshot{
  1441. valueKey: predicateKey, valueGen: predicateGen,
  1442. wildcardKey: wildcardKey, wildcardGen: wildcardGen,
  1443. }
  1444. m.cacheMu.RLock()
  1445. rowids := append([]int64(nil), m.indexCache[indexKey][valueKey]...)
  1446. primaryKeys := make([]string, 0, len(rowids))
  1447. missingRowID := int64(0)
  1448. missingRowKey := false
  1449. for _, rowid := range rowids {
  1450. if primaryKey, ok := m.rowKeyCache[tableKey][rowid]; ok {
  1451. primaryKeys = append(primaryKeys, primaryKey)
  1452. } else {
  1453. missingRowID = rowid
  1454. missingRowKey = true
  1455. break
  1456. }
  1457. }
  1458. m.cacheMu.RUnlock()
  1459. if missingRowKey {
  1460. return nil, indexPredicateSnapshot{}, fmt.Errorf("index %s is missing rowid %d", indexName, missingRowID)
  1461. }
  1462. // If no rowids found, return empty result
  1463. if len(rowids) == 0 {
  1464. return []indexedRowVersion{}, snapshot, nil
  1465. }
  1466. rows := make([]indexedRowVersion, 0, len(primaryKeys))
  1467. err = m.pool.WithClient(func(client *KVClient) error {
  1468. keys := make([][]byte, len(primaryKeys))
  1469. for i, primaryKey := range primaryKeys {
  1470. keys[i] = []byte(m.dataKey(table, primaryKey))
  1471. }
  1472. results, err := client.MultiGet(keys)
  1473. if err != nil {
  1474. return err
  1475. }
  1476. for i, result := range results {
  1477. if !result.Found {
  1478. primaryKey := primaryKeys[i]
  1479. return fmt.Errorf("index %s references missing primary key %s", indexName, primaryKey)
  1480. }
  1481. row, err := decodeRow(result.Value)
  1482. if err != nil {
  1483. return err
  1484. }
  1485. rows = append(rows, indexedRowVersion{
  1486. row: row,
  1487. key: m.dataKey(table, primaryKeys[i]),
  1488. lsn: result.LSN,
  1489. })
  1490. }
  1491. return nil
  1492. })
  1493. if err != nil {
  1494. return nil, indexPredicateSnapshot{}, err
  1495. }
  1496. return rows, snapshot, nil
  1497. }
  1498. // SelectByIndex retrieves rows using an in-memory equality index followed by a
  1499. // single MultiGet for the matching primary keys.
  1500. func (m *TableManager) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
  1501. versions, _, err := m.selectByIndexWithLSN(table, indexName, colValue)
  1502. if err != nil {
  1503. return nil, err
  1504. }
  1505. rows := make([]Row, len(versions))
  1506. for i := range versions {
  1507. rows[i] = versions[i].row
  1508. }
  1509. return rows, nil
  1510. }