246 lines
7.3 KiB
Nim
246 lines
7.3 KiB
Nim
## Ingest a file's content into Sonic
|
|
import os
|
|
import strutils
|
|
import sonic
|
|
|
|
var
|
|
verbose: bool = false
|
|
|
|
const
|
|
USAGE_DOC = staticRead("static/usage.txt")
|
|
|
|
proc showUsage(msg = "") =
|
|
## Display a usage message and quit
|
|
if msg != "":
|
|
stderr.writeLine "ERROR: $#\n" % msg
|
|
quit USAGE_DOC.replace("${app}", getAppFilename()),
|
|
(if msg == "": 0 else: 10)
|
|
|
|
proc envParam(name: string; default = ""): string =
|
|
## Return an environment parameter or return an error
|
|
let eVal = getEnv(name, default)
|
|
if eVal == "":
|
|
quit("Expected a value: $" & name, 1)
|
|
return eVal
|
|
|
|
proc getChannel(mode = SonicChannel.Ingest): Sonic =
|
|
## Return a channel based on command line params
|
|
let
|
|
host = envParam("SONIC_HOST")
|
|
port = envParam("SONIC_PORT")
|
|
secret = envParam("SONIC_SECRET")
|
|
try:
|
|
return open(host, port.parseInt(), secret, mode)
|
|
except ValueError:
|
|
let err = getCurrentExceptionMsg()
|
|
quit "Not a valid number: $#\n$#" % [port, err], 1
|
|
|
|
proc consolidate(channel: Sonic) =
|
|
## Trigger a consolidation
|
|
discard channel.execCommand("TRIGGER", @["consolidate"])
|
|
|
|
proc close(chn: Sonic) =
|
|
## Close a sonic channel
|
|
let outp = chn.quit()
|
|
if verbose:
|
|
stderr.writeLine "<closing $# channel: #$>" % [$(chn.channel), outp]
|
|
|
|
proc intAt(args: openArray[string]; pos: int; default: int): int =
|
|
## Parse a positional parameter if it exists, if not use default
|
|
if args.len-1 < pos or args[pos] == "":
|
|
return default
|
|
try:
|
|
return args[pos].parseInt()
|
|
except:
|
|
let err = getCurrentExceptionMsg()
|
|
quit("Not a valid number: $#\n$#" % [args[pos], err], 1)
|
|
|
|
################################################################
|
|
## Execute the user-facing commands
|
|
##
|
|
proc cmdPing() =
|
|
## Execute the "ping" command
|
|
let
|
|
chn = getChannel(SonicChannel.Control)
|
|
response = chn.execCommand("PING")
|
|
chn.close()
|
|
quit response, if response == "PONG": 0 else: 1
|
|
|
|
proc cmdCount(collection, bucket, objId: string) =
|
|
## Return indexed search data count for collection/bucket/objId
|
|
let
|
|
chn = getChannel(SonicChannel.Search)
|
|
response = chn.count(collection, bucket, objId)
|
|
chn.close()
|
|
quit $response, 0
|
|
|
|
proc cmdPush(collection, bucket, obj, data: string) =
|
|
## Ingest a file's content into a Sonic instance
|
|
var
|
|
justOne = false
|
|
stream = stdin
|
|
|
|
if data.len == 0:
|
|
quit "Data is an empty string", 1
|
|
elif data == "-":
|
|
stderr.writeLine "Reading from <stdin>"
|
|
elif data[0] == '@':
|
|
let filename = data[1 ..< data.len]
|
|
try:
|
|
stream = open(filename, bufSize=8000)
|
|
except:
|
|
let err = getCurrentExceptionMsg()
|
|
quit "Can't open \"$#\":\n$#" % [filename, err], 10
|
|
else:
|
|
justOne = true
|
|
|
|
var
|
|
ingCh = getChannel(SonicChannel.Ingest)
|
|
ctlCh = getChannel(SonicChannel.Control)
|
|
rMsg = ""
|
|
rCode = 0
|
|
|
|
if justOne:
|
|
let pushedOk = ingCh.push(collection, bucket, obj, data)
|
|
ctlCh.consolidate()
|
|
rMsg = if pushedOk: "" else: "push command returned a warning"
|
|
rCode = if pushedOk: 0 else: 1
|
|
else:
|
|
let
|
|
objIdPrefix = if data[0] == '@': "$#/$#:" % [obj, data[1 ..< data.len]]
|
|
else: obj & ":"
|
|
var
|
|
line = newStringOfCap(256)
|
|
count = 0
|
|
stderr.write("push: ")
|
|
while stream.readLine(line):
|
|
let
|
|
objId = objIdPrefix & $count
|
|
pushed = ingCh.push(collection, bucket, objId, line)
|
|
stderr.write(if pushed: "." else: "x")
|
|
inc count
|
|
if (count mod 31) == 0:
|
|
ctlCh.consolidate()
|
|
stderr.write("#")
|
|
stderr.write("\n")
|
|
if data[0] == '@':
|
|
close(stream)
|
|
ctlCh.close()
|
|
ingCh.close()
|
|
quit rMsg, rCode
|
|
|
|
proc cmdPop(collection, bucket, obj, data: string) =
|
|
## Pop search data from the given collection/bucket/obj
|
|
if data.len > 0:
|
|
let
|
|
ingCh = getChannel(SonicChannel.Ingest)
|
|
ctlCh = getChannel(SonicChannel.Control)
|
|
popOut = ingCh.pop(collection, bucket, obj, data)
|
|
ctlCh.consolidate()
|
|
ctlCh.close()
|
|
ingCh.close()
|
|
quit $popOut, 0
|
|
quit "Data is an empty string", 1
|
|
|
|
proc cmdQuery(collection, bucket, terms: string; limit, offset: int) =
|
|
## Query the indexes, echo the results to stdout
|
|
let
|
|
srChn = getChannel(SonicChannel.Search)
|
|
results = srChn.query(collection, bucket, terms, limit, offset)
|
|
srChn.close()
|
|
quit(results.join("\n"), 0)
|
|
|
|
proc cmdSuggest(collection, bucket, word: string; limit: int) =
|
|
## Query suggestions based on the word
|
|
let
|
|
srChn = getChannel(SonicChannel.Search)
|
|
results = srChn.suggest(collection, bucket, word, limit)
|
|
srChn.close()
|
|
quit(results.join("\n"), 0)
|
|
|
|
proc cmdFlush(collection: string; bucket=""; objId="") =
|
|
## Flushes all indexed data for the given collection, bucket or object
|
|
let
|
|
cnChn = getChannel(SonicChannel.Control)
|
|
results = cnChn.flush(collection, bucket, objId)
|
|
cnChn.close()
|
|
quit($results, 0)
|
|
|
|
################################################################
|
|
## Parse the command line and dispatch appropriate actions
|
|
##
|
|
proc main() =
|
|
## Parse the command line, dispatch the appropriate actions
|
|
var args = commandLineParams()
|
|
|
|
if args.len == 0 or "-h" in args or "--help" in args:
|
|
showUsage()
|
|
|
|
verbose = ("-v" in args) or ("--verbose" in args)
|
|
if verbose:
|
|
while "-v" in args:
|
|
args.delete(args.find("-v"))
|
|
while "--verbose" in args:
|
|
args.delete(args.find("--verbose"))
|
|
|
|
case args[0]:
|
|
of "help":
|
|
showUsage()
|
|
|
|
of "ping": # ping
|
|
if args.len > 1:
|
|
showUsage("'ping' takes no arguments")
|
|
cmdPing()
|
|
|
|
of "count": # count <collection> [bucket [object]]
|
|
let
|
|
bucket = (if args.len >= 3: args[2] else: "")
|
|
objId = (if args.len >= 4: args[3] else: "")
|
|
if args.len > 4:
|
|
showUsage("Too many arguments for 'count'")
|
|
cmdCount(args[1], bucket, objId)
|
|
|
|
of "push": # push <collection> <bucket> <object> "data|@filename|-"
|
|
if args.len != 5:
|
|
let pre = if args.len < 5: "Missing" else: "Too many"
|
|
showUsage(pre & " arguments for 'push'")
|
|
cmdPush(args[1], args[2], args[3], args[4])
|
|
|
|
of "pop": # pop <collection> <bucket> <object> "data"
|
|
if args.len != 5:
|
|
let pre = if args.len < 5: "Missing" else: "Too many"
|
|
showUsage(pre & " arguments for 'pop'")
|
|
cmdPop(args[1], args[2], args[3], args[4])
|
|
|
|
of "query": # query <collection> <bucket> "terms" [limit=10] [offset=0]
|
|
if args.len < 4 or args.len > 6:
|
|
let pre = if args.len < 4: "Missing" else: "Too many"
|
|
showUsage(pre & " arguments for 'query'")
|
|
let
|
|
limit = args.intAt(4, 10)
|
|
offset = args.intAt(5, 0)
|
|
cmdQuery(args[1], args[2], args[3], limit, offset)
|
|
|
|
of "suggest": # query <collection> <bucket> "word" [limit=10]
|
|
if args.len < 4 or args.len > 5:
|
|
let pre = if args.len < 4: "Missing" else: "Too many"
|
|
showUsage(pre & " arguments for 'suggest'")
|
|
let
|
|
limit = args.intAt(4, 10)
|
|
cmdSuggest(args[1], args[2], args[3], limit)
|
|
|
|
of "flush": # flush <collection> [bucket [object]]
|
|
if args.len < 3 or args.len > 5:
|
|
let pre = if args.len < 3: "Missing" else: "Too many"
|
|
showUsage(pre & " arguments for 'flush'")
|
|
let
|
|
bucket = (if args.len >= 3: args[3] else: "")
|
|
objId = (if args.len >= 4: args[4] else: "")
|
|
cmdFlush(args[2], bucket, objId)
|
|
|
|
else:
|
|
showUsage("Unknown command: " & args[0])
|
|
|
|
main()
|
|
# Fin
|