Ability to specify a table name for kv commands (#16450)

Closes #15786 

## Release notes summary - What our users need to know

`std-rfc/kv` commands now take an optional `--table (-t)` argument which
allows a custom table name to be used. If not specified, the current
`std-kv-store` table will be used.

## Tasks after submitting

- [x] Added in-help example to `kv set`. Did not add examples to the
other `kv` commands since they should be obvious based on the initial
`kv set`.

## Additional details

Example:

```nushell
use std-rfc/kv *
kv set -t foo bar 5
kv list
# Empty list, since the value was placed in a custom table
# => ╭────────────╮
# => │ empty list │
# => ╰────────────╯
kv list -t foo
# => ╭───┬─────┬───────╮
# => │ # │ key │ value │
# => ├───┼─────┼───────┤
# => │ 0 │ bar │     5 │
# => ╰───┴─────┴───────╯
```

To protect against injection attacks, table names can only include
alphanumeric characters with `_` and `-`.

crates/nu-std/std-rfc/kv/mod.nu

@@ -1,3 +1,5 @@
+const default_table_name = 'std_kv_store'
+const valid_table_regex = '^[a-zA-Z0-9_]+$'
 # kv module
 #
 # use std-rfc/kv *
@@ -21,6 +23,9 @@
 @example "Store a number" {
   kv set foo 5
 }
+@example "Store a number in a custom table named foo" {
+  kv set -t foo bar 5
+}
 @example "Update a number and return it" {
   let $new_foo = (kv get foo | kv set foo { $in + 1 } --return value)
 }
@@ -34,6 +39,7 @@ export def "kv set" [
   value_or_closure?: any
   --return (-r): string   # Whether and what to return to the pipeline output
   --universal (-u)        # Store the key-value pair in a universal database
+  --table (-t): string = $default_table_name    # Optional table name
 ] {
   # Pipeline input is preferred, but prioritize
   # parameter if present. This allows $in to be
@@ -54,15 +60,15 @@ export def "kv set" [
     value: ($value | to nuon)
   }
 
-  let db_open = (db_setup --universal=$universal)
+  let db_open = (db_setup --universal=$universal --table=$table)
   try {
     # Delete the existing key if it does exist
-    do $db_open | query db "DELETE FROM std_kv_store WHERE key = :key" --params { key: $key }
+    do $db_open | query db $"DELETE FROM ($table) WHERE key = :key" --params { key: $key }
   }
 
   match $universal {
-    true  => { $kv_pair | into sqlite (universal_db_path) -t std_kv_store }
-    false => { $kv_pair | stor insert -t std_kv_store }
+    true  => { $kv_pair | into sqlite (universal_db_path) -t $table }
+    false => { $kv_pair | stor insert -t $table }
   }
 
   # The value that should be returned from `kv set`
@@ -80,8 +86,8 @@ export def "kv set" [
   # ---
   # all: The entire contents of the existing kv table are returned
   match ($return | default 'input') {
-    'all' => (kv list --universal=$universal)
-    'a' => (kv list --universal=$universal)
+    'all' => (kv list --universal=$universal --table=$table)
+    'a' => (kv list --universal=$universal --table=$table)
     'value' => $value
     'v' => $value
     'input' => $input
@@ -108,10 +114,11 @@ export def "kv set" [
 export def "kv get" [
   key: string       # Key of the kv-pair to retrieve
   --universal (-u)  # Whether to use the universal db
+  --table (-t): string = $default_table_name    # Optional table name
 ] {
-  let db_open = (db_setup --universal=$universal)
+  let db_open = (db_setup --universal=$universal --table=$table)
   do $db_open
-    | query db "SELECT value FROM std_kv_store WHERE key = :key" --params { key: $key }
+    | query db $"SELECT value FROM ($table) WHERE key = :key" --params { key: $key }
     | match $in {
       # Match should be exactly one row
       [$el] => { $el.value | from nuon }
@@ -125,10 +132,10 @@ export def "kv get" [
 # Returns results as the Nushell value rather than the stored nuon.
 export def "kv list" [
   --universal (-u)  # Whether to use the universal db
+  --table (-t): string = $default_table_name    # Optional table name
 ] {
-  let db_open = (db_setup --universal=$universal)
-
-  do $db_open | $in.std_kv_store? | each {|kv_pair|
+  let db_open = (db_setup --universal=$universal --table=$table)
+  do $db_open | $in | get -o $table | each {|kv_pair|
     {
       key: $kv_pair.key
       value: ($kv_pair.value | from nuon )
@@ -140,15 +147,16 @@ export def "kv list" [
 export def --env "kv drop" [
   key: string       # Key of the kv-pair to drop
   --universal (-u)  # Whether to use the universal db
+  --table (-t): string = $default_table_name    # Optional table name
 ] {
-  let db_open = (db_setup --universal=$universal)
+  let db_open = (db_setup --universal=$universal --table=$table)
 
-  let value = (kv get --universal=$universal $key)
+  let value = (kv get --universal=$universal --table=$table $key)
 
   try {
     do $db_open
       # Hack to turn a SQLiteDatabase into a table
-      | query db "DELETE FROM std_kv_store WHERE key = :key" --params { key: $key }
+      | query db $"DELETE FROM ($table) WHERE key = :key" --params { key: $key }
   }
 
   if $universal and ($env.NU_KV_UNIVERSALS? | default false) {
@@ -166,8 +174,20 @@ def universal_db_path [] {
 }
 
 def db_setup [
-  --universal   # Whether to use the universal db
+  --universal       # Whether to use the universal db
+  --table: string   # The table name to use
 ] : nothing -> closure {
+
+  if $table not-like $valid_table_regex {
+    error make {
+      msg: "Invalid table name"
+      label: {
+        text: "Table name must match regex: ${valid_table_regex}"
+        span: (metadata $table).span
+      }
+    }
+  }
+
   try {
     match $universal {
       true  => {
@@ -178,12 +198,12 @@ def db_setup [
           key: $uuid
           value: ''
         }
-        $dummy_record | into sqlite (universal_db_path) -t std_kv_store
-        open (universal_db_path) | query db "DELETE FROM std_kv_store WHERE key = :key" --params { key: $uuid }
+        $dummy_record | into sqlite (universal_db_path) -t $table
+        open (universal_db_path) | query db $"DELETE FROM ($table) WHERE key = :key" --params { key: $uuid }
       }
       false => {
         # Create the stor table if it doesn't exist
-        stor create -t std_kv_store -c {session: str, key: str, value: str} | ignore
+        stor create -t $table -c {session: str, key: str, value: str} | ignore
       }
     }
   }

crates/nu-std/tests/test_std-rfc_kv.nu

@@ -302,3 +302,17 @@ def universal-return_value_only [] {
     kv drop --universal $key
     rm $env.NU_UNIVERSAL_KV_PATH
 }
+
+@test
+def simple-local-set-table [] {
+    if ('sqlite' not-in (version).features) { return }
+
+    let key = (random uuid)
+
+    kv set -t foo $key 42
+    let actual = (kv get -t foo $key)
+    let expected = 42
+    assert equal $actual $expected
+
+    kv drop -t foo $key | ignore
+}