Merge pull request #284 from datastax/PYTHON-241

PYTHON-241 - Meta indexes for Index Meta

cassandra/metadata.py

@@ -122,10 +122,8 @@ class Metadata(object):
             keyspace_col_rows = col_def_rows.get(keyspace_meta.name, {})
             keyspace_trigger_rows = trigger_rows.get(keyspace_meta.name, {})
             for table_row in cf_def_rows.get(keyspace_meta.name, []):
-                table_meta = self._build_table_metadata(
+                table_meta = self._build_table_metadata(keyspace_meta, table_row, keyspace_col_rows, keyspace_trigger_rows)
-                    keyspace_meta, table_row, keyspace_col_rows,
+                keyspace_meta._add_table_metadata(table_meta)
-                    keyspace_trigger_rows)
-                keyspace_meta.tables[table_meta.name] = table_meta
 
             for usertype_row in usertype_rows.get(keyspace_meta.name, []):
                 usertype = self._build_usertype(keyspace_meta.name, usertype_row)
@@ -160,6 +158,7 @@ class Metadata(object):
         if old_keyspace_meta:
             keyspace_meta.tables = old_keyspace_meta.tables
             keyspace_meta.user_types = old_keyspace_meta.user_types
+            keyspace_meta.indexes = old_keyspace_meta.indexes
             if (keyspace_meta.replication_strategy != old_keyspace_meta.replication_strategy):
                 self._keyspace_updated(keyspace)
         else:
@@ -184,12 +183,11 @@ class Metadata(object):
 
         if not cf_results:
             # the table was removed
-            keyspace_meta.tables.pop(table, None)
+            keyspace_meta._drop_table_metadata(table)
         else:
             assert len(cf_results) == 1
-            keyspace_meta.tables[table] = self._build_table_metadata(
+            table_meta = self._build_table_metadata(keyspace_meta, cf_results[0], {table: col_results}, {table: triggers_result})
-                keyspace_meta, cf_results[0], {table: col_results},
+            keyspace_meta._add_table_metadata(table_meta)
-                {table: triggers_result})
 
     def _keyspace_added(self, ksname):
         if self.token_map:
@@ -385,6 +383,8 @@ class Metadata(object):
         column_meta = ColumnMetadata(table_metadata, name, data_type, is_static=is_static)
         index_meta = self._build_index_metadata(column_meta, row)
         column_meta.index = index_meta
+        if index_meta:
+            table_metadata.indexes[index_meta.name] = index_meta
         return column_meta
 
     def _build_index_metadata(self, column_metadata, row):
@@ -733,6 +733,11 @@ class KeyspaceMetadata(object):
     A map from table names to instances of :class:`~.TableMetadata`.
     """
 
+    indexes = None
+    """
+    A dict mapping index names to :class:`.IndexMetadata` instances.
+    """
+
     user_types = None
     """
     A map from user-defined type names to instances of :class:`~cassandra.metadata..UserType`.
@@ -745,6 +750,7 @@ class KeyspaceMetadata(object):
         self.durable_writes = durable_writes
         self.replication_strategy = ReplicationStrategy.create(strategy_class, strategy_options)
         self.tables = {}
+        self.indexes = {}
         self.user_types = {}
 
     def export_as_string(self):
@@ -780,6 +786,18 @@ class KeyspaceMetadata(object):
                 self.resolve_user_types(field_type.typename, types, user_type_strings)
         user_type_strings.append(user_type.as_cql_query(formatted=True))
 
+    def _add_table_metadata(self, table_metadata):
+        self._drop_table_metadata(table_metadata.name)
+
+        self.tables[table_metadata.name] = table_metadata
+        for index_name, index_metadata in six.iteritems(table_metadata.indexes):
+            self.indexes[index_name] = index_metadata
+
+    def _drop_table_metadata(self, table_name):
+        table_meta = self.tables.pop(table_name, None)
+        if table_meta:
+            for index_name in table_meta.indexes:
+                self.indexes.pop(index_name, None)
 
 class UserType(object):
     """
@@ -884,6 +902,11 @@ class TableMetadata(object):
     A dict mapping column names to :class:`.ColumnMetadata` instances.
     """
 
+    indexes = None
+    """
+    A dict mapping index names to :class:`.IndexMetadata` instances.
+    """
+
     is_compact_storage = False
 
     options = None
@@ -945,6 +968,7 @@ class TableMetadata(object):
         self.partition_key = [] if partition_key is None else partition_key
         self.clustering_key = [] if clustering_key is None else clustering_key
         self.columns = OrderedDict() if columns is None else columns
+        self.indexes = {}
         self.options = options
         self.comparator = None
         self.triggers = OrderedDict() if triggers is None else triggers
@@ -1242,6 +1266,12 @@ class IndexMetadata(object):
                 protect_name(self.column.name),
                 self.index_options["class_name"])
 
