com.datastax.driver.core.Session Java Examples

@Test(expected = IllegalStateException.class /* the datastax driver complains that page fetching will cause a deadlock */ )
public void paginatedResultInCassandraThreadsThrowsException() {
    int numberOfRecords = 500;
    int numberOfPages = 10;
    Session session = cassandraCQLUnit.getSession();
    PreparedStatement insertStmt = session.prepare(insertInto("app_scale_jobs").values(
            Arrays.asList("job_id", "ref_id"),
            Arrays.asList(bindMarker(), bindMarker())
    ));

    for (int i = 0; i < numberOfRecords; i++) {
        ResultSet resultSet = session.execute(insertStmt.bind("job-" + i, UUID.randomUUID()));
        assertThat(resultSet.wasApplied()).isTrue();
    }

    PreparedStatement loadStmt = session.prepare(select("job_id", "ref_id").from("app_scale_jobs"));
    Observable<ResultSet> results = new CassStoreHelper(session, Schedulers.immediate()).execute(loadStmt.bind()
            // force pagination, and pages to be fetched on demand as the ResultSet is iterated on
            .setFetchSize(numberOfRecords / numberOfPages));

    results.doOnNext(rows -> rows.forEach(row ->
            assertThat(row.getString(0)).startsWith("job-"))
    ).toCompletable().await(1, TimeUnit.MINUTES);
}
 

@Inject
public CassandraMessageDAO(Session session, CassandraTypesProvider typesProvider, BlobStore blobStore,
                           BlobId.Factory blobIdFactory, CassandraConfiguration cassandraConfiguration,
                           CassandraConsistenciesConfiguration consistenciesConfiguration,
                           CassandraMessageId.Factory messageIdFactory) {
    this.cassandraAsyncExecutor = new CassandraAsyncExecutor(session);
    this.consistencyLevel = consistenciesConfiguration.getRegular();
    this.typesProvider = typesProvider;
    this.blobStore = blobStore;
    this.blobIdFactory = blobIdFactory;
    this.configuration = cassandraConfiguration;
    this.messageIdFactory = messageIdFactory;

    this.insert = prepareInsert(session);
    this.delete = prepareDelete(session);
    this.selectMetadata = prepareSelect(session, METADATA);
    this.selectHeaders = prepareSelect(session, HEADERS);
    this.selectFields = prepareSelect(session, FIELDS);
    this.selectBody = prepareSelect(session, BODY);
    this.selectAllMessagesWithAttachment = prepareSelectAllMessagesWithAttachment(session);
    this.cidParser = Cid.parser().relaxed();
}
 

@Test
public void testMapTupleToTupleValue() throws Exception {
    testInSession(new Callback() {
        @Override
        public void call(Session session) throws Exception {
            final DatastaxMapper<DbObjectsWithTupleValue> mapper = DatastaxMapperFactory.newInstance().mapTo(DbObjectsWithTupleValue.class);

            ResultSet rs = session.execute("select id, t from dbobjects_tuple");

            final Iterator<DbObjectsWithTupleValue> iterator = mapper.iterator(rs);

            DbObjectsWithTupleValue next = iterator.next();

            assertEquals(1, next.getId());
            assertEquals("t1", next.getT().getString(0));
            assertEquals(12l, next.getT().getLong(1));
            assertEquals(13, next.getT().getInt(2));

            assertFalse(iterator.hasNext());

        }
    });
}
 

@Test
public void initializeCqlTwice() throws TTransportException, IOException, InterruptedException {
  final InetSocketAddress contactPoint = CassandraTestHelper.getCassandraContactPoint();
  Cluster cluster = Cluster.builder()
      .addContactPointsWithPorts(Collections.singletonList(contactPoint))
      .build();
  Session session = cluster.connect();

  session.execute("DROP KEYSPACE IF EXISTS keyspaceTestCassandraSinkIT");
  Assert.assertNull(session.getCluster().getMetadata().getKeyspace("keyspaceTestCassandraSinkIT"));
  _do();
  Assert.assertNotNull(session.getCluster().getMetadata().getKeyspace("keyspaceTestCassandraSinkIT"));
  Assert.assertNotNull(session.getCluster().getMetadata().getKeyspace("keyspaceTestCassandraSinkIT")
      .getTable("tableTestCassandraSinkIT"));
  _do();
  Assert.assertNotNull(session.getCluster().getMetadata().getKeyspace("keyspaceTestCassandraSinkIT"));
  Assert.assertNotNull(session.getCluster().getMetadata().getKeyspace("keyspaceTestCassandraSinkIT")
      .getTable("tableTestCassandraSinkIT"));
  session.execute("DROP KEYSPACE IF EXISTS keyspaceTestCassandraSinkIT");

  session.close();
  cluster.close();
}
 

