From 42d2f519997dd50ee70af3625cc21257b6e4ee0e Mon Sep 17 00:00:00 2001
From: Tyson Dai <100662199+Murkeee@users.noreply.github.com>
Date: Mon, 12 Aug 2024 16:55:36 +1000
Subject: [PATCH] Add apiKey.js script to add API keys to database and update
 usage count. (#43)

* API key database handling

- Add tools/apiKey.js
- Modified database/index.js
- Modified database/types.js

Signed-off-by: Tyson Dai <100662199+Murkeee@users.noreply.github.com>

* Fixed function name in apiKey.js

changed function name from test() to add keyToTable()

Signed-off-by: Tyson Dai <100662199+Murkeee@users.noreply.github.com>

---------

Signed-off-by: Tyson Dai <100662199+Murkeee@users.noreply.github.com>
---
 database/index.js | 11 +++++++---
 database/types.js |  3 ++-
 tools/apiKey.js   | 55 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 tools/apiKey.js

diff --git a/database/index.js b/database/index.js
index 70f249a..f367b92 100644
--- 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
index 2e04f93..97f975a 100644
--- 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';
\ No newline at end of file
+export const DATASET_TABLE = 'dataset';
+export const API_KEY_TABLE = 'api_tokens'
\ No newline at end of file
diff --git a/tools/apiKey.js b/tools/apiKey.js
new file mode 100644
index 0000000..3d272b5
--- /dev/null
+++ 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!"
+    }
+}
\ No newline at end of file