Added RiakClassLoader and helper (undocumented).

spring-data-riak/src/main/java/org/springframework/data/keyvalue/riak/util/RiakClassFileLoader.java

@@ -0,0 +1,153 @@
+/*
+ * Copyright (c) 2011 by J. Brisbin <jon@jbrisbin.com>
+ * Portions (c) 2011 by NPC International, Inc. or the
+ * original author(s).
+ *
+ * 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.
+ */
+
+package org.springframework.data.keyvalue.riak.util;
+
+import org.apache.commons.cli.*;
+import org.springframework.data.keyvalue.riak.core.RiakTemplate;
+
+import java.io.*;
+import java.net.URLEncoder;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+
+/**
+ * @author J. Brisbin <jon@jbrisbin.com>
+ */
+public class RiakClassFileLoader {
+
+  static Options opts = new Options();
+
+  static {
+    opts.addOption("v", false, "Verbose output");
+    opts.addOption("u",
+        true,
+        "URL to Riak (defaults to: 'http://localhost:8098/riak/{bucket}/{key}')");
+    opts.addOption("b", true, "Bucket to load class files into");
+    opts.addOption("k", true, "Key under which to store an individual class file");
+    opts.addOption("j", true, "JAR file to load into Riak");
+    opts.addOption("c", true, "Class file to load into Riak");
+    opts.addOption("d", true, "Directory from which to load all JAR files into Riak");
+  }
+
+  public static void main(String[] args) {
+    Parser p = new BasicParser();
+    CommandLine cl = null;
+    try {
+      cl = p.parse(opts, args);
+    } catch (ParseException e) {
+      System.err.println("Error parsing command line: " + e.getMessage());
+    }
+
+    if (null != cl) {
+      boolean verbose = cl.hasOption('v');
+      RiakTemplate riak = new RiakTemplate();
+      riak.getRestTemplate().setErrorHandler(new Ignore404sErrorHandler());
+      if (cl.hasOption('u')) {
+        riak.setDefaultUri(cl.getOptionValue('u'));
+      }
+      try {
+        riak.afterPropertiesSet();
+      } catch (Exception e) {
+        System.err.println("Error creating RiakTemplate: " + e.getMessage());
+      }
+      String[] files = cl.getOptionValues('j');
+      if (null != files) {
+        for (String file : files) {
+          if (verbose) {
+            System.out.println(String.format("Loading JAR file %s into Riak...", file));
+          }
+          try {
+            File zfile = new File(file);
+            ZipInputStream zin = new ZipInputStream(new FileInputStream(zfile));
+            ZipEntry entry;
+            while (null != (entry = zin.getNextEntry())) {
+              ByteArrayOutputStream bout = new ByteArrayOutputStream();
+              byte[] buff = new byte[16384];
+              for (int bytesRead = zin.read(buff); bytesRead > 0; bytesRead = zin.read(buff)) {
+                bout.write(buff, 0, bytesRead);
+              }
+
+              if (entry.getName().endsWith(".class")) {
+                String name = entry.getName().replaceAll("/", ".");
+                name = URLEncoder.encode(name.substring(0, name.length() - 6), "UTF-8");
+                String bucket;
+                if (cl.hasOption('b')) {
+                  bucket = cl.getOptionValue('b');
+                } else {
+                  bucket = URLEncoder.encode(zfile.getCanonicalFile().getName(), "UTF-8");
+                }
+                if (verbose) {
+                  System.out.println(String.format("Uploading to %s/%s", bucket, name));
+                }
+
+                // Load these bytes into Riak
+                riak.setAsBytes(bucket, name, bout.toByteArray());
+              }
+            }
+          } catch (FileNotFoundException e) {
+            System.err.println("Error reading JAR file: " + e.getMessage());
+          } catch (IOException e) {
+            System.err.println("Error reading JAR file: " + e.getMessage());
+          }
+        }
+      }
+
+      String[] classFiles = cl.getOptionValues('c');
+      if (null != classFiles) {
+        for (String classFile : classFiles) {
+          try {
+            FileInputStream fin = new FileInputStream(classFile);
+            ByteArrayOutputStream bout = new ByteArrayOutputStream();
+            byte[] buff = new byte[16384];
+            for (int bytesRead = fin.read(buff); bytesRead > 0; bytesRead = fin.read(buff)) {
+              bout.write(buff, 0, bytesRead);
+            }
+
+            String name;
+            if (cl.hasOption('k')) {
+              name = cl.getOptionValue('k');
+            } else {
+              throw new IllegalStateException(
+                  "Must specify a Riak key in which to store the data if loading individual class files.");
+            }
+            String bucket;
+            if (cl.hasOption('b')) {
+              bucket = cl.getOptionValue('b');
+            } else {
+              throw new IllegalStateException(
+                  "Must specify a Riak bucket in which to store the data if loading individual class files.");
+            }
+            if (verbose) {
+              System.out.println(String.format("Uploading to %s/%s", bucket, name));
+            }
+
+            // Load these bytes into Riak
+            riak.setAsBytes(bucket, name, bout.toByteArray());
+
+          } catch (FileNotFoundException e) {
+            System.err.println("Error reading class file: " + e.getMessage());
+          } catch (IOException e) {
+            System.err.println("Error reading class file: " + e.getMessage());
+          }
+        }
+      }
+    }
+  }
+
+}

