From ed28bc7755b1e08195699775ebc5d2b8aeca3108 Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Mon, 11 Mar 2024 22:08:18 -0700
Subject: [PATCH] [api] Allows to use relative jar uri for cache folder name

Fixes: #3025
---
 .../java/ai/djl/repository/JarRepository.java   | 17 +++++++++++++----
 .../djl/repository/RepositoryFactoryImpl.java   | 10 ++++++----
 .../ai/djl/repository/JarRepositoryTest.java    |  8 +++++++-
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/api/src/main/java/ai/djl/repository/JarRepository.java b/api/src/main/java/ai/djl/repository/JarRepository.java
index b1b443d01f0..1f610f99bb2 100644
--- a/api/src/main/java/ai/djl/repository/JarRepository.java
+++ b/api/src/main/java/ai/djl/repository/JarRepository.java
@@ -42,13 +42,16 @@ public class JarRepository extends AbstractRepository {
     private String artifactId;
     private String modelName;
     private String queryString;
+    private String originalUri;
 
     private Metadata metadata;
     private boolean resolved;
 
-    JarRepository(String name, URI uri, String fileName, String queryString) {
+    JarRepository(String name, URI uri, String fileName, URI realUri) {
         super(name, uri);
-        this.queryString = queryString;
+        this.uri = realUri;
+        queryString = uri.getRawQuery();
+        originalUri = uri.toString();
         modelName = arguments.get("model_name");
         artifactId = arguments.get("artifact_id");
         if (artifactId == null) {
@@ -123,8 +126,14 @@ private synchronized Metadata getMetadata() {
         metadata = new Metadata.MatchAllMetadata();
         metadata.setArtifactId(artifactId);
         metadata.setArtifacts(Collections.singletonList(artifact));
-        String hash =
-                Utils.hash(queryString == null ? uri.toString() : uri.toString() + queryString);
+        String hashKey;
+        if (Boolean.parseBoolean(arguments.get("ignore_real_uri"))) {
+            hashKey = originalUri;
+        } else {
+            hashKey = queryString == null ? uri.toString() : uri.toString() + queryString;
+        }
+
+        String hash = Utils.hash(hashKey);
         MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
         metadata.setRepositoryUri(mrl.toURI());
 
diff --git a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
index a730a57cb54..eb5ab68a703 100644
--- a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
+++ b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
@@ -144,7 +144,6 @@ private static final class JarRepositoryFactory implements RepositoryFactory {
         @Override
         public Repository newInstance(String name, URI uri) {
             String p = uri.getPath();
-            String queryString = uri.getRawQuery();
             if (p.startsWith("/")) {
                 p = p.substring(1);
             }
@@ -152,19 +151,22 @@ public Repository newInstance(String name, URI uri) {
             if (u == null) {
                 throw new IllegalArgumentException("Resource not found: " + uri);
             }
+
+            URI realUri;
             try {
-                uri = u.toURI();
+                // resolve real uri: jar:file:/path/my_lib.jar!/model.zip
+                realUri = u.toURI();
             } catch (URISyntaxException e) {
                 throw new IllegalArgumentException("Resource not found: " + uri, e);
             }
 
-            Path path = Paths.get(parseFilePath(uri));
+            Path path = Paths.get(parseFilePath(realUri));
             String fileName = path.toFile().getName();
             if (FilenameUtils.isArchiveFile(fileName)) {
                 fileName = FilenameUtils.getNamePart(fileName);
             }
 
-            return new JarRepository(name, uri, fileName, queryString);
+            return new JarRepository(name, uri, fileName, realUri);
         }
 
         /** {@inheritDoc} */
diff --git a/api/src/test/java/ai/djl/repository/JarRepositoryTest.java b/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
index 4599214fab5..c1370d1da69 100644
--- a/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
+++ b/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
@@ -45,7 +45,7 @@ public void testResource() throws IOException {
         URL[] url = {jarFile.toUri().toURL()};
         try {
             Thread.currentThread().setContextClassLoader(new URLClassLoader(url));
-            Repository repo = Repository.newInstance("test", "jar:///test.zip?hash=1");
+            Repository repo = Repository.newInstance("test", "jar:///test.zip");
             Assert.assertEquals("test", repo.getName());
             Assert.assertTrue(repo.isRemote());
 
@@ -55,6 +55,12 @@ public void testResource() throws IOException {
             Artifact artifact = repo.resolve(list.get(0), null);
             repo.prepare(artifact);
             Assert.assertEquals(1, artifact.getFiles().size());
+
+            repo = Repository.newInstance("test", "jar:///test.zip?ignore_real_uri=true");
+            list = repo.getResources();
+            artifact = repo.resolve(list.get(0), null);
+            Path p = repo.getResourceDirectory(artifact);
+            Assert.assertFalse(Files.exists(p));
         } finally {
             Thread.currentThread().setContextClassLoader(null);
         }