Add apiKey.js script to add API keys to database and update usage cou…

…nt. (#43) * API key database handling - Add tools/apiKey.js - Modified database/index.js - Modified database/types.js Signed-off-by: Tyson Dai <[email protected]> * Fixed function name in apiKey.js changed function name from test() to add keyToTable() Signed-off-by: Tyson Dai <[email protected]> --------- Signed-off-by: Tyson Dai <[email protected]>
SkywardAI · Aug 12, 2024 · 42d2f51 · 42d2f51
1 parent dcb5c00
commit 42d2f51
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 4 deletions.
diff --git a/database/index.js b/database/index.js
@@ -16,11 +16,11 @@
 import { connect } from "@lancedb/lancedb";
 import { 
     Schema, Field, FixedSizeList, 
-    Float32, Utf8, 
+    Float32, Utf8, Int32,
     // eslint-disable-next-line
     Table 
 } from "apache-arrow";
-import { DATASET_TABLE, SYSTEM_TABLE } from "./types.js";
+import { API_KEY_TABLE, DATASET_TABLE, SYSTEM_TABLE } from "./types.js";
 
 const uri = "/tmp/lancedb/";
 const db = await connect(uri);
@@ -38,7 +38,12 @@ export async function initDB(force = false) {
         new Field("dataset_name", new Utf8()),
         new Field("question", new Utf8()),
         new Field("answer", new Utf8())
-    ]), open_options)
+    ]), open_options);
+    // create or re-open api key table
+    await db.createEmptyTable(API_KEY_TABLE, new Schema([
+        new Field("api_key", new Utf8()),
+        new Field("usage", new Int32()),
+    ]), open_options);
 }
 
 /**

diff --git a/database/types.js b/database/types.js
@@ -14,4 +14,5 @@
 // limitations under the License.
 
 export const SYSTEM_TABLE = 'system';
-export const DATASET_TABLE = 'dataset';
+export const DATASET_TABLE = 'dataset';
+export const API_KEY_TABLE = 'api_tokens'
diff --git a/tools/apiKey.js b/tools/apiKey.js
@@ -0,0 +1,55 @@
+// coding=utf-8
+
+// Copyright [2024] [SkywardAI]
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//        http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { API_KEY_TABLE } from "../database/types.js";
+import { getTable } from "../database/index.js";
+
+const MAX_USAGE = 10;
+const tbl = await getTable(API_KEY_TABLE);
+
+function queryApiKeyTbl(key) {
+    const keyQuery = tbl.query()
+        .where('api_key = '+"'"+key+"'")
+        .limit(1)
+        .toArray();
+    return keyQuery
+}
+
+export async function addKeytoTable(key){
+    const keyQuery = await queryApiKeyTbl(key);
+    if ( keyQuery.length == 0){
+        await tbl.add([{api_key: key, usage: MAX_USAGE}]);
+    }
+    else {
+        return "Cannot add key to table: Key already exists";
+    }
+}
+
+export async function updateKeyUsage(key){
+    const keyQuery = await queryApiKeyTbl(key);
+    console.log(keyQuery)
+    if (keyQuery.length > 0){
+        let usage = keyQuery[0].usage
+        console.log(usage);
+        usage--;
+        await tbl.update({
+            where: 'api_key = ' + "'" + key + "'",
+            values: {api_key: key, usage: usage}
+        });
+    }
+    else {
+        return "Could not find key in table!"
+    }
+}