spring-data-riak/src/main/java/org/springframework/data/keyvalue/riak/util/RiakClassLoader.java

@@ -0,0 +1,186 @@
+/*
+ * Copyright (c) 2011 by J. Brisbin <jon@jbrisbin.com>
+ * Portions (c) 2011 by NPC International, Inc. or the
+ * original author(s).
+ *
+ * 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.
+ */
+
+package org.springframework.data.keyvalue.riak.util;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.data.keyvalue.riak.core.RiakTemplate;
+import org.springframework.http.HttpInputMessage;
+import org.springframework.http.HttpOutputMessage;
+import org.springframework.http.MediaType;
+import org.springframework.http.converter.HttpMessageConverter;
+import org.springframework.http.converter.HttpMessageNotReadableException;
+import org.springframework.http.converter.HttpMessageNotWritableException;
+import org.springframework.web.client.RestTemplate;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.UnsupportedEncodingException;
+import java.net.URLEncoder;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * @author J. Brisbin <jon@jbrisbin.com>
+ */
+public class RiakClassLoader extends ClassLoader {
+
+  protected final Log log = LogFactory.getLog(getClass());
+  protected Set<String> buckets = new LinkedHashSet<String>();
+  protected RiakTemplate riakTemplate;
+  protected String defaultBucket = null;
+
+  public RiakClassLoader(ClassLoader classLoader, RiakTemplate riakTemplate) {
+    super(classLoader);
+    init(riakTemplate);
+    loadBucketsFromClassPath();
+  }
+
+  public RiakClassLoader(RiakTemplate riakTemplate) {
+    init(riakTemplate);
+    loadBucketsFromClassPath();
+  }
+
+  public Set<String> getBuckets() {
+    return buckets;
+  }
+
+  public void setBuckets(Set<String> buckets) {
+    this.buckets = buckets;
+  }
+
+  public RiakTemplate getRiakTemplate() {
+    return riakTemplate;
+  }
+
+  public void setRiakTemplate(RiakTemplate riakTemplate) {
+    this.riakTemplate = riakTemplate;
+  }
+
+  public String getDefaultBucket() {
+    return defaultBucket;
+  }
+
+  public void setDefaultBucket(String defaultBucket) {
+    this.defaultBucket = defaultBucket;
+  }
+
+  @Override
+  protected Class<?> findClass(String s) throws ClassNotFoundException {
+    Class<?> c;
+    try {
+      c = super.findClass(s);
+      if (log.isDebugEnabled()) {
+        log.debug(String.format("Found class '%s' locally defined.", s));
+      }
+    } catch (Throwable t) {
+      // Class not defined in this ClassLoader yet
+    }
+
+    Set<String> buckets = new LinkedHashSet<String>(this.buckets);
+    if (null != defaultBucket) {
+      buckets.add(defaultBucket);
+    }
+    for (String bucket : buckets) {
+      if (bucket.indexOf("/") < 0) {
+        try {
+          if (log.isDebugEnabled()) {
+            log.debug(String.format("Class '%s' not locally defined, trying Riak.", s));
+          }
+          byte[] buff = riakTemplate.getAsBytes(URLEncoder.encode(bucket, "UTF-8"), s);
+          c = defineClass(s, buff, 0, buff.length);
+          if (null != c) {
+            return c;
+          }
+        } catch (ClassFormatError ignored) {
+        } catch (UnsupportedEncodingException e) {
+          log.error(e.getMessage(), e);
+        }
+      }
+    }
+
+    // Nothing found
+    throw new ClassNotFoundException("Class not found: " + s);
+  }
+
+
+  protected void loadBucketsFromClassPath() {
+    String classPath = System.getProperty("java.class.path");
+    String pathSep = System.getProperty("path.separator", ":");
+    if (null != classPath) {
+      String[] paths = classPath.split(pathSep);
+      for (String p : paths) {
+        buckets.add(p);
+      }
+    }
+  }
+
+  protected void init(RiakTemplate riakTemplate) {
+    this.riakTemplate = riakTemplate;
+    RestTemplate tmpl = this.riakTemplate.getRestTemplate();
+    tmpl.getMessageConverters().add(0, new JavaSerializationMessageHandler());
+    tmpl.setErrorHandler(new Ignore404sErrorHandler());
+  }
+
+  private class JavaSerializationMessageHandler implements HttpMessageConverter {
+
+    public boolean canRead(Class clazz, MediaType mediaType) {
+      return MediaType.APPLICATION_OCTET_STREAM.equals(mediaType);
+    }
+
+    public boolean canWrite(Class clazz, MediaType mediaType) {
+      return null != clazz;
+    }
+
+    public List<MediaType> getSupportedMediaTypes() {
+      List<MediaType> types = new ArrayList<MediaType>(1);
+      types.add(MediaType.APPLICATION_OCTET_STREAM);
+      return types;
+    }
+
+    public Object read(java.lang.Class clazz, HttpInputMessage inputMessage) throws
+        IOException,
+        HttpMessageNotReadableException {
+      ObjectInputStream oin = new ObjectInputStream(inputMessage.getBody());
+      try {
+        Class<?> c = (Class<?>) oin.readObject();
+        if (log.isDebugEnabled()) {
+          log.debug("Loaded class: " + c);
+        }
+        return c;
+      } catch (ClassNotFoundException e) {
+        throw new IllegalStateException(e.getMessage(), e);
+      }
+    }
+
+    public void write(Object o, MediaType contentType, HttpOutputMessage outputMessage) throws
+        IOException,
+        HttpMessageNotWritableException {
+      outputMessage.getHeaders().setContentType(MediaType.APPLICATION_OCTET_STREAM);
+      ObjectOutputStream oout = new ObjectOutputStream(outputMessage.getBody());
+      oout.writeObject(o);
+      oout.flush();
+    }
+
+  }
+
+}

