diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a0a66395..654ad5e4 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,6 +12,7 @@ Bug Fixes * Raise ValueError when tuple query parameters for prepared statements have extra items (PYTHON-98) * Correctly encode nested tuples and UDTs for non-prepared statements (PYTHON-100) +* Include User Defined Types in KeyspaceMetadata.export_as_string() (PYTHON-96) Other ----- diff --git a/cassandra/metadata.py b/cassandra/metadata.py index 46ee3a82..8aa77c38 100644 --- a/cassandra/metadata.py +++ b/cassandra/metadata.py @@ -450,7 +450,7 @@ class SimpleStrategy(ReplicationStrategy): while len(hosts) < self.replication_factor and j < len(ring): token = ring[(i + j) % len(ring)] host = token_to_host_owner[token] - if not host in hosts: + if host not in hosts: hosts.append(host) j += 1 @@ -597,7 +597,7 @@ class KeyspaceMetadata(object): self.user_types = {} def export_as_string(self): - return "\n".join([self.as_cql_query()] + [t.export_as_string() for t in self.tables.values()]) + return "\n\n".join([self.as_cql_query()] + self.user_type_strings() + [t.export_as_string() for t in self.tables.values()]) def as_cql_query(self): ret = "CREATE KEYSPACE %s WITH replication = %s " % ( @@ -605,6 +605,22 @@ class KeyspaceMetadata(object): self.replication_strategy.export_for_schema()) return ret + (' AND durable_writes = %s;' % ("true" if self.durable_writes else "false")) + def user_type_strings(self): + user_type_strings = [] + types = self.user_types.copy() + keys = sorted(types.keys()) + for k in keys: + if k in types: + self.resolve_user_types(k, types, user_type_strings) + return user_type_strings + + def resolve_user_types(self, key, types, user_type_strings): + user_type = types.pop(key) + for field_type in user_type.field_types: + if field_type.cassname == 'UserType' and field_type.typename in types: + self.resolve_user_types(field_type.typename, types, user_type_strings) + user_type_strings.append(user_type.as_cql_query(formatted=True)) + class UserType(object): """ @@ -664,7 +680,7 @@ class UserType(object): fields.append("%s %s" % (protect_name(field_name), field_type.cql_parameterized_type())) ret += field_join.join("%s%s" % (padding, field) for field in fields) - ret += "\n)" if formatted else ")" + ret += "\n);" if formatted else ");" return ret diff --git a/tests/unit/test_metadata.py b/tests/unit/test_metadata.py index dd42040a..28d56cff 100644 --- a/tests/unit/test_metadata.py +++ b/tests/unit/test_metadata.py @@ -15,7 +15,9 @@ try: import unittest2 as unittest except ImportError: - import unittest # noqa + import unittest # noqa + +from mock import Mock import cassandra from cassandra.cqltypes import IntegerType, AsciiType, TupleType @@ -24,7 +26,7 @@ from cassandra.metadata import (Murmur3Token, MD5Token, NetworkTopologyStrategy, SimpleStrategy, LocalStrategy, NoMurmur3, protect_name, protect_names, protect_value, is_valid_name, - UserType) + UserType, KeyspaceMetadata) from cassandra.policies import SimpleConvictionPolicy from cassandra.pool import Host @@ -183,12 +185,12 @@ class TestNameEscaping(unittest.TestCase): 'tests ?!@#$%^&*()', '1' ]), - [ - 'tests', - "\"test's\"", - '"tests ?!@#$%^&*()"', - '"1"' - ]) + [ + 'tests', + "\"test's\"", + '"tests ?!@#$%^&*()"', + '"1"' + ]) def test_protect_value(self): """ @@ -252,19 +254,61 @@ class TestTokens(unittest.TestCase): pass +class TestKeyspaceMetadata(unittest.TestCase): + + def test_export_as_string_user_types(self): + keyspace_name = 'test' + keyspace = KeyspaceMetadata(keyspace_name, True, 'SimpleStrategy', dict(replication_factor=3)) + keyspace.user_types['a'] = UserType(keyspace_name, 'a', ['one', 'two'], + [self.mock_user_type('UserType', 'c'), + self.mock_user_type('IntType', 'int')]) + keyspace.user_types['b'] = UserType(keyspace_name, 'b', ['one', 'two', 'three'], + [self.mock_user_type('UserType', 'd'), + self.mock_user_type('IntType', 'int'), + self.mock_user_type('UserType', 'a')]) + keyspace.user_types['c'] = UserType(keyspace_name, 'c', ['one'], + [self.mock_user_type('IntType', 'int')]) + keyspace.user_types['d'] = UserType(keyspace_name, 'd', ['one'], + [self.mock_user_type('UserType', 'c')]) + + self.assertEqual("""CREATE KEYSPACE test WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '3'} AND durable_writes = true; + +CREATE TYPE test.c ( + one int +); + +CREATE TYPE test.a ( + one c, + two int +); + +CREATE TYPE test.d ( + one c +); + +CREATE TYPE test.b ( + one d, + two int, + three a +);""", keyspace.export_as_string()) + + def mock_user_type(self, cassname, typename): + return Mock(**{'cassname': cassname, 'typename': typename, 'cql_parameterized_type.return_value': typename}) + + class TestUserTypes(unittest.TestCase): def test_as_cql_query(self): field_types = [IntegerType, AsciiType, TupleType.apply_parameters([IntegerType, AsciiType])] udt = UserType("ks1", "mytype", ["a", "b", "c"], field_types) - self.assertEqual("CREATE TYPE ks1.mytype (a varint, b ascii, c tuple)", udt.as_cql_query(formatted=False)) + self.assertEqual("CREATE TYPE ks1.mytype (a varint, b ascii, c tuple);", udt.as_cql_query(formatted=False)) self.assertEqual("""CREATE TYPE ks1.mytype ( a varint, b ascii, c tuple -)""", udt.as_cql_query(formatted=True)) +);""", udt.as_cql_query(formatted=True)) def test_as_cql_query_name_escaping(self): udt = UserType("MyKeyspace", "MyType", ["AbA", "keyspace"], [AsciiType, AsciiType]) - self.assertEqual('CREATE TYPE "MyKeyspace"."MyType" ("AbA" ascii, "keyspace" ascii)', udt.as_cql_query(formatted=False)) + self.assertEqual('CREATE TYPE "MyKeyspace"."MyType" ("AbA" ascii, "keyspace" ascii);', udt.as_cql_query(formatted=False))