@Test
public void testGetBigInteger() throws Exception {
    testInSession(new Callback() {
        @Override
        public void call(Session session) throws Exception {
            final BigIntegerObject nObject = DatastaxMapperFactory.newInstance().mapTo(BigIntegerObject.class).iterator(session.execute("select * from test_number")).next();

            assertEquals(1, nObject.bi.longValue());
            assertEquals(2, nObject.i.longValue());
            assertEquals(3, nObject.vi.longValue());
            assertEquals(3, nObject.dec.longValue());
            assertEquals(3, nObject.f.longValue());
            assertEquals(3, nObject.d.longValue());
        }
    });

}
 

static Session createSession() throws Exception {
    Cluster cluster = Cluster.builder().addContactPoint("127.0.0.1")
            // long read timeout is sometimes needed on slow travis ci machines
            .withSocketOptions(new SocketOptions().setReadTimeoutMillis(30000))
            .withQueryOptions(getQueryOptions())
            .build();
    Session session = cluster.connect();
    session.execute("CREATE KEYSPACE IF NOT EXISTS test WITH REPLICATION ="
            + " { 'class' : 'SimpleStrategy', 'replication_factor' : 1 }");
    session.execute("CREATE TABLE IF NOT EXISTS test.users"
            + " (id int PRIMARY KEY, fname text, lname text)");
    try {
        session.execute("TRUNCATE test.users");
    } catch (NoHostAvailableException e) {
        // sometimes slow, so give it a second chance
        session.execute("TRUNCATE test.users");
    }
    for (int i = 0; i < 10; i++) {
        session.execute("INSERT INTO test.users (id, fname, lname) VALUES (" + i + ", 'f" + i
                + "', 'l" + i + "')");
    }
    return session;
}
 

@Test
public void testGetFloat() throws Exception {
    testInSession(new Callback() {
        @Override
        public void call(Session session) throws Exception {
            final FloatObject nObject = DatastaxMapperFactory.newInstance().mapTo(FloatObject.class).iterator(session.execute("select * from test_number")).next();

            assertEquals(1.0, nObject.bi, 0.001);
            assertEquals(2.0, nObject.i, 0.001);
            assertEquals(3.0, nObject.vi, 0.001);
            assertEquals(3.5, nObject.dec, 0.001);
            assertEquals(3.7, nObject.f, 0.001);
            assertEquals(3.9, nObject.d, 0.001);
        }
    });

}
 

private PreparedStatement prepareInsert(Session session) {
    return session.prepare(insertInto(TABLE_NAME)
        .value(MESSAGE_ID, bindMarker(MESSAGE_ID))
        .value(INTERNAL_DATE, bindMarker(INTERNAL_DATE))
        .value(BODY_START_OCTET, bindMarker(BODY_START_OCTET))
        .value(FULL_CONTENT_OCTETS, bindMarker(FULL_CONTENT_OCTETS))
        .value(BODY_OCTECTS, bindMarker(BODY_OCTECTS))
        .value(BODY_CONTENT, bindMarker(BODY_CONTENT))
        .value(HEADER_CONTENT, bindMarker(HEADER_CONTENT))
        .value(PROPERTIES, bindMarker(PROPERTIES))
        .value(TEXTUAL_LINE_COUNT, bindMarker(TEXTUAL_LINE_COUNT))
        .value(ATTACHMENTS, bindMarker(ATTACHMENTS)));
}
 

public MergeTokenRangesFunction(Comparable maxValue,
                                Comparable minValue,
                                Session session,
                                IPartitioner partitioner,
                                Iterable<Comparable> allRanges) {
    this.maxValue = maxValue;
    this.minValue = minValue;
    this.session = session;
    this.partitioner = partitioner;
    this.allRanges = allRanges;
}
 

static void executeAndUpdateTimer(Session session, Statement stmt, Timer timer) {
	long startTime = System.nanoTime();
	try{
		session.execute(stmt);
	}finally{
		if(timer != null) {
			timer.update(System.nanoTime() - startTime, TimeUnit.NANOSECONDS);
		}
	}
}
 

