Fixes #484, support arbitary URL scheme for repository

Change-Id: Idffb89945aef3eb191c3e005e84f1b7a9820ea1b
deepjavalibrary · Jan 5, 2021 · 6c26228 · 6c26228
1 parent ddda35c
commit 6c26228
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 39 deletions.
diff --git a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
@@ -103,15 +103,13 @@ public Repository newInstance(String name, String url) {
  }
  }
  return new SimpleRepository(name, path, names[0], names[1]);
- } else if ("http".equalsIgnoreCase(scheme) || "https".equalsIgnoreCase(scheme)) {
- if (FilenameUtils.isArchiveFile(fileName)) {
- return new SimpleUrlRepository(name, uri, names[0], names[1]);
- }
  } else if ("jar".equals(scheme)) {
  if (!FilenameUtils.isArchiveFile(fileName)) {
  throw new IllegalArgumentException("Only archive file is supported for res URL.");
  }
  return new JarRepository(name, uri, names[0], names[1]);
+ } else if (FilenameUtils.isArchiveFile(fileName)) {
+ return new SimpleUrlRepository(name, uri, names[0], names[1]);
  }
  return new RemoteRepository(name, uri);
  }

diff --git a/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java b/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java
@@ -116,43 +116,48 @@ private synchronized Metadata getMetadata() throws IOException {
  return metadata;
  }
 
- HttpURLConnection conn = null;
- try {
- resolved = true;
- conn = (HttpURLConnection) uri.toURL().openConnection();
- conn.setRequestMethod("HEAD");
- int code = conn.getResponseCode();
- if (code != 200) {
- logger.info("request error: " + code);
- return null;
- }
- int contentLength = conn.getContentLength();
-
- Artifact artifact = new Artifact();
- Map<String, Artifact.Item> files = new ConcurrentHashMap<>();
- Artifact.Item item = new Artifact.Item();
- item.setUri(uri.getPath());
- item.setName(""); // avoid creating extra folder
- item.setArtifact(artifact);
- item.setSize(contentLength);
- files.put(artifactId, item);
- artifact.setFiles(files);
- artifact.setName(modelName);
-
- metadata = new Metadata.MatchAllMetadata();
- metadata.setApplication(Application.UNDEFINED);
- metadata.setGroupId(DefaultModelZoo.GROUP_ID);
- metadata.setArtifactId(artifactId);
- metadata.setArtifacts(Collections.singletonList(artifact));
- String hash = md5hash(uri.toString());
- MRL mrl = MRL.model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
- metadata.setRepositoryUri(mrl.toURI());
+ Artifact artifact = new Artifact();
+ Map<String, Artifact.Item> files = new ConcurrentHashMap<>();
+ Artifact.Item item = new Artifact.Item();
+ item.setUri(uri.getPath());
+ item.setName(""); // avoid creating extra folder
+ item.setArtifact(artifact);
+ item.setSize(getContentLength());
+ files.put(artifactId, item);
+ artifact.setFiles(files);
+ artifact.setName(modelName);
+
+ metadata = new Metadata.MatchAllMetadata();
+ metadata.setApplication(Application.UNDEFINED);
+ metadata.setGroupId(DefaultModelZoo.GROUP_ID);
+ metadata.setArtifactId(artifactId);
+ metadata.setArtifacts(Collections.singletonList(artifact));
+ String hash = md5hash(uri.toString());
+ MRL mrl = MRL.model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
+ metadata.setRepositoryUri(mrl.toURI());
+ return metadata;
+ }
 
- return metadata;
- } finally {
- if (conn != null) {
- conn.disconnect();
+ private long getContentLength() throws IOException {
+ String scheme = uri.getScheme();
+ if ("http".equalsIgnoreCase(scheme) || "https".equalsIgnoreCase(scheme)) {
+ HttpURLConnection conn = null;
+ try {
+ resolved = true;
+ conn = (HttpURLConnection) uri.toURL().openConnection();
+ conn.setRequestMethod("HEAD");
+ int code = conn.getResponseCode();
+ if (code != 200) {
+ logger.info("request error: " + code);
+ return -1;
+ }
+ return conn.getContentLength();
+ } finally {
+ if (conn != null) {
+ conn.disconnect();
+ }
  }
  }
+ return -1;
  }
 }