add virtual metadata to parser

mambocab · mambocab · commit a6d0626205d9 · 2018-08-17T10:52:28.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,5 +1,10 @@
 3.15.0
 ======
+
+Features
+--------
+* Parse Virtual Keyspace Metadata (PYTHON-992)
+
 Bug Fixes
 ---------
 * Tokenmap.get_replicas returns the wrong value if token coincides with the end of the range (PYTHON-978)
diff --git a/cassandra/metadata.py b/cassandra/metadata.py
@@ -639,6 +639,14 @@ class KeyspaceMetadata(object):
     A dict mapping view names to :class:`.MaterializedViewMetadata` instances.
     """
 
+    virtual = False
+    """
+    A boolean indicating if this is a virtual keyspace or not. Always ``False``
+    for clusters running pre-4.0 versions of Cassandra.
+
+    .. versionadded:: 3.15
+    """
+
     _exc_info = None
     """ set if metadata parsing failed """
 
@@ -671,13 +679,20 @@ def export_as_string(self):
                 ret += line
             ret += "\nApproximate structure, for reference:\n(this should not be used to reproduce this schema)\n\n%s\n*/" % cql
             return ret
+        if self.virtual:
+            return ("/*\nWarning: Keyspace {ks} is a virtual keyspace and cannot be recreated with CQL.\n"
+                    "Structure, for reference:*/\n"
+                    "{cql}\n"
+                    "").format(ks=self.name, cql=cql)
         return cql
 
     def as_cql_query(self):
         """
         Returns a CQL query string that can be used to recreate just this keyspace,
         not including user-defined types and tables.
         """
+        if self.virtual:
+            return "// VIRTUAL KEYSPACE {}".format(protect_name(self.name))
         ret = "CREATE KEYSPACE %s WITH replication = %s " % (
             protect_name(self.name),
             self.replication_strategy.export_for_schema())
@@ -1065,11 +1080,21 @@ def primary_key(self):
     _exc_info = None
     """ set if metadata parsing failed """
 
+    virtual = False
+    """
+    A boolean indicating if this is a virtual table or not. Always ``False``
+    for clusters running pre-4.0 versions of Cassandra.
+
+    .. versionadded:: 3.15
+    """
+
     @property
     def is_cql_compatible(self):
         """
         A boolean indicating if this table can be represented as CQL in export
         """
+        if self.virtual:
+            return False
         comparator = getattr(self, 'comparator', None)
         if comparator:
             # no compact storage with more than one column beyond PK if there
@@ -1086,7 +1111,7 @@ def is_cql_compatible(self):
     Metadata describing configuration for table extensions
     """
 
-    def __init__(self, keyspace_name, name, partition_key=None, clustering_key=None, columns=None, triggers=None, options=None):
+    def __init__(self, keyspace_name, name, partition_key=None, clustering_key=None, columns=None, triggers=None, options=None, virtual=False):
         self.keyspace_name = keyspace_name
         self.name = name
         self.partition_key = [] if partition_key is None else partition_key