private PreparedStatement prepareConditionalInsert(Session session) {
    return session.prepare(
        insertInto(CassandraACLTable.TABLE_NAME)
            .value(CassandraACLTable.ID, bindMarker(CassandraACLTable.ID))
            .value(CassandraACLTable.ACL, bindMarker(CassandraACLTable.ACL))
            .value(CassandraACLTable.VERSION, INITIAL_VALUE)
            .ifNotExists());
}
 

public JobsManager(Session session) {
    RxSession rxSession = new RxSessionImpl(session);
    configurationService = new ConfigurationService();
    configurationService.init(rxSession);
    // The hostname is not important here as this class is only used during initialization to persist
    // reoccurring jobs
    scheduler = new SchedulerImpl(rxSession, "localhost");
}
 

private PreparedStatement prepareAdd(Session session) {
    return session.prepare(insertInto(TABLE)
        .value(OWNER, bindMarker(OWNER))
        .value(MESSAGE_ID, bindMarker(MESSAGE_ID))
        .value(BUCKET_NAME, bindMarker(BUCKET_NAME))
        .value(BLOB_ID, bindMarker(BLOB_ID)));
}
 

public CassandraBaseDAO(Session session, ObjectMapper objectMapper, CassandraConfiguration config) {
    this.session = session;
    this.objectMapper = objectMapper;
    this.config = config;

    init();
}
 

public Response getRecords(
    String keySpace, String table, Map<String, Object> filters, List<String> fields) {
  Response response = new Response();
  Session session = connectionManager.getSession(keySpace);
  try {
    Select select;
    if (CollectionUtils.isNotEmpty(fields)) {
      select = QueryBuilder.select((String[]) fields.toArray()).from(keySpace, table);
    } else {
      select = QueryBuilder.select().all().from(keySpace, table);
    }

    if (MapUtils.isNotEmpty(filters)) {
      Select.Where where = select.where();
      for (Map.Entry<String, Object> filter : filters.entrySet()) {
        Object value = filter.getValue();
        if (value instanceof List) {
          where = where.and(QueryBuilder.in(filter.getKey(), ((List) filter.getValue())));
        } else {
          where = where.and(QueryBuilder.eq(filter.getKey(), filter.getValue()));
        }
      }
    }

    ResultSet results = null;
    results = session.execute(select);
    response = CassandraUtil.createResponse(results);
  } catch (Exception e) {
    ProjectLogger.log(Constants.EXCEPTION_MSG_FETCH + table + " : " + e.getMessage(), e);
    throw new ProjectCommonException(
        ResponseCode.SERVER_ERROR.getErrorCode(),
        ResponseCode.SERVER_ERROR.getErrorMessage(),
        ResponseCode.SERVER_ERROR.getResponseCode());
  }
  return response;
}
 

/**
 * Gets a Cassandra connection holder
 *
 * @param read indicates that the application is reading from the Database, this should be true when running AIDA and false when building the repository
 */
public static synchronized CassandraConnectionHolder getCassandraConnectionHolder(boolean read) throws IOException {
  if (cassandraConnectionHolder == null) {
    cassandraConnectionHolder = new CassandraConnectionHolder(read, CassandraConfig.FILE_NAME);

    if (!read) {
      cassandraConnectionHolder.createKeyspaceIfNotExists(CassandraConfig.get(CassandraConfig.FILE_NAME, CassandraConfig.KEYSPACE));
    }

    Session session = cassandraConnectionHolder.getSession();
    session.execute("USE " + CassandraConfig.get(CassandraConfig.FILE_NAME, CassandraConfig.KEYSPACE));
    slogger_.info("Using keyspace " + CassandraConfig.get(CassandraConfig.FILE_NAME, CassandraConfig.KEYSPACE));
  }
  return cassandraConnectionHolder;
}
 

@Test
public void testClassNameMatching() throws Exception {
    testInSession(new Callback() {
        @Override
        public void call(Session session) throws Exception {
            DatastaxCrud<DbObjects, Long> crud =
                    DatastaxMapperFactory.newInstance().crud(DbObjects.class, Long.class).to(session);
            testCrudDbObject(crud, DbObject.newInstance(new DbObjects()));
        }
    });

}
 

