def do_this():
    for keyspace in keyspaces:
        factory = ManagedCassandraClientFactory(keyspace)
        reactor.connectTCP("localhost", 9160, factory)
        client = CassandraClient(factory)
        k =  yield client.describe_keyspace(keyspace)
        for cf in [c.name for c in k.cf_defs]:
            n = 0
            db_rows = yield client.get_range_slices(column_family=cf, count=10000000)
            with gzip.GzipFile("%s.%s.csv.gz" % (keyspace, cf), "w") as f:
                out = csv.DictWriter(f, fieldnames=["keyspace", "cf", "key", "col", "val"])
                out.writeheader()
                for row in db_rows:
                    for col in row.columns:
                        out.writerow({"keyspace": keyspace,
                                      "cf": cf,
                                      "key": row.key.encode("string_escape"),
                                      "col": col.column.name,
                                      "val": col.column.value.encode("string_escape")})
                    n += 1
                print "Successfully backed up %d rows from %s.%s" % (n, keyspace, cf)
    reactor.stop()
class CassandraClientTest(unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self):
        self.cmanager = ManagedCassandraClientFactory(keyspace='system')
        self.client = CassandraClient(self.cmanager)
        for i in xrange(CONNS):
            reactor.connectTCP(HOST, PORT, self.cmanager)
        yield self.cmanager.deferred

        remote_ver = yield self.client.describe_version()
        self.version = getAPIVersion(remote_ver)

        self.my_keyspace = KsDef(
            name=KEYSPACE,
            strategy_class='org.apache.cassandra.locator.SimpleStrategy',
            strategy_options={'replication_factor': '1'},
            cf_defs=[
                CfDef(
                    keyspace=KEYSPACE,
                    name=CF,
                    column_type='Standard'
                ),
                CfDef(
                    keyspace=KEYSPACE,
                    name=SCF,
                    column_type='Super'
                ),
                CfDef(
                    keyspace=KEYSPACE,
                    name=IDX_CF,
                    column_type='Standard',
                    comparator_type='org.apache.cassandra.db.marshal.UTF8Type',
                    column_metadata=[
                        ColumnDef(
                            name='col1',
                            validation_class='org.apache.cassandra.db.marshal.UTF8Type',
                            index_type=IndexType.KEYS,
                            index_name='idxCol1')
                    ],
                    default_validation_class='org.apache.cassandra.db.marshal.BytesType'
                ),
            ]
        )
        if self.version == CASSANDRA_08_VERSION:
            self.my_keyspace.cf_defs.extend([
                CfDef(
                    keyspace=KEYSPACE,
                    name=COUNTER_CF,
                    column_type='Standard',
                    default_validation_class='org.apache.cassandra.db.marshal.CounterColumnType'
                ),
                CfDef(
                    keyspace=KEYSPACE,
                    name=SUPERCOUNTER_CF,
                    column_type='Super',
                    default_validation_class='org.apache.cassandra.db.marshal.CounterColumnType'
                ),
            ])

        yield self.client.system_add_keyspace(self.my_keyspace)
        yield self.client.set_keyspace(KEYSPACE)

    @defer.inlineCallbacks
    def tearDown(self):
        yield self.client.system_drop_keyspace(self.my_keyspace.name)
        self.cmanager.shutdown()
        for c in reactor.getDelayedCalls():
            c.cancel()
        reactor.removeAll()

    @defer.inlineCallbacks
    def test_insert_get(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        yield self.client.insert('test', SCF, 'superval', column=COLUMN, super_column=SCOLUMN)
        yield self.client.insert('test2', SCF, 'superval2', column=COLUMN,
                                 super_column=SCOLUMN)
        res = yield self.client.get('test', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        res = yield self.client.get('test2', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval2')
        res = yield self.client.get('test', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval')
        res = yield self.client.get('test2', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval2')

    @defer.inlineCallbacks
    def test_batch_insert_get_slice_and_count(self):
        yield self.client.batch_insert('test', CF,
                                       {COLUMN: 'test', COLUMN2: 'test2'})
        yield self.client.batch_insert('test', SCF,
                               {SCOLUMN: {COLUMN: 'test', COLUMN2: 'test2'}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_count('test', CF)
        self.assertEqual(res, 2)

    @defer.inlineCallbacks
    def test_batch_mutate_and_remove(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } }, 'test2': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } } })
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_remove({CF: ['test', 'test2']}, names=['test', 'test2'])
        yield self.client.batch_remove({SCF: ['test', 'test2']}, names=['test', 'test2'], supercolumn=SCOLUMN)

    @defer.inlineCallbacks
    def test_batch_mutate_with_deletion(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_mutate({'test': {CF: {COLUMN: None, COLUMN2: 'test3'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(len(res), 1)
        self.assertEqual(res[0].column.value, 'test3')

    @defer.inlineCallbacks
    def test_multiget_slice_remove(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        res = yield self.client.multiget_slice(['test', 'test2'], CF)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test'][1].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        yield self.client.remove('test', CF, column=COLUMN)
        yield self.client.remove('test2', CF, column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(len(res['test']), 0)
        self.assertEqual(len(res['test2']), 0)

    @defer.inlineCallbacks
    def test_range_slices(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        ks = yield self.client.get_range_slices(CF, start='', finish='')
        keys = [k.key for k in ks]
        for key in ['test', 'test2']:
            self.assertIn(key, keys)

    @defer.inlineCallbacks
    def test_indexed_slices(self):
        yield self.client.insert('test1', IDX_CF, 'one', column='col1')
        yield self.client.insert('test2', IDX_CF, 'two', column='col1')
        yield self.client.insert('test3', IDX_CF, 'three', column='col1')
        expressions = [IndexExpression('col1', IndexOperator.EQ, 'two')]
        res = yield self.client.get_indexed_slices(IDX_CF, expressions, start_key='')
        self.assertEquals(res[0].columns[0].column.value,'two')

    @defer.inlineCallbacks
    def test_counter_add(self):
        if self.version != CASSANDRA_08_VERSION:
            raise unittest.SkipTest('Counters are not supported in 0.7')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 2)

        # test super column counters
        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 2)

    @defer.inlineCallbacks
    def test_counter_remove(self):
        if self.version != CASSANDRA_08_VERSION:
            raise unittest.SkipTest('Counters are not supported in 0.7')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test', COUNTER_CF, column='col')
        yield self.assertFailure(self.client.get('test', COUNTER_CF, column='col'),
                                 NotFoundException)

        # test super column counters
        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test', SUPERCOUNTER_CF,
                                         column='col', super_column='scol')
        yield self.assertFailure(self.client.get('test', SUPERCOUNTER_CF,
                                                 column='col', super_column='scol'),
                                 NotFoundException)

    @defer.inlineCallbacks
    def test_cql(self):
        if self.version != CASSANDRA_08_VERSION:
            raise unittest.SkipTest('CQL is not supported in 0.7')

        yield self.client.insert('test', CF, 'testval', column='col1')
        res = yield self.client.get('test', CF, column='col1')
        self.assertEquals(res.column.value, 'testval')

        query = 'SELECT * from %s where KEY = %s' % (CF, 'test'.encode('hex'))
        uncompressed_result = yield self.client.execute_cql_query(query, Compression.NONE)
        self.assertEquals(uncompressed_result.rows[0].columns[0].name, 'col1')
        self.assertEquals(uncompressed_result.rows[0].columns[0].value, 'testval')

        compressed_query = zlib.compress(query)
        compressed_result = yield self.client.execute_cql_query(compressed_query, Compression.GZIP)
        self.assertEquals(uncompressed_result, compressed_result)

    def sleep(self, secs):
        d = defer.Deferred()
        reactor.callLater(secs, d.callback, None)
        return d

    @defer.inlineCallbacks
    def test_ttls(self):
        yield self.client.insert('test_ttls', CF, 'testval', column=COLUMN, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_insert('test_ttls', CF, {COLUMN:'testval'}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_mutate({'test_ttls': {CF: {COLUMN: 'testval'}}}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

    def compare_keyspaces(self, ks1, ks2):
        self.assertEqual(ks1.name, ks2.name)
        self.assertEqual(ks1.strategy_class, ks2.strategy_class)
        self.assertEqual(ks1.cf_defs, ks2.cf_defs)

        def get_rf(ksdef):
            rf = ksdef.replication_factor
            if ksdef.strategy_options and \
               'replication_factor' in ksdef.strategy_options:
                rf = int(ksdef.strategy_options['replication_factor'])
            return rf

        def strat_opts_no_rf(ksdef):
            if not ksdef.strategy_options:
                return {}
            opts = ksdef.strategy_options.copy()
            if 'replication_factor' in ksdef.strategy_options:
                del opts['replication_factor']
            return opts

        self.assertEqual(get_rf(ks1), get_rf(ks2))
        self.assertEqual(strat_opts_no_rf(ks1), strat_opts_no_rf(ks2))

    @defer.inlineCallbacks
    def test_keyspace_manipulation(self):
        try:
            yield self.client.system_drop_keyspace(T_KEYSPACE)
        except InvalidRequestException:
            pass
        ksdef = KsDef(name=T_KEYSPACE, strategy_class='org.apache.cassandra.locator.SimpleStrategy', strategy_options={'replication_factor': '1'}, cf_defs=[])
        yield self.client.system_add_keyspace(ksdef)
        ks2 = yield self.client.describe_keyspace(T_KEYSPACE)
        self.compare_keyspaces(ksdef, ks2)

        if DO_SYSTEM_RENAMING:
            newname = T_KEYSPACE + '2'
            yield self.client.system_rename_keyspace(T_KEYSPACE, newname)
            ks2 = yield self.client.describe_keyspace(newname)
            ksdef.name = newname
            self.compare_keyspaces(ksdef, ks2)
        yield self.client.system_drop_keyspace(ksdef.name)
        yield self.assertFailure(self.client.describe_keyspace(T_KEYSPACE), NotFoundException)
        if DO_SYSTEM_RENAMING:
            yield self.assertFailure(self.client.describe_keyspace(ksdef.name), NotFoundException)

    @defer.inlineCallbacks
    def test_column_family_manipulation(self):
        cfdef = CfDef(KEYSPACE, T_CF,
            column_type='Standard',
            comparator_type='org.apache.cassandra.db.marshal.BytesType',
            comment='foo',
            row_cache_size=0.0,
            key_cache_size=200000.0,
            read_repair_chance=1.0,
            column_metadata=[],
            gc_grace_seconds=86400,
            default_validation_class='org.apache.cassandra.db.marshal.BytesType',
            key_validation_class='org.apache.cassandra.db.marshal.BytesType',
            min_compaction_threshold=5,
            max_compaction_threshold=31,
            row_cache_save_period_in_seconds=0,
            key_cache_save_period_in_seconds=3600,
            memtable_flush_after_mins=60,
            memtable_throughput_in_mb=249,
            memtable_operations_in_millions=1.1671875,
            replicate_on_write=False,
            merge_shards_chance=0.10000000000000001,
            row_cache_provider=None,
            key_alias=None,
        )
        post_07_fields = ['replicate_on_write', 'merge_shards_chance',
                          'key_validation_class', 'row_cache_provider', 'key_alias']

        yield self.client.system_add_column_family(cfdef)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        cfdef2 = [c for c in ksdef.cf_defs if c.name == T_CF][0]

        for field in post_07_fields:
            # Most of these are ignored in 0.7, so we can't reliably compare them
            setattr(cfdef, field, None)
            setattr(cfdef2, field, None)

        # we don't know the id ahead of time. copy the new one so the equality
        # comparison won't fail
        cfdef.id = cfdef2.id
        self.assertEqual(cfdef, cfdef2)
        if DO_SYSTEM_RENAMING:
            newname = T_CF + '2'
            yield self.client.system_rename_column_family(T_CF, newname)
            ksdef = yield self.client.describe_keyspace(KEYSPACE)
            cfdef2 = [c for c in ksdef.cf_defs if c.name == newname][0]
            self.assertNotIn(T_CF, [c.name for c in ksdef.cf_defs])
            cfdef.name = newname
            self.assertEqual(cfdef, cfdef2)
        yield self.client.system_drop_column_family(cfdef.name)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        self.assertNotIn(cfdef.name, [c.name for c in ksdef.cf_defs])

    @defer.inlineCallbacks
    def test_describes(self):
        name = yield self.client.describe_cluster_name()
        self.assertIsInstance(name, str)
        self.assertNotEqual(name, '')
        partitioner = yield self.client.describe_partitioner()
        self.assert_(partitioner.startswith('org.apache.cassandra.'),
                     msg='partitioner is %r' % partitioner)
        snitch = yield self.client.describe_snitch()
        self.assert_(snitch.startswith('org.apache.cassandra.'),
                     msg='snitch is %r' % snitch)
        version = yield self.client.describe_version()
        self.assertIsInstance(version, str)
        self.assertIn('.', version)
        schemavers = yield self.client.describe_schema_versions()
        self.assertIsInstance(schemavers, dict)
        self.assertNotEqual(schemavers, {})
        ring = yield self.client.describe_ring(KEYSPACE)
        self.assertIsInstance(ring, list)
        self.assertNotEqual(ring, [])
        for r in ring:
            self.assertIsInstance(r.start_token, str)
            self.assertIsInstance(r.end_token, str)
            self.assertIsInstance(r.endpoints, list)
            self.assertNotEqual(r.endpoints, [])
            for ep in r.endpoints:
                self.assertIsInstance(ep, str)

    @defer.inlineCallbacks
    def test_errback(self):
        yield self.client.remove('poiqwe', CF)
        try:
            yield self.client.get('poiqwe', CF, column='foo')
        except Exception, e:
            pass
Example #3
0
class CassandraClientTest(unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self):
        self.cmanager = ManagedCassandraClientFactory(keyspace='system')
        self.client = CassandraClient(self.cmanager)
        for i in xrange(CONNS):
            reactor.connectTCP(HOST, PORT, self.cmanager)
        yield self.cmanager.deferred

        remote_ver = yield self.client.describe_version()
        self.version = thrift_api_ver_to_cassandra_ver(remote_ver)

        self.my_keyspace = KsDef(
            name=KEYSPACE,
            strategy_class='org.apache.cassandra.locator.SimpleStrategy',
            strategy_options={'replication_factor': '1'},
            cf_defs=[
                CfDef(keyspace=KEYSPACE, name=CF, column_type='Standard'),
                CfDef(keyspace=KEYSPACE, name=SCF, column_type='Super'),
                CfDef(
                    keyspace=KEYSPACE,
                    name=IDX_CF,
                    column_type='Standard',
                    comparator_type='org.apache.cassandra.db.marshal.UTF8Type',
                    column_metadata=[
                        ColumnDef(name='col1',
                                  validation_class=
                                  'org.apache.cassandra.db.marshal.UTF8Type',
                                  index_type=IndexType.KEYS,
                                  index_name='idxCol1')
                    ],
                    default_validation_class=
                    'org.apache.cassandra.db.marshal.BytesType'),
            ])
        if self.version == CASSANDRA_08_VERSION:
            self.my_keyspace.cf_defs.extend([
                CfDef(keyspace=KEYSPACE,
                      name=COUNTER_CF,
                      column_type='Standard',
                      default_validation_class=
                      'org.apache.cassandra.db.marshal.CounterColumnType'),
                CfDef(keyspace=KEYSPACE,
                      name=SUPERCOUNTER_CF,
                      column_type='Super',
                      default_validation_class=
                      'org.apache.cassandra.db.marshal.CounterColumnType'),
            ])

        yield self.client.system_add_keyspace(self.my_keyspace)
        yield self.client.set_keyspace(KEYSPACE)

    @defer.inlineCallbacks
    def tearDown(self):
        yield self.client.system_drop_keyspace(self.my_keyspace.name)
        self.cmanager.shutdown()
        for c in reactor.getDelayedCalls():
            c.cancel()
        reactor.removeAll()

    @defer.inlineCallbacks
    def test_insert_get(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        yield self.client.insert('test',
                                 SCF,
                                 'superval',
                                 column=COLUMN,
                                 super_column=SCOLUMN)
        yield self.client.insert('test2',
                                 SCF,
                                 'superval2',
                                 column=COLUMN,
                                 super_column=SCOLUMN)
        res = yield self.client.get('test', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        res = yield self.client.get('test2', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval2')
        res = yield self.client.get('test',
                                    SCF,
                                    column=COLUMN,
                                    super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval')
        res = yield self.client.get('test2',
                                    SCF,
                                    column=COLUMN,
                                    super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval2')

    @defer.inlineCallbacks
    def test_batch_insert_get_slice_and_count(self):
        yield self.client.batch_insert('test', CF, {
            COLUMN: 'test',
            COLUMN2: 'test2'
        })
        yield self.client.batch_insert(
            'test', SCF, {SCOLUMN: {
                COLUMN: 'test',
                COLUMN2: 'test2'
            }})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test',
                                          SCF,
                                          names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_count('test', CF)
        self.assertEqual(res, 2)

    @defer.inlineCallbacks
    def test_batch_mutate_and_remove(self):
        yield self.client.batch_mutate({
            'test': {
                CF: {
                    COLUMN: 'test',
                    COLUMN2: 'test2'
                },
                SCF: {
                    SCOLUMN: {
                        COLUMN: 'test',
                        COLUMN2: 'test2'
                    }
                }
            },
            'test2': {
                CF: {
                    COLUMN: 'test',
                    COLUMN2: 'test2'
                },
                SCF: {
                    SCOLUMN: {
                        COLUMN: 'test',
                        COLUMN2: 'test2'
                    }
                }
            }
        })
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test',
                                          SCF,
                                          names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2',
                                          SCF,
                                          names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_remove({CF: ['test', 'test2']},
                                       names=['test', 'test2'])
        yield self.client.batch_remove({SCF: ['test', 'test2']},
                                       names=['test', 'test2'],
                                       supercolumn=SCOLUMN)

    @defer.inlineCallbacks
    def test_batch_mutate_with_deletion(self):
        yield self.client.batch_mutate(
            {'test': {
                CF: {
                    COLUMN: 'test',
                    COLUMN2: 'test2'
                }
            }})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_mutate(
            {'test': {
                CF: {
                    COLUMN: None,
                    COLUMN2: 'test3'
                }
            }})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(len(res), 1)
        self.assertEqual(res[0].column.value, 'test3')

    @defer.inlineCallbacks
    def test_multiget_slice_remove(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        res = yield self.client.multiget_slice(['test', 'test2'], CF)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test'][1].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        yield self.client.remove('test', CF, column=COLUMN)
        yield self.client.remove('test2', CF, column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(len(res['test']), 0)
        self.assertEqual(len(res['test2']), 0)

    @defer.inlineCallbacks
    def test_range_slices(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        ks = yield self.client.get_range_slices(CF, start='', finish='')
        keys = [k.key for k in ks]
        for key in ['test', 'test2']:
            self.assertIn(key, keys)

    @defer.inlineCallbacks
    def test_indexed_slices(self):
        yield self.client.insert('test1', IDX_CF, 'one', column='col1')
        yield self.client.insert('test2', IDX_CF, 'two', column='col1')
        yield self.client.insert('test3', IDX_CF, 'three', column='col1')
        expressions = [IndexExpression('col1', IndexOperator.EQ, 'two')]
        res = yield self.client.get_indexed_slices(IDX_CF,
                                                   expressions,
                                                   start_key='')
        self.assertEquals(res[0].columns[0].column.value, 'two')

    @defer.inlineCallbacks
    def test_counter_add(self):
        if self.version != CASSANDRA_08_VERSION:
            raise unittest.SkipTest('Counters are not supported in 0.7')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 2)

        # test super column counters
        yield self.client.add('test',
                              SUPERCOUNTER_CF,
                              1,
                              column='col',
                              super_column='scol')
        res = yield self.client.get('test',
                                    SUPERCOUNTER_CF,
                                    column='col',
                                    super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test',
                              SUPERCOUNTER_CF,
                              1,
                              column='col',
                              super_column='scol')
        res = yield self.client.get('test',
                                    SUPERCOUNTER_CF,
                                    column='col',
                                    super_column='scol')
        self.assertEquals(res.counter_column.value, 2)

    @defer.inlineCallbacks
    def test_counter_remove(self):
        if self.version != CASSANDRA_08_VERSION:
            raise unittest.SkipTest('Counters are not supported in 0.7')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test', COUNTER_CF, column='col')
        yield self.assertFailure(
            self.client.get('test', COUNTER_CF, column='col'),
            NotFoundException)

        # test super column counters
        yield self.client.add('test',
                              SUPERCOUNTER_CF,
                              1,
                              column='col',
                              super_column='scol')
        res = yield self.client.get('test',
                                    SUPERCOUNTER_CF,
                                    column='col',
                                    super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test',
                                         SUPERCOUNTER_CF,
                                         column='col',
                                         super_column='scol')
        yield self.assertFailure(
            self.client.get('test',
                            SUPERCOUNTER_CF,
                            column='col',
                            super_column='scol'), NotFoundException)

    def sleep(self, secs):
        d = defer.Deferred()
        reactor.callLater(secs, d.callback, None)
        return d

    @defer.inlineCallbacks
    def test_ttls(self):
        yield self.client.insert('test_ttls',
                                 CF,
                                 'testval',
                                 column=COLUMN,
                                 ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(
            self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_insert('test_ttls',
                                       CF, {COLUMN: 'testval'},
                                       ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(
            self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_mutate(
            {'test_ttls': {
                CF: {
                    COLUMN: 'testval'
                }
            }}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(
            self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

    def compare_keyspaces(self, ks1, ks2):
        self.assertEqual(ks1.name, ks2.name)
        self.assertEqual(ks1.strategy_class, ks2.strategy_class)
        self.assertEqual(ks1.cf_defs, ks2.cf_defs)

        def get_rf(ksdef):
            rf = ksdef.replication_factor
            if ksdef.strategy_options and \
               'replication_factor' in ksdef.strategy_options:
                rf = int(ksdef.strategy_options['replication_factor'])
            return rf

        def strat_opts_no_rf(ksdef):
            if not ksdef.strategy_options:
                return {}
            opts = ksdef.strategy_options.copy()
            if 'replication_factor' in ksdef.strategy_options:
                del opts['replication_factor']
            return opts

        self.assertEqual(get_rf(ks1), get_rf(ks2))
        self.assertEqual(strat_opts_no_rf(ks1), strat_opts_no_rf(ks2))

    @defer.inlineCallbacks
    def test_keyspace_manipulation(self):
        try:
            yield self.client.system_drop_keyspace(T_KEYSPACE)
        except InvalidRequestException:
            pass
        ksdef = KsDef(
            name=T_KEYSPACE,
            strategy_class='org.apache.cassandra.locator.SimpleStrategy',
            strategy_options={'replication_factor': '1'},
            cf_defs=[])
        yield self.client.system_add_keyspace(ksdef)
        ks2 = yield self.client.describe_keyspace(T_KEYSPACE)
        self.compare_keyspaces(ksdef, ks2)

        if DO_SYSTEM_RENAMING:
            newname = T_KEYSPACE + '2'
            yield self.client.system_rename_keyspace(T_KEYSPACE, newname)
            ks2 = yield self.client.describe_keyspace(newname)
            ksdef.name = newname
            self.compare_keyspaces(ksdef, ks2)
        yield self.client.system_drop_keyspace(ksdef.name)
        yield self.assertFailure(self.client.describe_keyspace(T_KEYSPACE),
                                 NotFoundException)
        if DO_SYSTEM_RENAMING:
            yield self.assertFailure(self.client.describe_keyspace(ksdef.name),
                                     NotFoundException)

    @defer.inlineCallbacks
    def test_column_family_manipulation(self):
        cfdef = CfDef(
            KEYSPACE,
            T_CF,
            column_type='Standard',
            comparator_type='org.apache.cassandra.db.marshal.BytesType',
            comment='foo',
            row_cache_size=0.0,
            key_cache_size=200000.0,
            read_repair_chance=1.0,
            column_metadata=[],
            gc_grace_seconds=86400,
            default_validation_class=
            'org.apache.cassandra.db.marshal.BytesType',
            key_validation_class='org.apache.cassandra.db.marshal.BytesType',
            min_compaction_threshold=5,
            max_compaction_threshold=31,
            row_cache_save_period_in_seconds=0,
            key_cache_save_period_in_seconds=3600,
            memtable_flush_after_mins=60,
            memtable_throughput_in_mb=249,
            memtable_operations_in_millions=1.1671875,
            replicate_on_write=False,
            merge_shards_chance=0.10000000000000001,
            row_cache_provider=None,
            key_alias=None,
        )
        post_07_fields = [
            'replicate_on_write', 'merge_shards_chance',
            'key_validation_class', 'row_cache_provider', 'key_alias'
        ]
        post_08_fields = [
            'memtable_throughput_in_mb', 'memtable_flush_after_mins',
            'memtable_operations_in_millions'
        ]

        yield self.client.system_add_column_family(cfdef)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        cfdef2 = [c for c in ksdef.cf_defs if c.name == T_CF][0]

        for field in post_07_fields:
            # Most of these are ignored in 0.7, so we can't reliably compare them
            setattr(cfdef, field, None)
            setattr(cfdef2, field, None)

        for field in post_08_fields:
            # These fields change from 0.8 to 1.0
            setattr(cfdef, field, None)
            setattr(cfdef2, field, None)

        # we don't know the id ahead of time. copy the new one so the equality
        # comparison won't fail
        cfdef.id = cfdef2.id
        self.assertEqual(cfdef, cfdef2)
        if DO_SYSTEM_RENAMING:
            newname = T_CF + '2'
            yield self.client.system_rename_column_family(T_CF, newname)
            ksdef = yield self.client.describe_keyspace(KEYSPACE)
            cfdef2 = [c for c in ksdef.cf_defs if c.name == newname][0]
            self.assertNotIn(T_CF, [c.name for c in ksdef.cf_defs])
            cfdef.name = newname
            self.assertEqual(cfdef, cfdef2)
        yield self.client.system_drop_column_family(cfdef.name)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        self.assertNotIn(cfdef.name, [c.name for c in ksdef.cf_defs])

    @defer.inlineCallbacks
    def test_describes(self):
        name = yield self.client.describe_cluster_name()
        self.assertIsInstance(name, str)
        self.assertNotEqual(name, '')
        partitioner = yield self.client.describe_partitioner()
        self.assert_(partitioner.startswith('org.apache.cassandra.'),
                     msg='partitioner is %r' % partitioner)
        snitch = yield self.client.describe_snitch()
        self.assert_(snitch.startswith('org.apache.cassandra.'),
                     msg='snitch is %r' % snitch)
        version = yield self.client.describe_version()
        self.assertIsInstance(version, str)
        self.assertIn('.', version)
        schemavers = yield self.client.describe_schema_versions()
        self.assertIsInstance(schemavers, dict)
        self.assertNotEqual(schemavers, {})
        ring = yield self.client.describe_ring(KEYSPACE)
        self.assertIsInstance(ring, list)
        self.assertNotEqual(ring, [])
        for r in ring:
            self.assertIsInstance(r.start_token, str)
            self.assertIsInstance(r.end_token, str)
            self.assertIsInstance(r.endpoints, list)
            self.assertNotEqual(r.endpoints, [])
            for ep in r.endpoints:
                self.assertIsInstance(ep, str)

    @defer.inlineCallbacks
    def test_errback(self):
        yield self.client.remove('poiqwe', CF)
        try:
            yield self.client.get('poiqwe', CF, column='foo')
        except Exception, e:
            pass
Example #4
0
class CassandraSchemaManager(object):
    """Manages creation and destruction of cassandra schemas.

    Useful for both testing and production
    """
    def __init__(self, keyspace_def, error_if_existing=False):
        self.keyspace_def = keyspace_def
        self.error_if_existing = error_if_existing
        self.created_keyspace = False
        self.created_cfs = []

        self.client = None
        self.manager = None
        self.connector = None

    def connect(self, host=None, port=9160, username=None, password=None):
        if not host:
            host, port = get_host_port()

        if username or password:
            if not (username and password):
                raise CassandraConfigurationError(
                    "Specify both username and password or neither")
        else:
            username, password = get_credentials()
        authz = dict(username=username, password=password)

        self.manager = ManagedCassandraClientFactory(credentials=authz,
                                                     check_api_version=True)
        self.connector = reactor.connectTCP(host, port, self.manager)
        self.client = CassandraClient(self.manager)

    def disconnect(self):
        if self.manager:
            self.manager.shutdown()
        if self.connector:
            self.connector.disconnect()

    @timeout(DEFAULT_CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def create(self, truncate=False):
        if not self.client:
            self.connect()

        keyspace = self.keyspace_def

        try:
            existing = yield self.client.describe_keyspace(keyspace.name)
        except NotFoundException:
            existing = None

        # keyspace already exists
        if existing:
            yield self.client.set_keyspace(keyspace.name)
            _compare_ks_properties(existing, keyspace)

            existing_cfs = dict((cf.name, cf) for cf in existing.cf_defs)

            for cf in keyspace.cf_defs:
                if cf.name in existing_cfs:

                    if truncate:
                        # in truncate mode we drop and readd any existing CFs.
                        yield self.client.system_drop_column_family(cf.name)
                        yield self.client.system_add_column_family(cf)
                    else:
                        _compare_cf_properties(existing_cfs[cf.name], cf)
                else:
                    if cf.keyspace != keyspace.name:
                        raise CassandraSchemaError(
                            "CF %s has wrong keyspace name", cf.name)
                    self.created_cfs.append(cf.name)
                    yield self.client.system_add_column_family(cf)
        else:
            self.created_keyspace = True
            yield self.client.system_add_keyspace(keyspace)
            yield self.client.set_keyspace(keyspace.name)

    @timeout(DEFAULT_CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def teardown(self):
        if self.created_keyspace:
            yield self.client.system_drop_keyspace(self.keyspace_def.name)

        elif self.created_cfs:
            for cf in self.created_cfs:
                yield self.client.system_drop_column_family(cf)
class CassandraClientTest(unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self):
        self.cmanager = ManagedCassandraClientFactory(keyspace='system')
        self.client = CassandraClient(self.cmanager)
        for i in xrange(CONNS):
            reactor.connectTCP(HOST, PORT, self.cmanager)
        yield self.cmanager.deferred

        self.my_keyspace = KsDef(
            name=KEYSPACE,
            strategy_class='org.apache.cassandra.locator.SimpleStrategy',
            replication_factor=1,
            cf_defs=[
                CfDef(
                    keyspace=KEYSPACE,
                    name=CF,
                    column_type='Standard'
                ),
                CfDef(
                    keyspace=KEYSPACE,
                    name=SCF,
                    column_type='Super'
                ),
                CfDef(
                    keyspace=KEYSPACE,
                    name=IDX_CF,
                    column_type='Standard',
                    comparator_type='org.apache.cassandra.db.marshal.UTF8Type',
                    column_metadata=[
                        ColumnDef(
                            name='col1',
                            validation_class='org.apache.cassandra.db.marshal.UTF8Type',
                            index_type=IndexType.KEYS,
                            index_name='idxCol1')
                    ],
                    default_validation_class='org.apache.cassandra.db.marshal.BytesType')
            ]
        )
        yield self.client.system_add_keyspace(self.my_keyspace)
        yield self.client.set_keyspace(KEYSPACE)
    
    @defer.inlineCallbacks
    def tearDown(self):
        yield self.client.system_drop_keyspace(self.my_keyspace.name)
        self.cmanager.shutdown()
        for c in reactor.getDelayedCalls():
            c.cancel()
        reactor.removeAll()
    
    @defer.inlineCallbacks
    def test_insert_get(self): 
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        yield self.client.insert('test', SCF, 'superval', column=COLUMN, super_column=SCOLUMN)
        yield self.client.insert('test2', SCF, 'superval2', column=COLUMN,
                                 super_column=SCOLUMN)
        res = yield self.client.get('test', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        res = yield self.client.get('test2', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval2')
        res = yield self.client.get('test', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval')
        res = yield self.client.get('test2', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval2')

    @defer.inlineCallbacks
    def test_batch_insert_get_slice_and_count(self):
        yield self.client.batch_insert('test', CF,
                                       {COLUMN: 'test', COLUMN2: 'test2'})
        yield self.client.batch_insert('test', SCF,
                               {SCOLUMN: {COLUMN: 'test', COLUMN2: 'test2'}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2)) 
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_count('test', CF)
        self.assertEqual(res, 2)
        
    @defer.inlineCallbacks
    def test_batch_mutate_and_remove(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } }, 'test2': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } } })
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_remove({CF: ['test', 'test2']}, names=['test', 'test2'])
        yield self.client.batch_remove({SCF: ['test', 'test2']}, names=['test', 'test2'], supercolumn=SCOLUMN)

    @defer.inlineCallbacks
    def test_batch_mutate_with_deletion(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_mutate({'test': {CF: {COLUMN: None, COLUMN2: 'test3'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(len(res), 1)
        self.assertEqual(res[0].column.value, 'test3')

    @defer.inlineCallbacks
    def test_multiget_slice_remove(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        res = yield self.client.multiget_slice(['test', 'test2'], CF)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test'][1].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        yield self.client.remove('test', CF, column=COLUMN)
        yield self.client.remove('test2', CF, column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(len(res['test']), 0)
        self.assertEqual(len(res['test2']), 0)
        
    @defer.inlineCallbacks
    def test_range_slices(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        ks = yield self.client.get_range_slices(CF, start='', finish='')
        keys = [k.key for k in ks]
        for key in ['test', 'test2']:
            self.assertIn(key, keys)

    @defer.inlineCallbacks
    def test_indexed_slices(self):
        yield self.client.insert('test1', IDX_CF, 'one', column='col1')
        yield self.client.insert('test2', IDX_CF, 'two', column='col1')
        yield self.client.insert('test3', IDX_CF, 'three', column='col1')
        expressions = [IndexExpression('col1', IndexOperator.EQ, 'two')]
        res = yield self.client.get_indexed_slices(IDX_CF, expressions, start_key='')
        self.assertEquals(res[0].columns[0].column.value,'two')

    def sleep(self, secs):
        d = defer.Deferred()
        reactor.callLater(secs, d.callback, None)
        return d

    @defer.inlineCallbacks
    def test_ttls(self):
        yield self.client.insert('test_ttls', CF, 'testval', column=COLUMN, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_insert('test_ttls', CF, {COLUMN:'testval'}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

        yield self.client.batch_mutate({'test_ttls': {CF: {COLUMN: 'testval'}}}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), NotFoundException)

    @defer.inlineCallbacks
    def test_keyspace_manipulation(self):
        ksdef = KsDef(name=T_KEYSPACE, strategy_class='org.apache.cassandra.locator.SimpleStrategy', replication_factor=1, cf_defs=[])
        yield self.client.system_add_keyspace(ksdef)
        ks2 = yield self.client.describe_keyspace(T_KEYSPACE)
        self.assertEqual(ksdef, ks2)
        if DO_SYSTEM_RENAMING:
            newname = T_KEYSPACE + '2'
            yield self.client.system_rename_keyspace(T_KEYSPACE, newname)
            ks2 = yield self.client.describe_keyspace(newname)
            ksdef.name = newname
            self.assertEqual(ksdef, ks2)
        yield self.client.system_drop_keyspace(ksdef.name)
        yield self.assertFailure(self.client.describe_keyspace(T_KEYSPACE), NotFoundException)
        if DO_SYSTEM_RENAMING:
            yield self.assertFailure(self.client.describe_keyspace(ksdef.name), NotFoundException)

    @defer.inlineCallbacks
    def test_column_family_manipulation(self):
        cfdef = CfDef(KEYSPACE, T_CF,
            column_type='Standard',
            comparator_type='org.apache.cassandra.db.marshal.BytesType',
            comment='foo',
            row_cache_size=0.0,
            key_cache_size=200000.0,
            read_repair_chance=1.0,
            column_metadata=[],
            gc_grace_seconds=86400,
            default_validation_class='org.apache.cassandra.db.marshal.BytesType',
            min_compaction_threshold=5,
            max_compaction_threshold=31,
            row_cache_save_period_in_seconds=0,
            key_cache_save_period_in_seconds=3600,
            memtable_flush_after_mins=60,
            memtable_throughput_in_mb=249,
            memtable_operations_in_millions=1.1671875,
        )
        yield self.client.system_add_column_family(cfdef)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        cfdef2 = [c for c in ksdef.cf_defs if c.name == T_CF][0]
        # we don't know the id ahead of time. copy the new one so the equality
        # comparison won't fail
        cfdef.id = cfdef2.id
        self.assertEqual(cfdef, cfdef2)
        if DO_SYSTEM_RENAMING:
            newname = T_CF + '2'
            yield self.client.system_rename_column_family(T_CF, newname)
            ksdef = yield self.client.describe_keyspace(KEYSPACE)
            cfdef2 = [c for c in ksdef.cf_defs if c.name == newname][0]
            self.assertNotIn(T_CF, [c.name for c in ksdef.cf_defs])
            cfdef.name = newname
            self.assertEqual(cfdef, cfdef2)
        yield self.client.system_drop_column_family(cfdef.name)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        self.assertNotIn(cfdef.name, [c.name for c in ksdef.cf_defs])

    @defer.inlineCallbacks
    def test_describes(self):
        name = yield self.client.describe_cluster_name()
        self.assertIsInstance(name, str)
        self.assertNotEqual(name, '')
        partitioner = yield self.client.describe_partitioner()
        self.assert_(partitioner.startswith('org.apache.cassandra.'),
                     msg='partitioner is %r' % partitioner)
        snitch = yield self.client.describe_snitch()
        self.assert_(snitch.startswith('org.apache.cassandra.'),
                     msg='snitch is %r' % snitch)
        version = yield self.client.describe_version()
        self.assertIsInstance(version, str)
        self.assertIn('.', version)
        schemavers = yield self.client.describe_schema_versions()
        self.assertIsInstance(schemavers, dict)
        self.assertNotEqual(schemavers, {})
        ring = yield self.client.describe_ring(KEYSPACE)
        self.assertIsInstance(ring, list)
        self.assertNotEqual(ring, [])
        for r in ring:
            self.assertIsInstance(r.start_token, str)
            self.assertIsInstance(r.end_token, str)
            self.assertIsInstance(r.endpoints, list)
            self.assertNotEqual(r.endpoints, [])
            for ep in r.endpoints:
                self.assertIsInstance(ep, str)

    @defer.inlineCallbacks
    def test_errback(self):
        yield self.client.remove('poiqwe', CF)
        try:
            yield self.client.get('poiqwe', CF, column='foo')
        except Exception, e:
            pass
Example #6
0
class CassandraClientTest(unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self):
        self.cmanager = ManagedCassandraClientFactory(keyspace='system')
        self.client = CassandraClient(self.cmanager)
        for i in xrange(CONNS):
            reactor.connectTCP(HOST, PORT, self.cmanager)
        yield self.cmanager.deferred

        remote_ver = yield self.client.describe_version()
        self.version = tuple(map(int, remote_ver.split('.')))

        self.my_keyspace = ttypes.KsDef(
            name=KEYSPACE,
            strategy_class='org.apache.cassandra.locator.SimpleStrategy',
            strategy_options={},
            cf_defs=[
                ttypes.CfDef(
                    keyspace=KEYSPACE,
                    name=CF,
                    column_type='Standard'
                ),
                ttypes.CfDef(
                    keyspace=KEYSPACE,
                    name=SCF,
                    column_type='Super'
                ),
                ttypes.CfDef(
                    keyspace=KEYSPACE,
                    name=IDX_CF,
                    column_type='Standard',
                    comparator_type='org.apache.cassandra.db.marshal.UTF8Type',
                    column_metadata=[
                        ttypes.ColumnDef(
                            name='col1',
                            validation_class='org.apache.cassandra.db.marshal.UTF8Type',
                            index_type=ttypes.IndexType.KEYS,
                            index_name='idxCol1')
                    ],
                    default_validation_class='org.apache.cassandra.db.marshal.BytesType'
                ),
            ]
        )

        if self.version <= KS_RF_ATTRIBUTE:
            self.my_keyspace.replication_factor = 1
        else:
            self.my_keyspace.strategy_options['replication_factor'] = '1'

        if self.version >= COUNTERS_SUPPORTED_API:
            self.my_keyspace.cf_defs.extend([
                ttypes.CfDef(
                    keyspace=KEYSPACE,
                    name=COUNTER_CF,
                    column_type='Standard',
                    default_validation_class='org.apache.cassandra.db.marshal.CounterColumnType'
                ),
                ttypes.CfDef(
                    keyspace=KEYSPACE,
                    name=SUPERCOUNTER_CF,
                    column_type='Super',
                    default_validation_class='org.apache.cassandra.db.marshal.CounterColumnType'
                ),
            ])

        yield self.client.system_add_keyspace(self.my_keyspace)
        yield self.client.set_keyspace(KEYSPACE)

    @defer.inlineCallbacks
    def tearDown(self):
        yield self.client.system_drop_keyspace(self.my_keyspace.name)
        self.cmanager.shutdown()
        for c in reactor.getDelayedCalls():
            c.cancel()
        reactor.removeAll()

    @defer.inlineCallbacks
    def test_insert_get(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        yield self.client.insert('test', SCF, 'superval', column=COLUMN, super_column=SCOLUMN)
        yield self.client.insert('test2', SCF, 'superval2', column=COLUMN,
                                 super_column=SCOLUMN)
        res = yield self.client.get('test', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        res = yield self.client.get('test2', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval2')
        res = yield self.client.get('test', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval')
        res = yield self.client.get('test2', SCF, column=COLUMN, super_column=SCOLUMN)
        self.assertEqual(res.column.value, 'superval2')

    @defer.inlineCallbacks
    def test_batch_insert_get_slice_and_count(self):
        yield self.client.batch_insert('test', CF,
                                       {COLUMN: 'test', COLUMN2: 'test2'})
        yield self.client.batch_insert('test', SCF,
                               {SCOLUMN: {COLUMN: 'test', COLUMN2: 'test2'}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_count('test', CF)
        self.assertEqual(res, 2)

    @defer.inlineCallbacks
    def test_batch_mutate_and_remove(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } }, 'test2': {CF: {COLUMN: 'test', COLUMN2: 'test2'}, SCF: { SCOLUMN: { COLUMN: 'test', COLUMN2: 'test2'} } } })
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        res = yield self.client.get_slice('test2', SCF, names=(COLUMN, COLUMN2),
                                          super_column=SCOLUMN)
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_remove({CF: ['test', 'test2']}, names=['test', 'test2'])
        yield self.client.batch_remove({SCF: ['test', 'test2']}, names=['test', 'test2'], supercolumn=SCOLUMN)

    @defer.inlineCallbacks
    def test_batch_mutate_with_deletion(self):
        yield self.client.batch_mutate({'test': {CF: {COLUMN: 'test', COLUMN2: 'test2'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(res[0].column.value, 'test')
        self.assertEqual(res[1].column.value, 'test2')
        yield self.client.batch_mutate({'test': {CF: {COLUMN: None, COLUMN2: 'test3'}}})
        res = yield self.client.get_slice('test', CF, names=(COLUMN, COLUMN2))
        self.assertEqual(len(res), 1)
        self.assertEqual(res[0].column.value, 'test3')

    @defer.inlineCallbacks
    def test_multiget_slice_remove(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        res = yield self.client.multiget_slice(['test', 'test2'], CF)
        self.assertEqual(res['test'][0].column.value, 'testval')
        self.assertEqual(res['test'][1].column.value, 'testval')
        self.assertEqual(res['test2'][0].column.value, 'testval2')
        yield self.client.remove('test', CF, column=COLUMN)
        yield self.client.remove('test2', CF, column=COLUMN)
        res = yield self.client.multiget(['test', 'test2'], CF, column=COLUMN)
        self.assertEqual(len(res['test']), 0)
        self.assertEqual(len(res['test2']), 0)

    @defer.inlineCallbacks
    def test_range_slices(self):
        yield self.client.insert('test', CF, 'testval', column=COLUMN)
        yield self.client.insert('test', CF, 'testval', column=COLUMN2)
        yield self.client.insert('test2', CF, 'testval2', column=COLUMN)
        ks = yield self.client.get_range_slices(CF, start='', finish='')
        keys = [k.key for k in ks]
        for key in ['test', 'test2']:
            self.assertIn(key, keys)

    @defer.inlineCallbacks
    def test_indexed_slices(self):
        yield self.client.insert('test1', IDX_CF, 'one', column='col1')
        yield self.client.insert('test2', IDX_CF, 'two', column='col1')
        yield self.client.insert('test3', IDX_CF, 'three', column='col1')
        expressions = [ttypes.IndexExpression('col1', ttypes.IndexOperator.EQ, 'two')]
        res = yield self.client.get_indexed_slices(IDX_CF, expressions, start_key='')
        self.assertEquals(res[0].columns[0].column.value,'two')

    @defer.inlineCallbacks
    def test_counter_add(self):
        if self.version < COUNTERS_SUPPORTED_API:
            raise unittest.SkipTest('Counters are not supported before 0.8')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 2)

        # test super column counters
        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 2)

    @defer.inlineCallbacks
    def test_counter_remove(self):
        if self.version < COUNTERS_SUPPORTED_API:
            raise unittest.SkipTest('Counters are not supported before 0.8')

        # test standard column counter
        yield self.client.add('test', COUNTER_CF, 1, column='col')
        res = yield self.client.get('test', COUNTER_CF, column='col')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test', COUNTER_CF, column='col')
        yield self.assertFailure(self.client.get('test', COUNTER_CF, column='col'),
                                 ttypes.NotFoundException)

        # test super column counters
        yield self.client.add('test', SUPERCOUNTER_CF, 1, column='col', super_column='scol')
        res = yield self.client.get('test', SUPERCOUNTER_CF, column='col', super_column='scol')
        self.assertEquals(res.counter_column.value, 1)

        yield self.client.remove_counter('test', SUPERCOUNTER_CF,
                                         column='col', super_column='scol')
        yield self.assertFailure(self.client.get('test', SUPERCOUNTER_CF,
                                                 column='col', super_column='scol'),
                                 ttypes.NotFoundException)

    def sleep(self, secs):
        d = defer.Deferred()
        reactor.callLater(secs, d.callback, None)
        return d

    @defer.inlineCallbacks
    def test_ttls(self):
        yield self.client.insert('test_ttls', CF, 'testval', column=COLUMN, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), ttypes.NotFoundException)

        yield self.client.batch_insert('test_ttls', CF, {COLUMN:'testval'}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), ttypes.NotFoundException)

        yield self.client.batch_mutate({'test_ttls': {CF: {COLUMN: 'testval'}}}, ttl=1)
        res = yield self.client.get('test_ttls', CF, column=COLUMN)
        self.assertEqual(res.column.value, 'testval')
        yield self.sleep(2)
        yield self.assertFailure(self.client.get('test_ttls', CF, column=COLUMN), ttypes.NotFoundException)

    def compare_keyspaces(self, ks1, ks2):
        self.assertEqual(ks1.name, ks2.name)
        self.assertEqual(ks1.strategy_class, ks2.strategy_class)
        self.assertEqual(ks1.cf_defs, ks2.cf_defs)

        def get_rf(ksdef):
            rf = ksdef.replication_factor
            if ksdef.strategy_options and \
               'replication_factor' in ksdef.strategy_options:
                rf = int(ksdef.strategy_options['replication_factor'])
            return rf

        def strat_opts_no_rf(ksdef):
            if not ksdef.strategy_options:
                return {}
            opts = ksdef.strategy_options.copy()
            if 'replication_factor' in ksdef.strategy_options:
                del opts['replication_factor']
            return opts

        self.assertEqual(get_rf(ks1), get_rf(ks2))
        self.assertEqual(strat_opts_no_rf(ks1), strat_opts_no_rf(ks2))

    @defer.inlineCallbacks
    def test_keyspace_manipulation(self):
        try:
            yield self.client.system_drop_keyspace(T_KEYSPACE)
        except ttypes.InvalidRequestException:
            pass

        ksdef = ttypes.KsDef(name=T_KEYSPACE, strategy_class='org.apache.cassandra.locator.SimpleStrategy', strategy_options={}, cf_defs=[])
        if self.version <= KS_RF_ATTRIBUTE:
            ksdef.replication_factor = 1
        else:
            ksdef.strategy_options['replication_factor'] = '1'

        yield self.client.system_add_keyspace(ksdef)
        ks2 = yield self.client.describe_keyspace(T_KEYSPACE)
        self.compare_keyspaces(ksdef, ks2)

        if DO_SYSTEM_RENAMING:
            newname = T_KEYSPACE + '2'
            yield self.client.system_rename_keyspace(T_KEYSPACE, newname)
            ks2 = yield self.client.describe_keyspace(newname)
            ksdef.name = newname
            self.compare_keyspaces(ksdef, ks2)
        yield self.client.system_drop_keyspace(ksdef.name)
        yield self.assertFailure(self.client.describe_keyspace(T_KEYSPACE), ttypes.NotFoundException)
        if DO_SYSTEM_RENAMING:
            yield self.assertFailure(self.client.describe_keyspace(ksdef.name), ttypes.NotFoundException)

    @defer.inlineCallbacks
    def test_column_family_manipulation(self):
        # CfDef attributes present in all supported c*/thrift-api versions
        common_attrs = (
            ('column_type', 'Standard'),
            ('comparator_type', 'org.apache.cassandra.db.marshal.BytesType'),
            ('comment', 'foo'),
            ('read_repair_chance', 1.0),
            ('column_metadata', []),
            ('gc_grace_seconds', 86400),
            ('default_validation_class', 'org.apache.cassandra.db.marshal.BytesType'),
            ('min_compaction_threshold', 5),
            ('max_compaction_threshold', 31),
        )
        cfdef = ttypes.CfDef(KEYSPACE, T_CF)
        for attr, val in common_attrs:
            setattr(cfdef, attr, val)

        yield self.client.system_add_column_family(cfdef)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        cfdefs = [c for c in ksdef.cf_defs if c.name == T_CF]
        self.assertEqual(len(cfdefs), 1)
        cfdef2 = cfdefs[0]

        for attr, val in common_attrs:
            val1 = getattr(cfdef, attr)
            val2 = getattr(cfdef2, attr)
            self.assertEqual(val1, val2, 'attribute %s mismatch: %r != %r' % (attr, val1, val2))

        if DO_SYSTEM_RENAMING:
            newname = T_CF + '2'
            yield self.client.system_rename_column_family(T_CF, newname)
            ksdef = yield self.client.describe_keyspace(KEYSPACE)
            cfdef2 = [c for c in ksdef.cf_defs if c.name == newname][0]
            self.assertNotIn(T_CF, [c.name for c in ksdef.cf_defs])
            cfdef.name = newname
            self.assertEqual(cfdef, cfdef2)
        yield self.client.system_drop_column_family(cfdef.name)
        ksdef = yield self.client.describe_keyspace(KEYSPACE)
        self.assertNotIn(cfdef.name, [c.name for c in ksdef.cf_defs])

    @defer.inlineCallbacks
    def test_describes(self):
        name = yield self.client.describe_cluster_name()
        self.assertIsInstance(name, str)
        self.assertNotEqual(name, '')
        partitioner = yield self.client.describe_partitioner()
        self.assert_(partitioner.startswith('org.apache.cassandra.'),
                     msg='partitioner is %r' % partitioner)
        snitch = yield self.client.describe_snitch()
        self.assert_(snitch.startswith('org.apache.cassandra.'),
                     msg='snitch is %r' % snitch)
        version = yield self.client.describe_version()
        self.assertIsInstance(version, str)
        self.assertIn('.', version)
        schemavers = yield self.client.describe_schema_versions()
        self.assertIsInstance(schemavers, dict)
        self.assertNotEqual(schemavers, {})
        ring = yield self.client.describe_ring(KEYSPACE)
        self.assertIsInstance(ring, list)
        self.assertNotEqual(ring, [])
        for r in ring:
            self.assertIsInstance(r.start_token, str)
            self.assertIsInstance(r.end_token, str)
            self.assertIsInstance(r.endpoints, list)
            self.assertNotEqual(r.endpoints, [])
            for ep in r.endpoints:
                self.assertIsInstance(ep, str)

    @defer.inlineCallbacks
    def test_errback(self):
        yield self.client.remove('poiqwe', CF)
        try:
            yield self.client.get('poiqwe', CF, column='foo')
        except Exception, e:
            pass
Example #7
0
class CassandraSchemaManager(object):
    """Manages creation and destruction of cassandra schemas.

    Useful for both testing and production
    """

    def __init__(self, keyspace_def, error_if_existing=False):
        self.keyspace_def = keyspace_def
        self.error_if_existing=error_if_existing
        self.created_keyspace = False
        self.created_cfs = []

        self.client = None
        self.manager = None
        self.connector = None

    def connect(self, host=None, port=9160, username=None, password=None):
        if not host:
            host, port = get_host_port()

        if username or password:
            if not (username and password):
                raise CassandraConfigurationError(
                    "Specify both username and password or neither")
        else:
            username, password = get_credentials()
        authz = dict(username=username, password=password)

        self.manager = ManagedCassandraClientFactory(credentials=authz,
                                                     check_api_version=True)
        self.connector = reactor.connectTCP(host, port, self.manager)
        self.client = CassandraClient(self.manager)

    def disconnect(self):
        if self.manager:
            self.manager.shutdown()
        if self.connector:
            self.connector.disconnect()

    @timeout(DEFAULT_CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def create(self, truncate=False):
        if not self.client:
            self.connect()

        keyspace = self.keyspace_def

        try:
            existing = yield self.client.describe_keyspace(keyspace.name)
        except NotFoundException:
            existing = None

        # keyspace already exists
        if existing:
            yield self.client.set_keyspace(keyspace.name)
            _compare_ks_properties(existing, keyspace)

            existing_cfs = dict((cf.name, cf) for cf in existing.cf_defs)

            for cf in keyspace.cf_defs:
                if cf.name in existing_cfs:

                    if truncate:
                        # in truncate mode we drop and readd any existing CFs.
                        yield self.client.system_drop_column_family(cf.name)
                        yield self.client.system_add_column_family(cf)
                    else:
                        _compare_cf_properties(existing_cfs[cf.name], cf)
                else:
                    if cf.keyspace != keyspace.name:
                        raise CassandraSchemaError(
                            "CF %s has wrong keyspace name", cf.name)
                    self.created_cfs.append(cf.name)
                    yield self.client.system_add_column_family(cf)
        else:
            self.created_keyspace = True
            yield self.client.system_add_keyspace(keyspace)
            yield self.client.set_keyspace(keyspace.name)

    @timeout(DEFAULT_CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def teardown(self):
        if self.created_keyspace:
            yield self.client.system_drop_keyspace(self.keyspace_def.name)

        elif self.created_cfs:
            for cf in self.created_cfs:
                yield self.client.system_drop_column_family(cf)
Example #8
0
class CassandraControllerStore(TCPConnection):
    """Cassandra persistence for EPU controller state

    All EPU controllers within a system share the same column families.


    The "known" CFs hold a list of known instance IDs and sensor IDs for each
    EPU controller. The value is irrelevant and inserts are of course
    idempotent. These are used for figuring out all information about a
    controller without walking the entire instances and sensors CFs.
    
    ControllerKnownInstances = {
        Controller1 = {
            instance_id_1,
            instance_id_2
        },
        Controller2 = {
            instance_id_1,
            instance_id_2
        }
    }

    ControllerKnownSensors = {
        Controller1 = {
            sensor_1,
            sensor_2
        },
        Controller2 = {
            sensor_1,
            sensor_2
        }
    }

    All records for an instance are held in a single row. The column keys
    are TimeUUIDs. Because instance records have more complicated ordering
    than time, it is necessary to ensure a record is actually new before
    inserting. This is safe since in the current architecture there is only
    one writer for controller.

    One limitation is that out-of-order records cannot be inserted in the
    store. So if the messaging layer provides the STARTED state record for
    an instance *after* the RUNNING record, it cannot be stored and must
    be dropped. The correct state will be preserved, but not all of history
    will be.

    ControllerInstances = {               #comparator = TimeUUIDType
        Controller1Instance1 = {
            TimeUUID1 : 'the actual record',
            TimeUUID2 : 'the actual record',
            TimeUUID3 : 'the actual record',
        },
        Controller2Instance1 = {
            TimeUUID1 : 'the actual record',
            TimeUUID2 : 'the actual record',
            TimeUUID3 : 'the actual record',
        }
    }


    Sensor records are stored similarly. Instead of a TimeUUID, they use
    longs as keys which are likely to be a timestamp. Again, the controller
    must check the timestamp and not treat the most recently arrived value
    as the latest. However it can still write older values as they will be
    correctly inserted into history.

    ControllerSensors = {                 #comparator = LongType
        Controller1Sensor1 = {
            timestamp1 : 'sensor message',
            timestamp2 : 'sensor message'
        }
    }

    The controller starts up with a dictionary of config values that are
    passed to the decision engine. These can be changed by calling the
    reconfigure operation. Reconfigured values are stored here in
    Cassandra and on reboot/recovery, they are folded into the original
    config before it is passed to the decision engine.

    Each controller gets a row. Each column represents a single key/value
    pair where the key is a string and the value is a JSON-encoded object.

    ControllerEngineConfig = {
        Controller1 = {
            key1 : 'value1',
            key2 : 'value2'
        },
        Controller2 = {
            key1 : 'value1',
            key2 : 'value2'
        }
    }
    """

    CONFIG_CF_NAME = "ControllerEngineConfig"
    INSTANCE_CF_NAME = "ControllerInstances"
    INSTANCE_ID_CF_NAME = "ControllerKnownInstances"
    SENSOR_CF_NAME = "ControllerSensors"
    SENSOR_ID_CF_NAME = "ControllerKnownSensors"

    _PAGE_SIZE = 100

    @classmethod
    def get_column_families(cls, keyspace=None, prefix=''):
        """Builds a list of column families needed by this store.
        @param keyspace Name of keyspace. If None, it must be added manually.
        @param prefix Optional prefix for cf names. Useful for testing.
        @retval list of CfDef objects
        """
        instance_cf=prefix+cls.INSTANCE_CF_NAME
        instance_id_cf=prefix+cls.INSTANCE_ID_CF_NAME
        sensor_cf=prefix+cls.SENSOR_CF_NAME
        sensor_id_cf=prefix+cls.SENSOR_ID_CF_NAME
        config_cf = prefix+cls.CONFIG_CF_NAME

        return [CfDef(keyspace, instance_cf,
                  comparator_type='org.apache.cassandra.db.marshal.TimeUUIDType'),
                CfDef(keyspace, instance_id_cf,
                  comparator_type='org.apache.cassandra.db.marshal.UTF8Type'),
                CfDef(keyspace, sensor_cf,
                  comparator_type='org.apache.cassandra.db.marshal.LongType'),
                CfDef(keyspace, sensor_id_cf,
                  comparator_type='org.apache.cassandra.db.marshal.UTF8Type'),
                CfDef(keyspace, config_cf,
                  comparator_type='org.apache.cassandra.db.marshal.UTF8Type'),
                ]

    def __init__(self, controller_name, host, port, username, password,
                 keyspace, instance_factory, sensor_item_factory,
                 prefix=''):

        self.controller_name = str(controller_name)
        # keep a set of known instances and sensors so we can save on
        # unnecessary inserts to the controller instance/sensor lists
        self.seen_instances = set()
        self.seen_sensors = set()

        self.instance_factory = instance_factory
        self.sensor_item_factory = sensor_item_factory

        authorization_dictionary = {'username': username, 'password': password}

        self.manager = ManagedCassandraClientFactory(
                credentials=authorization_dictionary,
                check_api_version=True, keyspace=keyspace)

        TCPConnection.__init__(self, host, port, self.manager)
        self.client = CassandraClient(self.manager)

        self.instance_cf = prefix + self.INSTANCE_CF_NAME
        self.instance_id_cf = prefix + self.INSTANCE_ID_CF_NAME
        self.sensor_cf = prefix + self.SENSOR_CF_NAME
        self.sensor_id_cf = prefix + self.SENSOR_ID_CF_NAME
        self.config_cf = prefix + self.CONFIG_CF_NAME

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def check_schema(self):
        ks = yield self.client.describe_keyspace(self.manager.keyspace)
        cfs = dict((cf.name,cf) for cf in ks.cf_defs)

        missing = [cf for cf in (self.instance_cf, self.instance_id_cf,
                                 self.sensor_cf, self.sensor_id_cf,
                                 self.config_cf)
                   if cf in cfs]
        if missing:
            error = "EPU Controller is missing Cassandra column families: %s"
            raise Exception(error % ", ".join(missing))

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def add_instance(self, instance):
        """Adds a new instance object to persistence
        @param instance Instance to add
        @retval Deferred
        """

        instance_id = str(instance.instance_id)
        if instance_id not in self.seen_instances:
            yield self.client.insert(self.controller_name, self.instance_id_cf,
                                     "", column=instance_id)
            self.seen_instances.add(instance_id)

        key = self.controller_name + instance_id
        value = json.dumps(dict(instance.iteritems()))
        col = uuid.uuid1().bytes
        yield self.client.insert(key, self.instance_cf, value, column=col)

    @timeout(CASSANDRA_TIMEOUT)
    def get_instance_ids(self):
        """Retrieves a list of known instances

        @retval Deferred of list of instance IDs
        """
        return self._get_ids(self.instance_id_cf)

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def get_instance(self, instance_id):
        """Retrieves the latest instance object for the specified id
        @param instance_id ID of instance to retrieve
        @retval Deferred of Instance object or None
        """

        key = self.controller_name + str(instance_id)
        slice = yield self.client.get_slice(key, self.instance_cf,
                                          reverse=True, count=1)

        if slice:
            d = json.loads(slice[0].column.value)
            ret = self.instance_factory(**d)
        else:
            ret = None
        defer.returnValue(ret)

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def add_sensor(self, sensor):
        """Adds a new sensor object to persistence
        @param sensor Sensor to add
        @retval Deferred
        """

        sensor_id = str(sensor.sensor_id)
        if sensor_id not in self.seen_sensors:
            yield self.client.insert(self.controller_name, self.sensor_id_cf,
                                     "", column=sensor_id)
            self.seen_sensors.add(sensor_id)

        key = self.controller_name + sensor_id
        value = json.dumps(sensor.value)
        col = struct.pack('!Q', int(sensor.time))
        yield self.client.insert(key, self.sensor_cf, value, column=col)

    @timeout(CASSANDRA_TIMEOUT)
    def get_sensor_ids(self):
        """Retrieves a list of known sensors

        @retval Deferred of list of sensor IDs
        """
        return self._get_ids(self.sensor_id_cf)

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def get_sensor(self, sensor_id):
        """Retrieve the latest sensor item for the specified sensor

        @param sensor_id ID of the sensor item to retrieve
        @retval Deferred of SensorItem object or None
        """
        key = self.controller_name + str(sensor_id)
        slice = yield self.client.get_slice(key, self.sensor_cf,
                                          reverse=True, count=1)

        if slice:
            col = slice[0].column
            timestamp = struct.unpack("!Q", col.name)[0]
            val = json.loads(col.value)
            ret = self.sensor_item_factory(sensor_id, long(timestamp), val)
        else:
            ret = None

        defer.returnValue(ret)

    @timeout(CASSANDRA_TIMEOUT)
    @defer.inlineCallbacks
    def get_config(self, keys=None):
        """Retrieve the engine config dictionary.

        @param keys optional list of keys to retrieve
        @retval Deferred of config dictionary object
        """
        key = self.controller_name
        slice = yield self.client.get_slice(key, self.config_cf, names=keys)
        cfg = {}
        if slice:
            for col in slice:
                key = col.column.name
                val = col.column.value
                cfg[key] = json.loads(val)
        defer.returnValue(cfg)

    @timeout(CASSANDRA_TIMEOUT)
    def add_config(self, conf):
        """Store a dictionary of new engine conf values.

        These are folded into the existing configuration map. So for example
        if you first store {'a' : 1, 'b' : 1} and then store {'b' : 2},
        the result from get_config() will be {'a' : 1, 'b' : 2}.

        @param conf dictionary mapping strings to JSON-serializable objects
        @retval Deferred
        """
        d = dict((k, json.dumps(v)) for k,v in conf.iteritems())
        return self.client.batch_insert(self.controller_name, self.config_cf, d)

    @defer.inlineCallbacks
    def _get_ids(self, cf):
        """Retrieves IDs from either instance or sensor column families
        """
        # using a set because it isn't totally clear if the start parameter
        # to get_slice() is always inclusive or exclusive. Seeing mixed
        # messages on mailing list, so also not convinced if this isn't
        # something that has changed, or might.
        found_ids = set()
        start = ""
        done = False
        while not done:
            slice = yield self.client.get_slice(self.controller_name, cf,
                                                count=self._PAGE_SIZE,
                                                start=start)
            if slice:
                found_ids.update(col.column.name for col in slice)

                if len(slice) == self._PAGE_SIZE:
                    start = slice[-1].column.name
                else:
                    done = True
            else:
                done = True
        defer.returnValue(list(found_ids))

    def on_deactivate(self, *args, **kwargs):
        self.manager.shutdown()

    def on_terminate(self, *args, **kwargs):
        self.manager.shutdown()