@@ -1097,6 +1122,7 @@ def __init__(self, keyspace_name, name, partition_key=None, clustering_key=None,
         self.comparator = None
         self.triggers = OrderedDict() if triggers is None else triggers
         self.views = {}
+        self.virtual = virtual
 
     def export_as_string(self):
         """
@@ -1116,6 +1142,11 @@ def export_as_string(self):
             ret = "/*\nWarning: Table %s.%s omitted because it has constructs not compatible with CQL (was created via legacy API).\n" % \
                   (self.keyspace_name, self.name)
             ret += "\nApproximate structure, for reference:\n(this should not be used to reproduce this schema)\n\n%s\n*/" % self._all_as_cql()
+        elif self.virtual:
+            ret = ('/*\nWarning: Table {ks}.{tab} is a virtual table and cannot be recreated with CQL.\n'
+                   'Structure, for reference:\n'
+                   '{cql}\n*/').format(ks=self.keyspace_name, tab=self.name, cql=self._all_as_cql())
+
         else:
             ret = self._all_as_cql()
 
@@ -1150,7 +1181,8 @@ def as_cql_query(self, formatted=False):
         creations are not included).  If `formatted` is set to :const:`True`,
         extra whitespace will be added to make the query human readable.
         """
-        ret = "CREATE TABLE %s.%s (%s" % (
+        ret = "%s TABLE %s.%s (%s" % (
+            ('VIRTUAL' if self.virtual else 'CREATE'),
             protect_name(self.keyspace_name),
             protect_name(self.name),
             "\n" if formatted else "")
@@ -2245,28 +2277,32 @@ def _build_aggregate(aggregate_row):
                          aggregate_row['argument_types'], aggregate_row['state_func'], aggregate_row['state_type'],
                          aggregate_row['final_func'], aggregate_row['initcond'], aggregate_row['return_type'])
 
-    def _build_table_metadata(self, row, col_rows=None, trigger_rows=None, index_rows=None):
+    def _build_table_metadata(self, row, col_rows=None, trigger_rows=None, index_rows=None, virtual=False):
         keyspace_name = row["keyspace_name"]
         table_name = row[self._table_name_col]
 
         col_rows = col_rows or self.keyspace_table_col_rows[keyspace_name][table_name]
         trigger_rows = trigger_rows or self.keyspace_table_trigger_rows[keyspace_name][table_name]
         index_rows = index_rows or self.keyspace_table_index_rows[keyspace_name][table_name]
 
-        table_meta = TableMetadataV3(keyspace_name, table_name)
+        table_meta = TableMetadataV3(keyspace_name, table_name, virtual=virtual)
         try:
             table_meta.options = self._build_table_options(row)
             flags = row.get('flags', set())
             if flags:
                 compact_static = False
                 table_meta.is_compact_storage = 'dense' in flags or 'super' in flags or 'compound' not in flags
                 is_dense = 'dense' in flags
+            elif virtual:
+                compact_static = False
+                table_meta.is_compact_storage = False
+                is_dense = False
             else:
                 compact_static = True
                 table_meta.is_compact_storage = True
                 is_dense = False
 
-            self._build_table_columns(table_meta, col_rows, compact_static, is_dense)
+            self._build_table_columns(table_meta, col_rows, compact_static, is_dense, virtual)
 
             for trigger_row in trigger_rows:
                 trigger_meta = self._build_trigger_metadata(table_meta, trigger_row)
@@ -2288,7 +2324,7 @@ def _build_table_options(self, row):
         """ Setup the mostly-non-schema table options, like caching settings """
         return dict((o, row.get(o)) for o in self.recognized_table_options if o in row)
 
-    def _build_table_columns(self, meta, col_rows, compact_static=False, is_dense=False):
+    def _build_table_columns(self, meta, col_rows, compact_static=False, is_dense=False, virtual=False):
         # partition key
         partition_rows = [r for r in col_rows
                           if r.get('kind', None) == "partition_key"]
@@ -2432,6 +2468,116 @@ class SchemaParserV4(SchemaParserV3):
         )
     )
 