double trainAndValidate(int version) throws InstantiationException, IllegalAccessException, ClassNotFoundException {
	final ICollabFilterCassandra cfc;
	String className = "collabfilter.CollabFilterCassandra" + version;
	cfc = (ICollabFilterCassandra) Class.forName(className).newInstance();
	try (Session session = this.cassandraConnector.openSession()) {
		MatrixFactorizationModel model = cfc.train(this.sparkCtx, this.cassandraConnector);
		CassandraJavaRDD<CassandraRow> validationsCassRdd = javaFunctions(this.sparkCtx).cassandraTable(RatingDO.EMPLOYERRATINGS_KEYSPACE, RatingDO.VALIDATION_TABLE);
		JavaRDD<Rating> predictionJavaRdd = cfc.predict(model, validationsCassRdd);
		double rmse = cfc.validate(predictionJavaRdd, validationsCassRdd);
		System.out.println(cfc.resultsReport(predictionJavaRdd, validationsCassRdd, rmse));
		return rmse;
	}

}
 

@SuppressWarnings({"rawtypes", "unchecked"})
public State makeState(Map configuration, IMetricsContext metrics, int partitionIndex, int numPartitions) {

    if (clientFactory == null) {
        clientFactory = new MapConfiguredCqlClientFactory(configuration);
    }

    Session session;
    if(options.keyspace != null) {
        session = clientFactory.getSession(options.keyspace);
    } else {
        session = clientFactory.getSession();
    }

    CassandraCqlMapState state = new CassandraCqlMapState(session, mapper, options, configuration);
    state.registerMetrics(configuration, metrics, options.mapStateMetricName);

    CachedMap cachedMap = new CachedMap(state, options.localCacheSize);

    MapState mapState;
    if (stateType == StateType.NON_TRANSACTIONAL) {
        mapState = NonTransactionalMap.build(cachedMap);
    } else if (stateType == StateType.OPAQUE) {
        mapState = OpaqueMap.build(cachedMap);
    } else if (stateType == StateType.TRANSACTIONAL) {
        mapState = TransactionalMap.build(cachedMap);
    } else {
        throw new RuntimeException("Unknown state type: " + stateType);
    }

    return new SnapshottableMap(mapState, new Values(options.globalKey));
}
 

@Test
void initializeShouldExecuteCreateStatementAndReturnFullWhenTypeDoesNotExist() {
    KeyspaceMetadata keyspace = mock(KeyspaceMetadata.class);
    when(keyspace.getUserType(NAME)).thenReturn(null);
    Session session = mock(Session.class);

    assertThat(TYPE.initialize(keyspace, session))
            .isEqualByComparingTo(FULL);

    verify(keyspace).getUserType(NAME);
    verify(session).execute(STATEMENT);
}
 

public static void init(Cluster cluster) {
    try (Session systemSession = cluster.connect()) {
        String createKeyspace = String.format("CREATE KEYSPACE IF NOT EXISTS %s WITH replication = " +
                "{'class':'SimpleStrategy','replication_factor':'1'};", TEST_KEYSPACE);
        systemSession.execute(createKeyspace);
        String createTable = String.format("CREATE TABLE %s.%s (id text, value text, PRIMARY KEY(id))", TEST_KEYSPACE, TEST_TABLE);
        systemSession.execute(createTable);
    }
}
 

@Override
protected void configure() {
    bind(CassandraConfiguration.class).to(SystemPropertiesCassandraConfiguration.class);
    bind(Cluster.class).toProvider(CassandraClusterProvider.class).asEagerSingleton();
    bind(Session.class).toProvider(CassandraSessionProvider.class);

    bind(MetadataDAO.class).to(CassandraMetadataDAO.class);
    bind(ExecutionDAO.class).to(CassandraExecutionDAO.class);
    bind(EventHandlerDAO.class).to(CassandraEventHandlerDAO.class);
}
 

protected static PreparedStatement listByPlaceStatement(Session session) {
   CassandraQueryBuilder queryBuilder =
         CassandraQueryBuilder
            .select(RuleEnvironmentTable.NAME)
            .addWhereColumnEquals(RuleEnvironmentTable.Column.PLACE_ID.columnName());
   return RuleEnvironmentTable.addAllColumns(queryBuilder).prepare(session);
}
 

