Merge pull request datastax#978 from datastax/python-992_virtual-keyspaces

Python 992 virtual keyspaces

Author: mambocab

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)

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 "")
@@ -2064,12 +2096,16 @@ def _query_all(self):
             QueryMessage(query=self._SELECT_TRIGGERS, consistency_level=cl)
         ]
 
-        responses = self.connection.wait_for_responses(*queries, timeout=self.timeout, fail_on_error=False)
-        ((ks_success, ks_result), (table_success, table_result),
-         (col_success, col_result), (types_success, types_result),
+        ((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)) = responses
+         (triggers_success, triggers_result)) = (
+             self.connection.wait_for_responses(*queries, timeout=self.timeout,
+                                                fail_on_error=False)
+        )
 
         self.keyspaces_result = self._handle_results(ks_success, ks_result)
         self.tables_result = self._handle_results(table_success, table_result)
@@ -2241,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)
@@ -2284,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"]
@@ -2375,14 +2415,17 @@ def _query_all(self):
             QueryMessage(query=self._SELECT_VIEWS, consistency_level=cl)
         ]
 
-        responses = self.connection.wait_for_responses(*queries, timeout=self.timeout, fail_on_error=False)
-        ((ks_success, ks_result), (table_success, table_result),
-         (col_success, col_result), (types_success, types_result),
+        ((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)) = responses
+         (views_success, views_result)) = self.connection.wait_for_responses(
+             *queries, timeout=self.timeout, fail_on_error=False
+        )
 
         self.keyspaces_result = self._handle_results(ks_success, ks_result)
         self.tables_result = self._handle_results(table_success, table_result)
@@ -2425,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 = {}