+    _SELECT_VIRTUAL_KEYSPACES = 'SELECT * from system_virtual_schema.keyspaces'
+    _SELECT_VIRTUAL_TABLES = 'SELECT * from system_virtual_schema.tables'
+    _SELECT_VIRTUAL_COLUMNS = 'SELECT * from system_virtual_schema.columns'
+
+    def __init__(self, connection, timeout):
+        super(SchemaParserV4, self).__init__(connection, timeout)
+        self.virtual_keyspaces_rows = defaultdict(list)
+        self.virtual_tables_rows = defaultdict(list)
+        self.virtual_columns_rows = defaultdict(lambda: defaultdict(list))
+
+    def _query_all(self):
+        cl = ConsistencyLevel.ONE
+        # todo: this duplicates V3; we should find a way for _query_all methods
+        # to extend each other.
+        queries = [
+            # copied from V3
+            QueryMessage(query=self._SELECT_KEYSPACES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_TABLES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_COLUMNS, consistency_level=cl),
+            QueryMessage(query=self._SELECT_TYPES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_FUNCTIONS, consistency_level=cl),
+            QueryMessage(query=self._SELECT_AGGREGATES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_TRIGGERS, consistency_level=cl),
+            QueryMessage(query=self._SELECT_INDEXES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_VIEWS, consistency_level=cl),
+            # V4-only queries
+            QueryMessage(query=self._SELECT_VIRTUAL_KEYSPACES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_VIRTUAL_TABLES, consistency_level=cl),
+            QueryMessage(query=self._SELECT_VIRTUAL_COLUMNS, consistency_level=cl)
+        ]
+
+        responses = self.connection.wait_for_responses(
+            *queries, timeout=self.timeout, fail_on_error=False)
+        (
+            # copied from V3
+            (ks_success, ks_result),
+            (table_success, table_result),
+            (col_success, col_result),
+            (types_success, types_result),
+            (functions_success, functions_result),
+            (aggregates_success, aggregates_result),
+            (triggers_success, triggers_result),
+            (indexes_success, indexes_result),
+            (views_success, views_result),
+            # V4-only responses
+            (virtual_ks_success, virtual_ks_result),
+            (virtual_table_success, virtual_table_result),
+            (virtual_column_success, virtual_column_result)
+        ) = responses
+
+        # copied from V3
+        self.keyspaces_result = self._handle_results(ks_success, ks_result)
+        self.tables_result = self._handle_results(table_success, table_result)
+        self.columns_result = self._handle_results(col_success, col_result)
+        self.triggers_result = self._handle_results(triggers_success, triggers_result)
+        self.types_result = self._handle_results(types_success, types_result)
+        self.functions_result = self._handle_results(functions_success, functions_result)
+        self.aggregates_result = self._handle_results(aggregates_success, aggregates_result)
+        self.indexes_result = self._handle_results(indexes_success, indexes_result)
+        self.views_result = self._handle_results(views_success, views_result)
+        # V4-only results
+        self.virtual_keyspaces_result = self._handle_results(virtual_ks_success,
+                                                             virtual_ks_result)
+        self.virtual_tables_result = self._handle_results(virtual_table_success,
+                                                          virtual_table_result)
+        self.virtual_columns_result = self._handle_results(virtual_column_success,
+                                                           virtual_column_result)
+        self._aggregate_results()
+
+    def _aggregate_results(self):
+        super(SchemaParserV4, self)._aggregate_results()
+
+        m = self.virtual_tables_rows
+        for row in self.virtual_tables_result:
+            m[row["keyspace_name"]].append(row)
+
+        m = self.virtual_columns_rows
+        for row in self.virtual_columns_result:
+            ks_name = row['keyspace_name']
+            tab_name = row[self._table_name_col]
+            m[ks_name][tab_name].append(row)
+
+    def get_all_keyspaces(self):
+        for x in super(SchemaParserV4, self).get_all_keyspaces():
+            yield x
+
+        for row in self.virtual_keyspaces_result:
+            ks_name = row['keyspace_name']
+            keyspace_meta = self._build_keyspace_metadata(row)
+            keyspace_meta.virtual = True
+
+            for table_row in self.virtual_tables_rows.get(ks_name, []):
+                table_name = table_row[self._table_name_col]
+
+                col_rows = self.virtual_columns_rows[ks_name][table_name]
+                keyspace_meta._add_table_metadata(
+                    self._build_table_metadata(table_row,
+                                               col_rows=col_rows,
+                                               virtual=True)
+                )
+            yield keyspace_meta
+
+    @staticmethod
+    def _build_keyspace_metadata_internal(row):
+        # necessary fields that aren't int virtual ks
+        row["durable_writes"] = row.get("durable_writes", None)
+        row["replication"] = row.get("replication", {})
+        row["replication"]["class"] = row["replication"].get("class", None)
+        return super(SchemaParserV4, SchemaParserV4)._build_keyspace_metadata_internal(row)
+
 
 class TableMetadataV3(TableMetadata):
     compaction_options = {}
diff --git a/tests/integration/standard/test_metadata.py b/tests/integration/standard/test_metadata.py
@@ -17,7 +17,9 @@
 except ImportError:
     import unittest  # noqa
 
+from collections import defaultdict
 import difflib
+import logging
 import six
 import sys
 import time
