Mucking with Blueprints and Frames

BlueprintsGraphPersistenceRepository.java

package org.jboss.metrics.impl.persistence.blueprints;

import com.tinkerpop.blueprints.CloseableIterable;
import com.tinkerpop.blueprints.Graph;
import com.tinkerpop.blueprints.Index;
import com.tinkerpop.blueprints.IndexableGraph;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.frames.FramedGraph;
import com.tinkerpop.frames.FramedGraphFactory;
import org.jboss.metrics.spi.model.PersistentEntity;
import org.jboss.metrics.spi.persistence.NoSuchEntityException;
import org.jboss.metrics.spi.persistence.PersistenceRepository;
import org.jboss.metrics.spi.persistence.UnexpectedTypeException;

import java.util.Iterator;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Implementation of a {@link org.jboss.metrics.spi.persistence.PersistenceRepository}
 * backed by a Tinkerpop Blueprints Graph.  Threadsafe if supplied with a threadsafe backing
 * {@link com.tinkerpop.blueprints.Graph} on
 * {@link BlueprintsGraphPersistenceRepository#newInstance(com.tinkerpop.blueprints.IndexableGraph)}
 *
 * @author <a href="mailto:alr@jboss.org">Andrew Lee Rubinger</a>
 */
public class BlueprintsGraphPersistenceRepository implements PersistenceRepository {

    private static final Logger log = Logger.getLogger(BlueprintsGraphPersistenceRepository.class.getName());

    private static final String SUFFIX_INDEX = "-index";

    /**
     * Underlying storage engine
     */
    private final IndexableGraph graph;

    /**
     * Used for object/graph mapping
     */
    private final FramedGraphFactory framedGraphFactory;

    /**
     * Used in creating IDs for new entities
     */
    private final IdGenerator idGenerator;

    /**
     * Index for lookup by ID
     */
    private final Index<Vertex> idIndex;

    /**
     * Creates a new {@link BlueprintsGraphPersistenceRepository}
     * with the specified (required) backing
     * {@link com.tinkerpop.blueprints.Graph}
     *
     * @param graph The underlying graph
     * @return A new {@link BlueprintsGraphPersistenceRepository} instance
     * @throws IllegalArgumentException
     */
    public static BlueprintsGraphPersistenceRepository newInstance(final IndexableGraph graph) throws IllegalArgumentException {
        if (graph == null) {
            throw new IllegalArgumentException("graph must be specified");
        }
        return new BlueprintsGraphPersistenceRepository(graph);
    }

    private BlueprintsGraphPersistenceRepository(final IndexableGraph graph) {
        assert graph != null : "graph must be specified";
        this.graph = graph;
        this.framedGraphFactory = new FramedGraphFactory();
        this.idGenerator = new IdGenerator();
        // Apply an index to the IDs
        final String indexName = PersistentEntity.PROP_ID + SUFFIX_INDEX;
        Index<Vertex> idIndex = graph.getIndex(indexName, Vertex.class);
        if (idIndex == null) {
            idIndex = graph.createIndex(PersistentEntity.PROP_ID + SUFFIX_INDEX, Vertex.class);
            log.finer("Created ID index: " + idIndex);
        }
        this.idIndex = idIndex;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public <ENTITYTYPE extends PersistentEntity> ENTITYTYPE create(
            Class<ENTITYTYPE> typeToCreate) throws IllegalArgumentException {
        if (typeToCreate == null) {
            throw new IllegalArgumentException("type must be specified");
        }

        // Create a vertex
        final Vertex vertex = graph.addVertex(null);
        log.log(Level.FINER, "Created type " + typeToCreate.getName() + " with ID: " + vertex.getId());

        // Map to an object
        final ENTITYTYPE framedObject = this.frame(vertex, typeToCreate);

        // Set the ID
        final Object id = idGenerator.generateID(typeToCreate);
        framedObject.setId(id);

        // Apply indexing
        idIndex.put(PersistentEntity.PROP_ID, id, vertex);

        // Return
        return framedObject;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public <TYPE extends PersistentEntity> TYPE get(final Object id, final Class<TYPE> typeToRetrieve)
            throws IllegalArgumentException, UnexpectedTypeException, NoSuchEntityException {
        if (id == null) {
            throw new IllegalArgumentException("id must be specified");
        }

        final Vertex vertex = this.getVertex(id);
        final TYPE obj;
        obj = this.frame(vertex, typeToRetrieve);

        //TODO check for type assignment and throw UnexpectedTypeException if necessary
        return obj;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void remove(final Object id) throws IllegalArgumentException, NoSuchEntityException {
        if (id == null) {
            throw new IllegalArgumentException("id must be specified");
        }

        final Vertex vertex = this.getVertex(id);
        graph.removeVertex(vertex);
        idIndex.remove(PersistentEntity.PROP_ID,id,vertex);
    }

    private Vertex getVertex(final Object id) throws NoSuchEntityException {
        assert id != null : "id must be specified";
        // Find by ID in the index
        Vertex match = null;
        try (final CloseableIterable<Vertex> matches = idIndex.get(PersistentEntity.PROP_ID, id)) {
            final Iterator<Vertex> vertexes = matches.iterator();
            while (vertexes.hasNext()) {
                // non-unique fail-fast check
                if (match != null) {
                    throw new RuntimeException("Error exists in repo implementation; encountered non-unique ID: " + id);
                }
                match = vertexes.next();
            }
        }
        // None found
        if (match == null) {
            throw NoSuchEntityException.newInstance("No object with ID '" +
                    id.toString() + "' was found in the repository.");
        }

        return match;
    }

    /**
     * @param vertex Required
     * @param type   Required
     * @param <TYPE> Type of the object to be returned
     * @return
     */
    private <TYPE> TYPE frame(final Vertex vertex, Class<TYPE> type) {
        assert vertex != null : "vertex must be specified";
        assert type != null : "type must be specified";
        final FramedGraph<Graph> framedGraph = this.framedGraphFactory.create(graph);
        final TYPE obj = framedGraph.frame(vertex, type);
        return obj;
    }

}

BlueprintsGraphPersistenceRepositoryTestCase.java

package org.jboss.metrics.playground;

import com.tinkerpop.blueprints.IndexableGraph;
import com.tinkerpop.blueprints.impls.neo4j2.Neo4j2Graph;
import org.jboss.metrics.impl.jta.neo4j.ApplicationContext;
import org.jboss.metrics.impl.jta.neo4j.Neo4jExternalTransactionManager;
import org.jboss.metrics.impl.persistence.blueprints.BlueprintsGraphPersistenceRepository;
import org.jboss.metrics.playground.model.Organization;
import org.jboss.metrics.playground.model.Project;
import org.jboss.metrics.spi.persistence.NoSuchEntityException;
import org.jboss.metrics.spi.persistence.PersistenceRepository;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.neo4j.graphdb.GraphDatabaseService;
import org.neo4j.graphdb.factory.GraphDatabaseFactory;
import org.neo4j.graphdb.factory.GraphDatabaseSettings;
import org.neo4j.kernel.GraphDatabaseAPI;

import javax.transaction.HeuristicMixedException;
import javax.transaction.HeuristicRollbackException;
import javax.transaction.NotSupportedException;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import java.io.File;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Tests the Blueprints implementation of the {@link org.jboss.metrics.spi.persistence.PersistenceRepository}
 * in {@link org.jboss.metrics.impl.persistence.blueprints.BlueprintsGraphPersistenceRepository}.  Each test method
 * runs in its own transaction which is rolled back after invocation
 *
 * @author <a href="mailto:alr@jboss.org">ALR</a>
 */
public class BlueprintsGraphPersistenceRepositoryTestCase {

    private static final Logger log = Logger.getLogger(BlueprintsGraphPersistenceRepositoryTestCase.class.getName());

    private static final String FILE_NAME_STORAGE = "build/db/" + BlueprintsGraphPersistenceRepositoryTestCase.class.getSimpleName();

    /**
     * Persistent repo, our hook into storage
     */
    private static PersistenceRepository repo;

    /**
     * Reference to the underlying DB, used in cleanup
     */
    private static IndexableGraph underlyingStorageEngine;

    /**
     * Used to create a {@link javax.transaction.Transaction} for each test case
     */
    private static TransactionManager tm;

    /**
     * Tx used in the scope of the test; will be rolled back on test completion
     */
    private Transaction tx;

    /**
     * Initialize the database and create the persistent repo
     */
    @BeforeClass
    public static void init() {

        // TODO Is this having any effect on the TM?  Check w/ Narayana team; this came from an *old* quickstart
        System.setProperty("ObjectStoreBaseDir", "build");
        // Get a TM; we'll manually invoke it to mock actions the container would take
        final TransactionManager underlyingTm = com.arjuna.ats.jta.TransactionManager
                .transactionManager();

        // Set TM in test context
        ApplicationContext.TRANSACTION_MANAGER.set(underlyingTm);

        // Create the DB
        final GraphDatabaseService graphDatabaseService = new GraphDatabaseFactory().
                newEmbeddedDatabaseBuilder(FILE_NAME_STORAGE).setConfig(
                GraphDatabaseSettings.tx_manager_impl.name(),
                Neo4jExternalTransactionManager.Provider.NAME)
                .newGraphDatabase();


        // Create and init the underlying storage
        underlyingStorageEngine = new Neo4j2Graph(graphDatabaseService);

        tm = ((GraphDatabaseAPI) graphDatabaseService).getDependencyResolver().
                resolveDependency(TransactionManager.class);


        repo = BlueprintsGraphPersistenceRepository.newInstance(underlyingStorageEngine);

        // Neo4j2Graph autostarts a Tx (as will creating a new repo),
        // so make sure we're closing that up
        // if it's hanging around and start fresh
        try {
            tm.commit();
        } catch (final RollbackException | HeuristicRollbackException | HeuristicMixedException | SystemException e) {
            throw new RuntimeException("Error in test setup", e);
        }
    }

    /**
     * Starts a transaction before each test
     *
     * @throws NotSupportedException
     * @throws SystemException
     */
    @Before
    public void startTransaction() throws NotSupportedException, SystemException {
        tm.begin();
        tx = tm.getTransaction();
    }

    /**
     * Rolls back each transaction after each test so its state is never committed and made
     * available outside the context of the test method
     *
     * @throws SystemException
     * @throws HeuristicMixedException
     * @throws HeuristicRollbackException
     */
    @After
    public void rollbackTransaction() throws SystemException,
            HeuristicMixedException, HeuristicRollbackException {
        if (tx != null) {
            tx.setRollbackOnly();
            try {
                tm.commit();
            } catch (final RollbackException expected) {
                // Good
                return;
            }
            throw new IllegalStateException("should have thrown " + RollbackException.class.getName());
        }
    }

    /**
     * Cleanup
     */
    @AfterClass
    public static void shutdown() {
        // Shutdown the engine
        underlyingStorageEngine.shutdown();
        // Remove the datastore so this test has no persistent memory between runs
        final File datastoreFile = new File(FILE_NAME_STORAGE);
        TestUtils.recursiveRemove(datastoreFile);
    }

    /**
     * Ensures that the object/graph mapping and create/retrieval operations
     * of the {@link org.jboss.metrics.impl.persistence.blueprints.BlueprintsGraphPersistenceRepository}
     * are working as contracted
     */
    @Test
    public void createOperationRoundtrip() {

        // Set some properties
        final String orgName = "JBoss";
        final String projectName = "ShrinkWrap";

        // Create some projects/orgs and associate with one another
        final Project shrinkwrap = repo.create(Project.class);
        shrinkwrap.setName(projectName);
        final Organization jboss = repo.create(Organization.class);
        jboss.setName(orgName);
        jboss.addProject(shrinkwrap);
        shrinkwrap.setOrganization(jboss);
        final Object jbossId = jboss.getId();

        // Now find the JBoss org instance we persisted and map it back into a Java object
        final Organization jbossRoundtrip = repo.get(jbossId, Organization.class);

        // Ensure we have the associations intact
        log.log(Level.FINER, "Got roundtrip organization with name: " + jbossRoundtrip.getName());
        Assert.assertEquals("name of org is not as expected", orgName, jbossRoundtrip.getName());
        final Iterable<Project> roundtripProjects = jbossRoundtrip.getProjects();
        int counter = 0;
        for (final Project roundtripProject : roundtripProjects) {
            counter++;
            Assert.assertEquals("name of project is not as expected", projectName, roundtripProject.getName());
            log.log(Level.FINER, "\thas project: " + roundtripProject.getName());
        }
        Assert.assertEquals("should only be one project associated with the org", 1, counter);
    }

    @Test(expected = NoSuchEntityException.class)
    public void noSuchEntityExceptionOnGetWithNonexistentId() {
        repo.get("fakeID", Project.class);
    }

    @Test(expected = NoSuchEntityException.class)
    public void noSuchEntityExceptionOnRemoveWithNonexistentId() {
        repo.remove("fakeId");
    }

    @Test(expected = IllegalArgumentException.class)
    public void illegalArgumentExceptionOnNullArgToCreate() {
        repo.create(null);
    }

    @Test(expected = IllegalArgumentException.class)
    public void illegalArgumentExceptionOnNullArgToRemove() {
        repo.remove(null);
    }

    @Test
    public void remove() {
        // Create a Thing
        final Project project = repo.create(Project.class);

        // Get its ID
        final Object id = project.getId();
        log.finer("Stored ID: " + id);

        // Find it
        final Project roundtrip = repo.get(id, Project.class);

        // Ensure we got it
        Assert.assertNotNull("project should be retrieved and not null", roundtrip);

        // Delete it
        repo.remove(id);

        // Ensure we can't get it
        boolean gotExpectedException = false;
        try {
            repo.get(id, Project.class);
        } catch (final NoSuchEntityException nsee) {
            // Good
            gotExpectedException = true;
        }
        Assert.assertTrue("should not have obtained entity after deletion", gotExpectedException);
    }

}