feat(gate): add cache to store used urls and add file upload at tg_deploy

Author: destifo

typegate/src/engine/query_engine.ts

@@ -114,6 +114,13 @@ export interface EndpointToSchemaMap {
   [index: string]: { fnName: string; outputSchema: unknown };
 }
 
+export interface UploadUrlMeta {
+  fileName: string;
+  fileHash: string;
+  fileSizeInBytes: number;
+  urlUsed: boolean;
+}
+
 export class QueryEngine {
   name: string;
   queryCache: QueryCache;
@@ -131,6 +138,7 @@ export class QueryEngine {
       }
     >
   >;
+  fileUploadUrlCache: Map<string, UploadUrlMeta> = new Map();
 
   get rawName(): string {
     return this.tg.rawName;

typegate/src/services/file_upload_service.ts

@@ -1,50 +1,94 @@
 // Copyright Metatype OÜ, licensed under the Elastic License 2.0.
 // SPDX-License-Identifier: Elastic-2.0
 
+import { UploadUrlMeta } from "../engine/query_engine.ts";
+import { QueryEngine } from "../engine/query_engine.ts";
+
+interface TypegraphFile {
+  name: string;
+  fileHash: string;
+  fileSizeInBytes: number;
+}
+
 function createUploadPath(origin: string, typegraphName: string) {
   const rand_path = crypto.randomUUID();
-  return `${origin}/${typegraphName}/files/${rand_path}`;
+  return `${origin}/upload-files/${typegraphName}/files/${rand_path}`;
 }
 
-export function handleUploadUrl(request: Request, typegraphName: string) {
+export async function handleUploadUrl(request: Request, engine: QueryEngine) {
   const url = new URL(request.url);
   const origin = url.origin;
+  const { name, fileHash, fileSizeInBytes }: TypegraphFile = await request
+    .json();
 
-  const uploadUrl = createUploadPath(origin, typegraphName);
-  return new Response(uploadUrl);
-}
+  const uploadUrlMeta: UploadUrlMeta = {
+    fileName: name,
+    fileHash: fileHash,
+    fileSizeInBytes: fileSizeInBytes,
+    urlUsed: false,
+  };
 
-interface TypegraphFile {
-  name: string;
-  fileHash: string;
-  fileSizeInBytes: number;
-  file: string;
+  const uploadUrl = createUploadPath(origin, engine.name);
+
+  engine.fileUploadUrlCache.set(uploadUrl, uploadUrlMeta);
+
+  return new Response(JSON.stringify({ uploadUrl: [uploadUrl] }));
 }
 
 export async function handleFileUpload(
   request: Request,
-  typegraphName: string,
+  engine: QueryEngine,
 ) {
   const url = new URL(request.url);
-  if (request.method !== "POST") {
+  if (request.method !== "PUT") {
     throw new Error(
       `${url.pathname} does not support ${request.method} method`,
     );
   }
 
-  const fileStorageDir = `metatype_artifacts/${typegraphName}/files`;
-  const { name, fileHash, fileSizeInBytes, file }: TypegraphFile = await request
-    .json();
+  const uploadMeta = engine.fileUploadUrlCache.get(url.toString());
+
+  if (!uploadMeta) {
+    throw new Error(`Endpoint ${url.toString()} does not exist`);
+  }
 
-  const fileData = Uint8Array.from(atob(file), (c) => c.charCodeAt(0));
+  const { fileName, fileHash, fileSizeInBytes, urlUsed }: UploadUrlMeta =
+    uploadMeta!;
 
-  if (fileData.length !== fileSizeInBytes) {
-    throw new Error();
+  if (urlUsed) {
+    throw new Error(`Endpoint ${url.toString()} is disabled`);
   }
 
+  const reader = request.body?.getReader()!;
+
+  let fileData = new Uint8Array();
+  let bytesRead = 0;
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) {
+      break;
+    }
+    if (value) {
+      bytesRead += value.length;
+      const temp = new Uint8Array(fileData.length + value.length);
+      temp.set(fileData);
+      temp.set(value, fileData.length);
+      fileData = temp;
+    }
+  }
+
+  if (bytesRead !== fileSizeInBytes) {
+    throw new Error("File size does not match");
+  }
+
+  const fileStorageDir = `metatype_artifacts/${engine.name}/files`;
   await Deno.mkdir(fileStorageDir, { recursive: true });
-  const filePath = `${fileStorageDir}/${name}.${fileHash}`;
+  const filePath = `${fileStorageDir}/${fileName}.${fileHash}`;
   await Deno.writeFile(filePath, fileData);
 
+  // mark as the url used once the request completes.
+  uploadMeta.urlUsed = true;
+  engine.fileUploadUrlCache.set(url.toString(), uploadMeta);
+
   return new Response();
 }

typegate/src/typegate/mod.ts

@@ -134,11 +134,11 @@ export class Typegate {
       }
 
       if (serviceName === "get-upload-url") {
-        return handleUploadUrl(request, engineName);
+        return handleUploadUrl(request, engine);
       }
 
       if (serviceName === "upload-files") {
-        return handleFileUpload(request, engineName);
+        return handleFileUpload(request, engine);
       }
 
       if (serviceName === "auth") {

typegraph/python/typegraph/graph/tg_deploy.py

@@ -1,14 +1,14 @@
 # Copyright Metatype OÜ, licensed under the Mozilla Public License Version 2.0.
 # SPDX-License-Identifier: MPL-2.0
 
+import json
+from base64 import b64encode
 from dataclasses import dataclass
 from typing import Dict, Optional, Union
 from urllib import request
-import json
 from typegraph.graph.shared_types import BasicAuth
 from typegraph.wit import ArtifactResolutionConfig
 from typegraph.gen.types import Err
-
 from typegraph.graph.typegraph import TypegraphOutput
 from typegraph.gen.exports.utils import QueryDeployParams
 from typegraph.wit import store, wit_utils
@@ -39,6 +39,13 @@ class RemoveResult:
     typegate: Union[Dict[str, any], str]
 
 
+@dataclass
+class UploadArtifactMeta:
+    name: str
+    file_hash: str
+    file_size_in_bytes: int
+
+
 def tg_deploy(tg: TypegraphOutput, params: TypegraphDeployParams) -> DeployResult:
     sep = "/" if not params.base_url.endswith("/") else ""
     url = params.base_url + sep + "typegate"
@@ -47,6 +54,35 @@ def tg_deploy(tg: TypegraphOutput, params: TypegraphDeployParams) -> DeployResul
     if params.auth is not None:
         headers["Authorization"] = params.auth.as_header_value()
     serialized = tg.serialize(params.artifacts_config)
+
+    # upload the referred files
+    ref_files = core.get_ref_files(store)
+    if isinstance(ref_files, Err):
+        raise Exception(ref_files.value)
+
+    get_upload_url = params.base_url + sep + "get-upload-url"
+    for file_hash, file_path in ref_files.value:
+        with open(file_path, "rb") as file:
+            file_content = file.read()
+            artifact = UploadArtifactMeta(
+                name=file.name,
+                file_hash=file_hash,
+                file_size_in_bytes=len(file_content),
+            )
+            req = request.Request(
+                url=get_upload_url, method="GET", headers=headers, data=artifact
+            )
+
+            response = handle_response(request.urlopen(req).read().decode())
+            file_upload_url = response["uploadUrl"]
+            upload_req = request.Request(
+                url=file_upload_url,
+                method="PUT",
+                data=file_content,
+                headers={"Content-Type": "application/octet-stream"},
+            )
+            response = request.urlopen(upload_req)
+
     res = wit_utils.gql_deploy_query(
         store,
         params=QueryDeployParams(