+    def export_as_string(self):
+        """
+        Returns a CQL query string that can be used to recreate this index.
+        """
+        return self.as_cql_query() + ';'
+
 
 class TokenMap(object):
     """

tests/integration/standard/test_metadata.py

@@ -25,7 +25,7 @@ import sys
 from cassandra import AlreadyExists
 
 from cassandra.cluster import Cluster
-from cassandra.metadata import (Metadata, KeyspaceMetadata, TableMetadata,
+from cassandra.metadata import (Metadata, KeyspaceMetadata, TableMetadata, IndexMetadata,
                                 Token, MD5Token, TokenMap, murmur3)
 from cassandra.policies import SimpleConvictionPolicy
 from cassandra.pool import Host
@@ -900,3 +900,98 @@ class KeyspaceAlterMetadata(unittest.TestCase):
         new_keyspace_meta = self.cluster.metadata.keyspaces[name]
         self.assertNotEqual(original_keyspace_meta, new_keyspace_meta)
         self.assertEqual(new_keyspace_meta.durable_writes, False)
+
+
+class IndexMapTests(unittest.TestCase):
+
+    keyspace_name = 'index_map_tests'
+
+    @property
+    def table_name(self):
+        return self._testMethodName.lower()
+
+    @classmethod
+    def setup_class(cls):
+        cls.cluster = Cluster(protocol_version=PROTOCOL_VERSION)
+        cls.session = cls.cluster.connect()
+        try:
+            if cls.keyspace_name in cls.cluster.metadata.keyspaces:
+                cls.session.execute("DROP KEYSPACE %s" % cls.keyspace_name)
+
+            cls.session.execute(
+                """
+                CREATE KEYSPACE %s
+                WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'};
+                """ % cls.keyspace_name)
+            cls.session.set_keyspace(cls.keyspace_name)
+        except Exception:
+            cls.cluster.shutdown()
+            raise
+
+    @classmethod
+    def teardown_class(cls):
+        try:
+            cls.session.execute("DROP KEYSPACE %s" % cls.keyspace_name)
+        finally:
+            cls.cluster.shutdown()
+
+    def create_basic_table(self):
+        self.session.execute("CREATE TABLE %s (k int PRIMARY KEY, a int)" % self.table_name)
+
+    def drop_basic_table(self):
+        self.session.execute("DROP TABLE %s" % self.table_name)
+
+    def test_index_updates(self):
+        self.create_basic_table()
+
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        table_meta = ks_meta.tables[self.table_name]
+        self.assertNotIn('a_idx', ks_meta.indexes)
+        self.assertNotIn('b_idx', ks_meta.indexes)
+        self.assertNotIn('a_idx', table_meta.indexes)
+        self.assertNotIn('b_idx', table_meta.indexes)
+
+        self.session.execute("CREATE INDEX a_idx ON %s (a)" % self.table_name)
+        self.session.execute("ALTER TABLE %s ADD b int" % self.table_name)
+        self.session.execute("CREATE INDEX b_idx ON %s (b)" % self.table_name)
+
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        table_meta = ks_meta.tables[self.table_name]
+        self.assertIsInstance(ks_meta.indexes['a_idx'], IndexMetadata)
+        self.assertIsInstance(ks_meta.indexes['b_idx'], IndexMetadata)
+        self.assertIsInstance(table_meta.indexes['a_idx'], IndexMetadata)
+        self.assertIsInstance(table_meta.indexes['b_idx'], IndexMetadata)
+
+        # both indexes updated when index dropped
+        self.session.execute("DROP INDEX a_idx")
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        table_meta = ks_meta.tables[self.table_name]
+        self.assertNotIn('a_idx', ks_meta.indexes)
+        self.assertIsInstance(ks_meta.indexes['b_idx'], IndexMetadata)
+        self.assertNotIn('a_idx', table_meta.indexes)
+        self.assertIsInstance(table_meta.indexes['b_idx'], IndexMetadata)
+
+        # keyspace index updated when table dropped
+        self.drop_basic_table()
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        self.assertNotIn(self.table_name, ks_meta.tables)
+        self.assertNotIn('a_idx', ks_meta.indexes)
+        self.assertNotIn('b_idx', ks_meta.indexes)
+
+    def test_index_follows_alter(self):
+        self.create_basic_table()
+
+        idx = self.table_name + '_idx'
+        self.session.execute("CREATE INDEX %s ON %s (a)" % (idx, self.table_name))
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        table_meta = ks_meta.tables[self.table_name]
+        self.assertIsInstance(ks_meta.indexes[idx], IndexMetadata)
+        self.assertIsInstance(table_meta.indexes[idx], IndexMetadata)
+        self.session.execute('ALTER KEYSPACE %s WITH durable_writes = false' % self.keyspace_name)
+        old_meta = ks_meta
+        ks_meta = self.cluster.metadata.keyspaces[self.keyspace_name]
+        self.assertIsNot(ks_meta, old_meta)
+        table_meta = ks_meta.tables[self.table_name]
+        self.assertIsInstance(ks_meta.indexes[idx], IndexMetadata)
+        self.assertIsInstance(table_meta.indexes[idx], IndexMetadata)
+        self.drop_basic_table()