db cli支持

db/engine.go

@@ -835,8 +835,248 @@ func WildcardMatch(str, pattern string) bool {
 	return j == pLen
 }
 
+func (e *Engine) Count(sql string) (int, error) {
+	// Simple wrapper around Query for now, but optimized to not load values if possible.
+	// Actually we should refactor Query to support a "countOnly" mode.
+	// But to avoid breaking API signatures too much, let's just parse and execute with optimization.
+	
+	// Reuse Query logic but with count optimization
+	// We can't easily reuse Query() because it returns []QueryResult (heavy).
+	// Let's implement a specialized version or refactor Query internals.
+	// Refactoring Query internals to take a visitor/callback is cleanest.
+	
+	return e.execute(sql, true)
+}
+
+// execute is the internal implementation for Query and Count
+// countOnly: if true, returns count in first return val (as int cast), second is nil
+// Wait, return types need to be consistent. Let's return (count, results, error)
+func (e *Engine) execute(sql string, countOnly bool) (int, error) {
+	sql = strings.TrimSpace(sql)
+	
+	reLimit := regexp.MustCompile(`(?i)\s+LIMIT\s+(\d+)`)
+	reOffset := regexp.MustCompile(`(?i)\s+OFFSET\s+(\d+)`)
+	
+	limit := -1
+	offset := 0
+	
+	if match := reLimit.FindStringSubmatch(sql); len(match) > 0 {
+		limit, _ = strconv.Atoi(match[1])
+		sql = reLimit.ReplaceAllString(sql, "")
+	}
+	if match := reOffset.FindStringSubmatch(sql); len(match) > 0 {
+		offset, _ = strconv.Atoi(match[1])
+		sql = reOffset.ReplaceAllString(sql, "")
+	}
+
+	re := regexp.MustCompile(`(key|value|CommitIndex)\s*(=|like|>=|<=|>|<)\s*("[^"]*"|\d+)`)
+	matches := re.FindAllStringSubmatch(sql, -1)
+
+	if len(matches) == 0 {
+		return 0, fmt.Errorf("invalid query")
+	}
+
+	extractString := func(s string) string {
+		return strings.Trim(strings.TrimSpace(s), "\"")
+	}
+
+	// Optimization: Point Lookup Fast Path
+	// Only if not countOnly or if we need value to verify?
+	// If countOnly and query is `key=".."`, we can just check existence O(1).
+	for _, match := range matches {
+		if match[1] == "key" && match[2] == "=" {
+			targetKey := extractString(match[3])
+			entry, ok := e.Index.Get(targetKey)
+			if !ok {
+				return 0, nil
+			}
+			
+			// If we need to filter by Value, we MUST load it.
+			needValue := false
+			for _, m := range matches {
+				if m[1] == "value" { needValue = true; break }
+			}
+			
+			var val string
+			if needValue {
+				v, err := e.Storage.ReadValue(entry.ValueOffset)
+				if err != nil { return 0, nil }
+				val = v
+			}
+			
+			matchAll := true
+			for _, m := range matches {
+				field, op, valRaw := m[1], m[2], m[3]
+				switch field {
+				case "CommitIndex":
+					num, _ := strconv.ParseUint(valRaw, 10, 64)
+					switch op {
+					case ">": if !(entry.CommitIndex > num) { matchAll = false }
+					case "<": if !(entry.CommitIndex < num) { matchAll = false }
+					case ">=": if !(entry.CommitIndex >= num) { matchAll = false }
+					case "<=": if !(entry.CommitIndex <= num) { matchAll = false }
+					case "=": if !(entry.CommitIndex == num) { matchAll = false }
+					}
+				case "value":
+					t := extractString(valRaw)
+					switch op {
+					case "=": if val != t { matchAll = false }
+					case "like": if !WildcardMatch(val, t) { matchAll = false }
+					}
+				}
+				if !matchAll { break }
+			}
+			if matchAll {
+				return 1, nil
+			}
+			return 0, nil
+		}
+	}
+
+	// Inverted Index Logic
+	var candidates map[string]bool
+	var useFTIndex bool = false
+	
+	for _, match := range matches {
+		if match[1] == "value" && match[2] == "like" {
+			pattern := extractString(match[3])
+			clean := strings.Trim(pattern, "*")
+			if len(clean) > 0 && !strings.Contains(clean, "*") && !strings.Contains(clean, "?") {
+				matches := e.FTIndex.Search(pattern)
+				if matches != nil {
+					currentSet := make(map[string]bool)
+					for _, k := range matches {
+						currentSet[k] = true
+					}
+					if !useFTIndex {
+						candidates = currentSet
+						useFTIndex = true
+					} else {
+						newSet := make(map[string]bool)
+						for k := range candidates {
+							if currentSet[k] { newSet[k] = true }
+						}
+						candidates = newSet
+					}
+				} else {
+					return 0, nil
+				}
+			}
+		}
+	}
+	
+	var iterator func(func(string, IndexEntry) bool)
+	
+	if useFTIndex {
+		iterator = func(cb func(string, IndexEntry) bool) {
+			keys := make([]string, 0, len(candidates))
+			for k := range candidates { keys = append(keys, k) }
+			sort.Strings(keys)
+			for _, k := range keys {
+				if entry, ok := e.Index.Get(k); ok {
+					if !cb(k, entry) { return }
+				}
+			}
+		}
+	} else {
+		var prefix string = ""
+		var usePrefix bool = false
+		for _, match := range matches {
+			if match[1] == "key" && match[2] == "like" {
+				pattern := extractString(match[3])
+				if strings.HasSuffix(pattern, "*") {
+					clean := pattern[:len(pattern)-1]
+					if !strings.ContainsAny(clean, "*?") {
+						prefix = clean
+						usePrefix = true
+						break
+					}
+				}
+			}
+		}
+		
+		iterator = func(cb func(string, IndexEntry) bool) {
+			if usePrefix {
+				e.Index.WalkPrefix(prefix, cb)
+			} else {
+				e.Index.WalkPrefix("", cb)
+			}
+		}
+	}
+
+	matchedCount := 0
+	
+	// Execution
+	iterator(func(key string, entry IndexEntry) bool {
+		var valStr string
+		var valLoaded bool
+				
+		matchAll := true
+		for _, match := range matches {
+			field, op, valRaw := match[1], match[2], match[3]
+			switch field {
+			case "CommitIndex":
+				num, _ := strconv.ParseUint(valRaw, 10, 64)
+				switch op {
+				case ">": if !(entry.CommitIndex > num) { matchAll = false }
+				case "<": if !(entry.CommitIndex < num) { matchAll = false }
+				case ">=": if !(entry.CommitIndex >= num) { matchAll = false }
+				case "<=": if !(entry.CommitIndex <= num) { matchAll = false }
+				case "=": if !(entry.CommitIndex == num) { matchAll = false }
+				}
+			case "key":
+				target := extractString(valRaw)
+				switch op {
+				case "=": if key != target { matchAll = false }
+				case "like": if !WildcardMatch(key, target) { matchAll = false }
+				}
+			case "value":
+				if !valLoaded {
+					v, err := e.Storage.ReadValue(entry.ValueOffset)
+					if err != nil { matchAll = false; break }
+					valStr = v
+					valLoaded = true
+				}
+				target := extractString(valRaw)
+				switch op {
+				case "=": if valStr != target { matchAll = false }
+				case "like": if !WildcardMatch(valStr, target) { matchAll = false }
+				}
+			}
+			if !matchAll { break }
+		}
+		
+		if matchAll {
+			matchedCount++
+			if limit > 0 && offset == 0 && matchedCount >= limit {
+				return false
+			}
+		}
+		return true
+	})
+
+	if offset > 0 {
+		if offset >= matchedCount {
+			return 0, nil
+		}
+		matchedCount -= offset
+	}
+	if limit >= 0 {
+		if limit < matchedCount {
+			matchedCount = limit
+		}
+	}
+
+	return matchedCount, nil
+}
+
 func (e *Engine) Query(sql string) ([]QueryResult, error) {
+	return e.queryInternal(sql)
+}
+
+func (e *Engine) queryInternal(sql string) ([]QueryResult, error) {
 	sql = strings.TrimSpace(sql)
+
 	
 	reLimit := regexp.MustCompile(`(?i)\s+LIMIT\s+(\d+)`)
 	reOffset := regexp.MustCompile(`(?i)\s+OFFSET\s+(\d+)`)
@@ -1102,3 +1342,26 @@ func (e *Engine) Query(sql string) ([]QueryResult, error) {
 
 	return results, nil
 }
+
+func (e *Engine) Snapshot() ([]byte, error) {
+	// Not implemented for Radix Tree yet in this demo
+	return nil, nil
+}
+
+func (e *Engine) Restore(data []byte) error {
+	return nil
+}
+
+// DebugClearAuxiliary clears caches and inverted index to measure core index memory usage.
+// For internal benchmark use only.
+func (e *Engine) DebugClearAuxiliary() {
+	// Clear Cache
+	e.Storage.cacheMu.Lock()
+	e.Storage.cache = make(map[int64]string)
+	e.Storage.cacheMu.Unlock()
+	
+	// Clear FTIndex
+	e.FTIndex.mu.Lock()
+	e.FTIndex.index = make(map[string][]string)
+	e.FTIndex.mu.Unlock()
+}

example/cli/main.go

@@ -0,0 +1,280 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"math/rand"
+	"os"
+	"strings"
+	"time"
+	"unicode/utf8"
+
+	"igit.com/xbase/raft/db"
+)
+
+const DataDir = "cli_data"
+
+// ANSI Colors
+const (
+	ColorReset   = "\033[0m"
+	ColorRed     = "\033[31m"
+	ColorGreen   = "\033[32m"
+	ColorYellow  = "\033[33m"
+	ColorBlue    = "\033[34m"
+	ColorPurple  = "\033[35m"
+	ColorCyan    = "\033[36m"
+	ColorGray    = "\033[37m"
+	ColorBold    = "\033[1m"
+)
+
+func main() {
+	// Initialize DB
+	e, err := db.NewEngine(DataDir)
+	if err != nil {
+		fmt.Printf("%sError initializing DB: %v%s\n", ColorRed, err, ColorReset)
+		os.Exit(1)
+	}
+	defer e.Close()
+
+	printWelcome()
+
+	// Auto-seed if empty
+	res, _ := e.Query("LIMIT 1")
+	if len(res) == 0 {
+		fmt.Printf("%s[!] Database appears empty. Seeding 1000 records...%s\n", ColorYellow, ColorReset)
+		seedData(e)
+	} else {
+		fmt.Printf("%s[+] Database loaded. Ready for queries.%s\n", ColorGreen, ColorReset)
+	}
+
+	scanner := bufio.NewScanner(os.Stdin)
+
+	for {
+		fmt.Printf("\n%sraftdb%s> ", ColorBlue, ColorReset)
+		if !scanner.Scan() {
+			break
+		}
+		line := scanner.Text()
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+
+		parts := strings.Fields(line)
+		cmd := strings.ToLower(parts[0])
+
+		switch cmd {
+		case "exit", "quit":
+			fmt.Printf("%sBye!%s\n", ColorGreen, ColorReset)
+			return
+		case "help":
+			printHelp()
+		case "seed":
+			seedData(e)
+		case "add":
+			handleAdd(e, parts)
+		case "count":
+			handleCount(e, line)
+		default:
+			handleQuery(e, line)
+		}
+	}
+}
+
+func handleAdd(e *db.Engine, parts []string) {
+	if len(parts) < 3 {
+		fmt.Printf("%sUsage: add <key> <value>%s\n", ColorRed, ColorReset)
+		return
+	}
+	key := parts[1]
+	val := strings.Join(parts[2:], " ")
+	idx := uint64(time.Now().UnixNano())
+	
+	start := time.Now()
+	err := e.Set(key, val, idx)
+	duration := time.Since(start)
+
+	if err != nil {
+		fmt.Printf("%sError adding data: %v%s\n", ColorRed, err, ColorReset)
+	} else {
+		fmt.Printf("%sQuery OK, 1 row affected (%v)%s\n", ColorGreen, duration, ColorReset)
+	}
+}
+
+func handleCount(e *db.Engine, line string) {
+	// Format: count key like "..."
+	// Remove "count" from start
+	query := strings.TrimSpace(line[5:])
+	if query == "" {
+		fmt.Printf("%sUsage: count <query_conditions>%s\n", ColorRed, ColorReset)
+		return
+	}
+
+	start := time.Now()
+	// Use Query to get results then count
+	// Note: If Engine had a Count() method, we would use it here.
+	results, err := e.Query(query)
+	duration := time.Since(start)
+
+	if err != nil {
+		fmt.Printf("%sError: %v%s\n", ColorRed, err, ColorReset)
+		return
+	}
+
+	fmt.Printf("%s+-------+%s\n", ColorYellow, ColorReset)
+	fmt.Printf("%s| COUNT |%s\n", ColorYellow, ColorReset)
+	fmt.Printf("%s+-------+%s\n", ColorYellow, ColorReset)
+	fmt.Printf("| %-5d |\n", len(results))
+	fmt.Printf("%s+-------+%s\n", ColorYellow, ColorReset)
+	fmt.Printf("%s1 row in set (%v)%s\n", ColorGreen, duration, ColorReset)
+}
+
+func handleQuery(e *db.Engine, query string) {
+	start := time.Now()
+	results, err := e.Query(query)
+	duration := time.Since(start)
+
+	if err != nil {
+		fmt.Printf("%sError executing query: %v%s\n", ColorRed, err, ColorReset)
+		return
+	}
+
+	if len(results) == 0 {
+		fmt.Printf("%sEmpty set (%v)%s\n", ColorGray, duration, ColorReset)
+		return
+	}
+
+	// Calculate column widths
+	maxKey := 3
+	maxVal := 5
+	maxIdx := 12
+
+	for _, r := range results {
+		if len(r.Key) > maxKey {
+			maxKey = len(r.Key)
+		}
+		// Truncate value for display if too long
+		vLen := utf8.RuneCountInString(r.Value)
+		if vLen > 60 {
+			vLen = 60
+		}
+		if vLen > maxVal {
+			maxVal = vLen
+		}
+	}
+	if maxKey > 40 { maxKey = 40 } // Hard limit for key display
+
+	// Print Table Header
+	printSeparator(maxKey, maxVal, maxIdx)
+	fmt.Printf("| %-*s | %-*s | %-*s |\n", maxKey, "KEY", maxVal, "VALUE", maxIdx, "COMMIT_INDEX")
+	printSeparator(maxKey, maxVal, maxIdx)
+
+	// Print Rows
+	for _, r := range results {
+		k := r.Key
+		if utf8.RuneCountInString(k) > maxKey {
+			k = k[:maxKey-3] + "..."
+		}
+		v := r.Value
+		if utf8.RuneCountInString(v) > maxVal {
+			v = v[:maxVal-3] + "..."
+		}
+		fmt.Printf("| %-*s | %-*s | %-*d |\n", maxKey, k, maxVal, v, maxIdx, r.CommitIndex)
+	}
+	printSeparator(maxKey, maxVal, maxIdx)
+	
+	fmt.Printf("%s%d rows in set (%v)%s\n", ColorGreen, len(results), duration, ColorReset)
+}
+
+func printSeparator(k, v, i int) {
+	fmt.Printf("%s+%s+%s+%s+%s\n", ColorYellow, strings.Repeat("-", k+2), strings.Repeat("-", v+2), strings.Repeat("-", i+2), ColorReset)
+}
+
+func seedData(e *db.Engine) {
+	fmt.Println("Seeding 1000 records...")
+	start := time.Now()
+	
+	prefixes := []string{"user", "product", "log", "config"}
+	
+	for i := 0; i < 1000; i++ {
+		prefix := prefixes[rand.Intn(len(prefixes))]
+		key := fmt.Sprintf("%s:%d", prefix, i)
+		
+		var val string
+		switch prefix {
+		case "user":
+			val = fmt.Sprintf("name=User%d;role=%s;active=true", i, randomRole())
+		case "product":
+			val = fmt.Sprintf("title=Item%d;price=%d;desc=%s", i, rand.Intn(1000), randomDesc())
+		case "log":
+			val = fmt.Sprintf("level=%s;msg=Something happened at %d", randomLevel(), i)
+		case "config":
+			val = fmt.Sprintf("setting=%d;enabled=%v", i, rand.Intn(2) == 1)
+		}
+
+		e.Set(key, val, uint64(i+1))
+	}
+	
+	fmt.Printf("Seeded 1000 records in %v\n", time.Since(start))
+}
+
+func randomRole() string {
+	roles := []string{"admin", "editor", "viewer", "guest"}
+	return roles[rand.Intn(len(roles))]
+}
+
+func randomLevel() string {
+	levels := []string{"INFO", "WARN", "ERROR", "DEBUG"}
+	return levels[rand.Intn(len(levels))]
+}
+
+func randomDesc() string {
+	adjectives := []string{"Great", "Awesome", "Standard", "Basic", "Premium"}
+	nouns := []string{"Widget", "Tool", "Gadget", "Device"}
+	return fmt.Sprintf("%s %s", adjectives[rand.Intn(len(adjectives))], nouns[rand.Intn(len(nouns))])
+}
+
+func printWelcome() {
+	fmt.Printf(`%s
+    ____        ______   ____  ____ 
+   / __ \____ _/ __/ /_ / __ \/ __ )
+  / /_/ / __ '/ /_/ __// / / / __  |
+ / _, _/ /_/ / __/ /_ / /_/ / /_/ / 
+/_/ |_|\__,_/_/  \__//_____/_____/  
+                                    
+%sWelcome to RaftDB CLI Monitor.%s
+Type 'help' for commands.
+`, ColorCyan, ColorBold, ColorReset)
+}
+
+func printHelp() {
+	fmt.Printf(`
+%sCommands:%s
+  %shelp%s                     Show this help
+  %sseed%s                     Insert 1000 random records
+  %sadd <key> <value>%s        Add a new record
+  %scount <query>%s            Count records matching query
+  %s<query>%s                  Execute a RaftDB query
+  %sexit / quit%s              Exit the CLI
+
+%sQuery Examples:%s
+  key like "user:*"                    Find all users
+  key = "user:10"                      Find specific user
+  value like "*admin*"                 Find full-text matches (admin role)
+  value like "*ERROR*"                 Find error logs
+  key like "product:*" LIMIT 5         List 5 products
+  key like "log:*" OFFSET 10 LIMIT 5   Pagination example
+  
+%sCount Examples:%s
+  count key like "user:*"              Count total users
+  count value like "*admin*"           Count admins
+`, ColorYellow, ColorReset, 
+ColorGreen, ColorReset,
+ColorGreen, ColorReset,
+ColorGreen, ColorReset,
+ColorGreen, ColorReset,
+ColorGreen, ColorReset,
+ColorGreen, ColorReset,
+ColorYellow, ColorReset,
+ColorYellow, ColorReset)
+}