spring-data-riak/src/test/groovy/org/springframework/data/keyvalue/riak/core/RiakClassLoaderSpec.groovy

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2011 by J. Brisbin <jon@jbrisbin.com>
+ * Portions (c) 2011 by NPC International, Inc. or the
+ * original author(s).
+ *
+ * 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.
+ */
+
+package org.springframework.data.keyvalue.riak.core
+
+import org.springframework.data.keyvalue.riak.util.RiakClassFileLoader
+import org.springframework.data.keyvalue.riak.util.RiakClassLoader
+import spock.lang.Shared
+import spock.lang.Specification
+
+/**
+ * @author J. Brisbin <jon@jbrisbin.com>
+ */
+class RiakClassLoaderSpec extends Specification {
+
+  @Shared RiakTemplate riakTemplate = new RiakTemplate()
+
+  def setupSpec() {
+    RiakQosParameters qos = new RiakQosParameters()
+    qos.durableWriteThreshold = "all"
+    riakTemplate.defaultQosParameters = qos
+    riakTemplate.ignoreNotFound = true
+    riakTemplate.afterPropertiesSet()
+  }
+
+  def "Test load class file into Riak"() {
+
+    when:
+    def args = [
+        "-c", "src/test/classes/org/springframework/data/keyvalue/riak/core/ClassLoaderTest.class",
+        "-b", "test",
+        "-k", "org.springframework.data.keyvalue.riak.core.ClassLoaderTest"
+    ].toArray(new String[6])
+    RiakClassFileLoader.main(args)
+
+    then:
+    true
+
+  }
+
+  def "Test find class previously loaded into Riak"() {
+
+    given:
+    RiakClassLoader classLoader = new RiakClassLoader(riakTemplate)
+    classLoader.defaultBucket = "test"
+
+    when:
+    def clazz = Class.forName("org.springframework.data.keyvalue.riak.core.ClassLoaderTest", false, classLoader)
+    def inst = clazz?.newInstance()
+
+    then:
+    null != clazz
+    null != inst
+    inst.name == "ClassLoaderTest"
+
+  }
+
+}