pro.foundev.cassandra.commons.test.CassandraTestDB Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of cassandra-commons-test Show documentation
Session reuse, database cleanup, and persistence methods for use with integration tests in Apache Cassandra
The newest version!
/*
 * Copyright 2015 Ryan Svihla
 *
 * 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 pro.foundev.cassandra.commons.test;

import com.datastax.driver.core.*;
import pro.foundev.cassandra.commons.core.CassandraConfiguration;
import pro.foundev.cassandra.commons.core.CassandraSessionFactory;
import pro.foundev.cassandra.commons.core.config.ConfigFinder;
import pro.foundev.cassandra.commons.core.config.ConfigFinderImpl;

import java.io.Closeable;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

/**
 * Useful test class that provides very basic CRUD for testing support. It handles Session and Cluster setup and shutdown
 * and will when used with the SpringJUnit4ClassRunner by implementing DisposableBean.
 *
 * It also cleans up after itself by truncating tables automatically after a test. If you save to a table it will
 * truncate a table or if you explicitly call markForCleanup(keyspace, tableName). This feature is implemented by wiring up a
 * custom Test Listener
 */
public class CassandraTestDB implements Closeable {
    private final CassandraSessionFactory sessionFactory;
    protected final Session session;
    private PreparedStatement findTableStatement;
    private final Set tablesMarkedForCleanup = new HashSet<>();
    private Set keySpacesMarkedForCleanup = new HashSet<>();

    private synchronized Set getTablesMarkedForCleanup(){
       return tablesMarkedForCleanup;
    }
    private synchronized void resetTablesMarkedForCleanup() {
        tablesMarkedForCleanup.clear();
    }
    protected synchronized void markForCleanUp(FullTableName fullTableName){
        this.tablesMarkedForCleanup.add(fullTableName);
    }

    /**
     *
     * @param sessionFactory owns the session object and responsible for setup
     */
    public CassandraTestDB(CassandraSessionFactory sessionFactory){
        this.sessionFactory = sessionFactory;
        this.session = sessionFactory.getSession();
        findTableStatement = session.prepare("SELECT * FROM system.schema_columnfamilies " +
                "where keyspace_name=? and columnfamily_name=?");
    }

    /**
     * Factory method for CassandraTestDB using default ConfigFinderImpl
     * @return CassandraTestDB
     * @throws IOException
     */
    public static CassandraTestDB create() throws IOException {
        ConfigFinder configFinder = new ConfigFinderImpl();
        return create(configFinder);
    }

    /**
     * Factory method for CassandraTestDB with optional ConfigFinder override.
     * @return CassandraTestDB
     * @throws IOException when config file not found.
     */
    public static CassandraTestDB create(ConfigFinder configFinder) throws IOException {
        CassandraConfiguration configuration = CassandraConfiguration.parse(configFinder.find());
        CassandraSessionFactory sessionFactory = new CassandraSessionFactory(configuration);
        return new CassandraTestDB(sessionFactory);
    }

    /**
     * Cleanup strategy to that truncates known tables.
     * Depends on save or markForCleanup methods making the table known.
     * Today there is a precondition on the Keyspace and Table name being declared in the Mapping.
     * This may change later.
     */
    public void cleanUp() {
        getTablesMarkedForCleanup().forEach(t -> {
            Row row = session.execute(findTableStatement.bind(t.getKeyspaceName(), t.getTableName())).one();
            if (row != null) {
                session.execute("TRUNCATE " + t);
            }
        });
        resetTablesMarkedForCleanup();
        keySpacesMarkedForCleanup.forEach(k->session.execute("DROP KEYSPACE IF EXISTS " + k));
        keySpacesMarkedForCleanup.clear();
    }

    /**
     * Mark table for cleanup.
     * @param keyspace keyspace where the table lives
     * @param tableName table to truncate
     */
    public void markForCleanUp(String keyspace, String tableName){
        FullTableName fullTableName = new FullTableName();
        fullTableName.setKeyspaceName(keyspace);
        fullTableName.setTableName(tableName);
        markForCleanUp(fullTableName);
    }

    /**
     * assumes keyspace is already wired up
     * @param table table to mark for cleanup
     */
    public void markForCleanUp(String table) {
        String keyspace = session.getLoggedKeyspace();
        if(keyspace==null){
            //FIXME: better exceptions
            throw new RuntimeException("keyspace is null");
        }
       markForCleanUp(keyspace, table);
    }

    /**
     * Just returns the internal session object. Any queries that will
     * be executed against this session are not tracked by the listener.
     * @return
     */
    public Session getSession(){
        return session;
    }

    /**
     * creates a simple keyspace for testing purposes with SimpleStrategy and replication factor of 1
     * will blow it away on test teardown
     * @param keyspaceName
     */
    public void createTestKeyspace(String keyspaceName){
        keySpacesMarkedForCleanup.add(keyspaceName);
        session.execute("CREATE KEYSPACE " + keyspaceName+ " with replication " +
        "= {'class': 'SimpleStrategy', 'replication_factor':1 }");
    }

    /**
     * simple utility check to see if table exists, does mark table for cleanup if it does
     * @param keyspaceName
     * @param tableName
     * @return true if table is in system false if not
     */
    public boolean tableExists(String keyspaceName, String tableName){
        boolean exists = null != session.execute("SELECT * FROM system.schema_columnfamilies where " +
                "keyspace_name='"+keyspaceName+"' AND columnfamily_name='"+tableName+"'")
                .one();
        if(exists){
            FullTableName fullTableName = new FullTableName();
            fullTableName.setKeyspaceName(keyspaceName);
            fullTableName.setTableName(tableName);
        }
        return exists;
    }

    /**
     * drops the keyspace if it doesn't exist. great for testing methods that have manual keyspace resources
     * @param keyspaceName
     */
    public void deleteTestKeyspaceIfExists(String keyspaceName) {
        session.execute("DROP KEYSPACE IF EXISTS " + keyspaceName);
    }

    /**
     * Shuts down DataStax Java driver Session and Cluster objects safely
     */
    @Override
    public void close() throws IOException {
        sessionFactory.close();
    }

    /**
     * session.execute wrapper
     * @param statement statement to execute
     * @return
     */
    public ResultSet execute(Statement statement) {
       return session.execute(statement);
    }

    /**
     * count on a given table. It will page on versions of C* 2.0.0 and greater. Otherwise it'll pull the whole
     * results set into ram and offline the coordinator.
     * @param table table to get count on
     * @return total count of records in that table
     */
    public Long count(String table) {
        return session.execute("SELECT COUNT(*) FROM " + table)
                .one()
                .getLong(0);
    }


}