@Test
void initializeShouldReturnAlreadyDoneWhenTypeExists() {
    KeyspaceMetadata keyspace = mock(KeyspaceMetadata.class);
    when(keyspace.getUserType(NAME)).thenReturn(mock(UserType.class));
    Session session = mock(Session.class);

    assertThat(TYPE.initialize(keyspace, session))
            .isEqualByComparingTo(ALREADY_DONE);

    verify(keyspace).getUserType(NAME);
    verify(session, never()).execute(STATEMENT);
}
 

@Override
public Mapper apply(Session session) {
  if (mappingManager == null) {
    this.mappingManager = new MappingManager(session);
  }

  return new DefaultObjectMapper<T>(mappingManager.mapper(entity));
}
 

@Override
public PreparedStatement prepare(Session session, String query, ConsistencyLevel consistency, RetryPolicy retryPolicy) {
   PreparedStatement ps = session.prepare(query);
   ps.setConsistencyLevel(consistency);
   if(retryPolicy != null) {
      ps.setRetryPolicy(retryPolicy);
   }
   return ps;
}
 

private Response getRecordByIdentifier(
    String keyspaceName, String tableName, Object key, List<String> fields) {
  long startTime = System.currentTimeMillis();
  ProjectLogger.log(
      "Cassandra Service getRecordBy key method started at ==" + startTime, LoggerEnum.INFO);
  Response response = new Response();
  try {
    Session session = connectionManager.getSession(keyspaceName);
    Builder selectBuilder;
    if (CollectionUtils.isNotEmpty(fields)) {
      selectBuilder = QueryBuilder.select(fields.toArray(new String[fields.size()]));
    } else {
      selectBuilder = QueryBuilder.select().all();
    }
    Select selectQuery = selectBuilder.from(keyspaceName, tableName);
    Where selectWhere = selectQuery.where();
    if (key instanceof String) {
      selectWhere.and(eq(Constants.IDENTIFIER, key));
    } else if (key instanceof Map) {
      Map<String, Object> compositeKey = (Map<String, Object>) key;
      compositeKey
          .entrySet()
          .stream()
          .forEach(
              x -> {
                CassandraUtil.createQuery(x.getKey(), x.getValue(), selectWhere);
              });
    }
    ResultSet results = session.execute(selectWhere);
    response = CassandraUtil.createResponse(results);
  } catch (Exception e) {
    ProjectLogger.log(Constants.EXCEPTION_MSG_FETCH + tableName + " : " + e.getMessage(), e);
    throw new ProjectCommonException(
        ResponseCode.SERVER_ERROR.getErrorCode(),
        ResponseCode.SERVER_ERROR.getErrorMessage(),
        ResponseCode.SERVER_ERROR.getResponseCode());
  }
  logQueryElapseTime("getRecordByIdentifier", startTime);
  return response;
}
 

private PreparedStatement prepareInsert(Session session) {
    return session.prepare(
        insertInto(TABLE_NAME)
            .value(ATTACHMENT_ID_AS_UUID, bindMarker(ATTACHMENT_ID_AS_UUID))
            .value(ATTACHMENT_ID, bindMarker(ATTACHMENT_ID))
            .value(MESSAGE_ID, bindMarker(MESSAGE_ID)));
}
 

public DefaultCommandContext(CommandLine commandLine,
                             Session defaultSession,
                             Function<String, Session> sourceSessionFactory,
                             Function<String, Session> targetSessionFactory) {
    this.commandLine = commandLine;
    this.defaultSession = defaultSession;
    this.sourceSession = commandLine.hasOption('s')
            ? Optional.of(new CassSession(sourceSessionFactory, commandLine.getOptionValue('s')))
            : Optional.empty();
    this.targetSession = commandLine.hasOption('t')
            ? Optional.of(new CassSession(targetSessionFactory, commandLine.getOptionValue('t')))
            : Optional.empty();
}
 

@Test
void createSessionShouldKeepTheSetSchemaVersionWhenTypesAndTablesHaveNotChanged() {
    Session session = testee.get();
    assertThat(versionManager(session).computeVersion().block())
            .isEqualTo(MAX_VERSION);

    new CassandraTableManager(MODULE, session).clearAllTables();
    versionManagerDAO(session).updateVersion(MIN_VERSION);
    assertThat(versionManager(session).computeVersion().block())
            .isEqualTo(MIN_VERSION);

    assertThat(versionManager(testee.get()).computeVersion().block())
            .isEqualTo(MIN_VERSION);
}