update tgz artifact (#344)

* update tgz artifact * fix tests
kscalelabs · Aug 30, 2024 · 3a1a71b · 3a1a71b
1 parent 3f37292
commit 3a1a71b
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 34 deletions.
diff --git a/frontend/src/components/listing/artifacts/TgzArtifact.tsx b/frontend/src/components/listing/artifacts/TgzArtifact.tsx
@@ -1,33 +1,38 @@
-import React, { useState } from "react";
+import { useState } from "react";
 import { FaCheck, FaCopy } from "react-icons/fa";
 
+import { cx } from "class-variance-authority";
 import { components } from "gen/api";
 
 type SingleArtifactResponse = components["schemas"]["SingleArtifactResponse"];
 
-interface Props {
-  artifact: SingleArtifactResponse;
+interface ButtonProps {
+  command: string;
+  isFirst?: boolean;
+  prefix?: string;
 }
 
-const TgzArtifact = ({ artifact }: Props) => {
+const DownloadButton = ({ command, isFirst, prefix }: ButtonProps) => {
   const [copied, setCopied] = useState(false);
 
-  const command = `kscale urdf download ${artifact.artifact_id}`;
-
-  const copyToClipboard = () => {
-    navigator.clipboard.writeText(command).then(() => {
-      setCopied(true);
-      setTimeout(() => setCopied(false), 2000); // Reset after 2 seconds
-    });
+  const copyToClipboard = async (command: string) => {
+    setCopied(true);
+    await navigator.clipboard.writeText(command);
+    setTimeout(() => setCopied(false), 2000);
   };
 
   return (
     <button
-      className="bg-gray-100 p-3 rounded-md flex items-center justify-between cursor-pointer w-full text-left hover:bg-gray-200 transition-colors duration-200"
-      onClick={copyToClipboard}
+      className={cx(
+        "bg-gray-100 p-3 rounded-md flex items-center justify-between cursor-pointer w-full text-left hover:bg-gray-200 transition-colors duration-200",
+        isFirst ? "mt-0" : "mt-2",
+      )}
+      onClick={() => copyToClipboard(command)}
+      title={command}
+      disabled={copied}
     >
       <code className="text-xs text-gray-500 font-mono truncate flex-grow mr-2">
-        {command}
+        {prefix ? `${prefix}: ${command}` : command}
       </code>
       <span className="flex-shrink-0">
         {copied ? (
@@ -40,4 +45,24 @@ const TgzArtifact = ({ artifact }: Props) => {
   );
 };
 
+interface Props {
+  artifact: SingleArtifactResponse;
+}
+
+const TgzArtifact = ({ artifact }: Props) => {
+  const url = artifact.urls?.large;
+
+  return (
+    <>
+      <DownloadButton
+        command={`kscale urdf download ${artifact.artifact_id}`}
+        isFirst={true}
+      />
+      {url && (
+        <DownloadButton command={`${url}`} isFirst={false} prefix="URL" />
+      )}
+    </>
+  );
+};
+
 export default TgzArtifact;
diff --git a/store/app/crud/base.py b/store/app/crud/base.py
@@ -205,25 +205,13 @@ async def _list_me(
         sorted_items = sorted(response, key=sort_key, reverse=True)
         return sorted_items[(page - 1) * ITEMS_PER_PAGE : page * ITEMS_PER_PAGE], page * ITEMS_PER_PAGE < len(response)
 
-    async def _count_items(
-        self,
-        item_class: type[T],
-        expression_attribute_values: dict[str, Any] | None = None,
-    ) -> int:
+    async def _count_items(self, item_class: type[T]) -> int:
         table = await self.db.Table(TABLE_NAME)
-        if expression_attribute_values is None:
-            item_dict = await table.scan(
-                IndexName="type_index",
-                Select="COUNT",
-                FilterExpression=Key("type").eq(item_class.__name__),
-            )
-        else:
-            item_dict = await table.scan(
-                IndexName="type_index",
-                Select="COUNT",
-                FilterExpression=Key("type").eq(item_class.__name__),
-                ExpressionAttributeValues=expression_attribute_values,
-            )
+        item_dict = await table.scan(
+            IndexName="type_index",
+            Select="COUNT",
+            FilterExpression=Key("type").eq(item_class.__name__),
+        )
         return item_dict["Count"]
 
     def _validate_item(self, data: dict[str, Any], item_class: type[T]) -> T:

diff --git a/store/app/crud/users.py b/store/app/crud/users.py
@@ -4,9 +4,10 @@
 import warnings
 from typing import Any, Literal, overload
 
+from boto3.dynamodb.conditions import Key
 from botocore.exceptions import ClientError
 
-from store.app.crud.base import BaseCrud
+from store.app.crud.base import TABLE_NAME, BaseCrud
 from store.app.model import (
     APIKey,
     APIKeyPermissionSet,
@@ -148,13 +149,22 @@ async def add_api_key(
         source: APIKeySource,
         permissions: APIKeyPermissionSet,
     ) -> APIKey:
-        api_key_count = await self._count_items(APIKey, expression_attribute_values={":user_id": user_id})
+        api_key_count = await self.get_api_key_count(user_id)
         if api_key_count >= 10:
             raise ValueError("User has reached the maximum number of API keys (10)")
         api_key = APIKey.create(user_id=user_id, source=source, permissions=permissions)
         await self._add_item(api_key)
         return api_key
 
+    async def get_api_key_count(self, user_id: str) -> int:
+        table = await self.db.Table(TABLE_NAME)
+        item_dict = await table.scan(
+            IndexName="type_index",
+            Select="COUNT",
+            FilterExpression=Key("type").eq(APIKey.__name__) & Key("user_id").eq(user_id),
+        )
+        return item_dict["Count"]
+
     async def delete_api_key(self, token: APIKey | str) -> None:
         await self._delete_item(token)