@@ -36,11 +38,14 @@
                                BasicSegregatedKeyspaceUnitTestCase, BasicSharedKeyspaceUnitTestCase,
                                BasicExistingKeyspaceUnitTestCase, drop_keyspace_shutdown_cluster, CASSANDRA_VERSION,
                                get_supported_protocol_versions, greaterthanorequalcass30, lessthancass30, local,
-                               greaterthancass20)
+                               greaterthancass20, greaterthanorequalcass40)
 
 from tests.integration import greaterthancass21
 
 
+log = logging.getLogger(__name__)
+
+
 def setup_module():
     use_singledc()
 
@@ -2334,7 +2339,7 @@ def test_group_keys_by_host(self):
 
         stmt = """SELECT * FROM {}.{}
                                  WHERE k = ? """.format(self.ks_name, self.ks_name)
-        keys = ((1, ), (2, ), (2, ), (3, ))
+        keys = ((1,), (2,), (2,), (3,))
         self._assert_group_keys_by_host(keys, self.ks_name, stmt)
 
     def _assert_group_keys_by_host(self, keys, table_name, stmt):
@@ -2347,3 +2352,72 @@ def _assert_group_keys_by_host(self, keys, table_name, stmt):
             hosts = self.cluster.metadata.get_replicas(self.ks_name, routing_key)
             self.assertEqual(1, len(hosts))  # RF is 1 for this keyspace
             self.assertIn(key, keys_per_host[hosts[0]])
+
+
+class VirtualKeypaceTest(BasicSharedKeyspaceUnitTestCase):
+    virtual_ks_names = ('system_virtual_schema', 'system_views')
+
+    virtual_ks_structure = {
+        'system_views': {
+            # map from table names to sets of column names for unordered
+            # comparison
+            'caches': {'capacity_bytes', 'entry_count', 'hit_count',
+                       'hit_ratio', 'name', 'recent_hit_rate_per_second',
+                       'recent_request_rate_per_second', 'request_count',
+                       'size_bytes'},
+            'clients': {'address', 'connection_stage', 'driver_name',
+                        'driver_version', 'hostname', 'port',
+                        'protocol_version', 'request_count',
+                        'ssl_cipher_suite', 'ssl_enabled', 'ssl_protocol',
+                        'username'},
+            'sstable_tasks': {'keyspace_name', 'kind', 'progress',
+                              'table_name', 'task_id', 'total', 'unit'},
+            'thread_pools': {'active_tasks', 'active_tasks_limit',
+                             'blocked_tasks', 'blocked_tasks_all_time',
+                             'completed_tasks', 'name', 'pending_tasks'}
+        },
+        'system_virtual_schema': {
+            'columns': {'clustering_order', 'column_name',
+                        'column_name_bytes', 'keyspace_name', 'kind',
+                        'position', 'table_name', 'type'},
+            'keyspaces': {'keyspace_name'},
+            'tables': {'comment', 'keyspace_name', 'table_name'}
+        }
+    }
+
+    def test_existing_keyspaces_have_correct_virtual_tags(self):
+        for name, ks in self.cluster.metadata.keyspaces.items():
+            if name in self.virtual_ks_names:
+                self.assertTrue(
+                    ks.virtual,
+                    'incorrect .virtual value for {}'.format(name)
+                )
+            else:
+                self.assertFalse(
+                    ks.virtual,
+                    'incorrect .virtual value for {}'.format(name)
+                )
+
+    @greaterthanorequalcass40
+    def test_expected_keyspaces_exist_and_are_virtual(self):
+        for name in self.virtual_ks_names:
+            self.assertTrue(
+                self.cluster.metadata.keyspaces[name].virtual,
+                'incorrect .virtual value for {}'.format(name)
+            )
+
+    @greaterthanorequalcass40
+    def test_virtual_keyspaces_have_expected_schema_structure(self):
+        self.maxDiff = None
+
+        ingested_virtual_ks_structure = defaultdict(dict)
+        for ks_name, ks in self.cluster.metadata.keyspaces.items():
+            if not ks.virtual:
+                continue
+            for tab_name, tab in ks.tables.items():
+                ingested_virtual_ks_structure[ks_name][tab_name] = set(
+                    tab.columns.keys()
+                )
+
+        self.assertDictEqual(ingested_virtual_ks_structure,
+                             self.virtual_ks_structure)
