| # Copyright DataStax, Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| import unittest |
| |
| from datetime import datetime, date, time |
| from decimal import Decimal |
| from unittest.mock import Mock |
| from uuid import UUID, uuid4 |
| |
| from cassandra.cqlengine.models import Model |
| from cassandra.cqlengine.usertype import UserType, UserTypeDefinitionException |
| from cassandra.cqlengine import columns, connection |
| from cassandra.cqlengine.management import sync_table, drop_table, sync_type, create_keyspace_simple, drop_keyspace |
| from cassandra.cqlengine import ValidationError |
| from cassandra.util import Date, Time |
| |
| from tests.integration import PROTOCOL_VERSION |
| from tests.integration.cqlengine.base import BaseCassEngTestCase |
| from tests.integration.cqlengine import DEFAULT_KEYSPACE |
| |
| |
| class User(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| |
| |
| class UserModel(Model): |
| id = columns.Integer(primary_key=True) |
| info = columns.UserDefinedType(User) |
| |
| |
| class AllDatatypes(UserType): |
| a = columns.Ascii() |
| b = columns.BigInt() |
| c = columns.Blob() |
| d = columns.Boolean() |
| e = columns.DateTime() |
| f = columns.Decimal() |
| g = columns.Double() |
| h = columns.Float() |
| i = columns.Inet() |
| j = columns.Integer() |
| k = columns.Text() |
| l = columns.TimeUUID() |
| m = columns.UUID() |
| n = columns.VarInt() |
| |
| |
| class AllDatatypesModel(Model): |
| id = columns.Integer(primary_key=True) |
| data = columns.UserDefinedType(AllDatatypes) |
| |
| |
| class UserDefinedTypeTests(BaseCassEngTestCase): |
| |
| def setUp(self): |
| if PROTOCOL_VERSION < 3: |
| raise unittest.SkipTest("UDTs require native protocol 3+, currently using: {0}".format(PROTOCOL_VERSION)) |
| |
| def test_can_create_udts(self): |
| class User(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| |
| sync_type(DEFAULT_KEYSPACE, User) |
| user = User(age=42, name="John") |
| self.assertEqual(42, user.age) |
| self.assertEqual("John", user.name) |
| |
| # Add a field |
| class User(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| gender = columns.Text() |
| |
| sync_type(DEFAULT_KEYSPACE, User) |
| user = User(age=42) |
| user["name"] = "John" |
| user["gender"] = "male" |
| self.assertEqual(42, user.age) |
| self.assertEqual("John", user.name) |
| self.assertEqual("male", user.gender) |
| |
| # Remove a field |
| class User(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| |
| sync_type(DEFAULT_KEYSPACE, User) |
| user = User(age=42, name="John", gender="male") |
| with self.assertRaises(AttributeError): |
| user.gender |
| |
| def test_can_insert_udts(self): |
| |
| sync_table(UserModel) |
| self.addCleanup(drop_table, UserModel) |
| |
| user = User(age=42, name="John") |
| UserModel.create(id=0, info=user) |
| |
| self.assertEqual(1, UserModel.objects.count()) |
| |
| john = UserModel.objects.first() |
| self.assertEqual(0, john.id) |
| self.assertTrue(type(john.info) is User) |
| self.assertEqual(42, john.info.age) |
| self.assertEqual("John", john.info.name) |
| |
| def test_can_update_udts(self): |
| sync_table(UserModel) |
| self.addCleanup(drop_table, UserModel) |
| |
| user = User(age=42, name="John") |
| created_user = UserModel.create(id=0, info=user) |
| |
| john_info = UserModel.objects.first().info |
| self.assertEqual(42, john_info.age) |
| self.assertEqual("John", john_info.name) |
| |
| created_user.info = User(age=22, name="Mary") |
| created_user.update() |
| |
| mary_info = UserModel.objects.first().info |
| self.assertEqual(22, mary_info["age"]) |
| self.assertEqual("Mary", mary_info["name"]) |
| |
| def test_can_update_udts_with_nones(self): |
| sync_table(UserModel) |
| self.addCleanup(drop_table, UserModel) |
| |
| user = User(age=42, name="John") |
| created_user = UserModel.create(id=0, info=user) |
| |
| john_info = UserModel.objects.first().info |
| self.assertEqual(42, john_info.age) |
| self.assertEqual("John", john_info.name) |
| |
| created_user.info = None |
| created_user.update() |
| |
| john_info = UserModel.objects.first().info |
| self.assertIsNone(john_info) |
| |
| def test_can_create_same_udt_different_keyspaces(self): |
| sync_type(DEFAULT_KEYSPACE, User) |
| |
| create_keyspace_simple("simplex", 1) |
| sync_type("simplex", User) |
| drop_keyspace("simplex") |
| |
| def test_can_insert_partial_udts(self): |
| class UserGender(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| gender = columns.Text() |
| |
| class UserModelGender(Model): |
| id = columns.Integer(primary_key=True) |
| info = columns.UserDefinedType(UserGender) |
| |
| sync_table(UserModelGender) |
| self.addCleanup(drop_table, UserModelGender) |
| |
| user = UserGender(age=42, name="John") |
| UserModelGender.create(id=0, info=user) |
| |
| john_info = UserModelGender.objects.first().info |
| self.assertEqual(42, john_info.age) |
| self.assertEqual("John", john_info.name) |
| self.assertIsNone(john_info.gender) |
| |
| user = UserGender(age=42) |
| UserModelGender.create(id=0, info=user) |
| |
| john_info = UserModelGender.objects.first().info |
| self.assertEqual(42, john_info.age) |
| self.assertIsNone(john_info.name) |
| self.assertIsNone(john_info.gender) |
| |
| def test_can_insert_nested_udts(self): |
| class Depth_0(UserType): |
| age = columns.Integer() |
| name = columns.Text() |
| |
| class Depth_1(UserType): |
| value = columns.UserDefinedType(Depth_0) |
| |
| class Depth_2(UserType): |
| value = columns.UserDefinedType(Depth_1) |
| |
| class Depth_3(UserType): |
| value = columns.UserDefinedType(Depth_2) |
| |
| class DepthModel(Model): |
| id = columns.Integer(primary_key=True) |
| v_0 = columns.UserDefinedType(Depth_0) |
| v_1 = columns.UserDefinedType(Depth_1) |
| v_2 = columns.UserDefinedType(Depth_2) |
| v_3 = columns.UserDefinedType(Depth_3) |
| |
| sync_table(DepthModel) |
| self.addCleanup(drop_table, DepthModel) |
| |
| udts = [Depth_0(age=42, name="John")] |
| udts.append(Depth_1(value=udts[0])) |
| udts.append(Depth_2(value=udts[1])) |
| udts.append(Depth_3(value=udts[2])) |
| |
| DepthModel.create(id=0, v_0=udts[0], v_1=udts[1], v_2=udts[2], v_3=udts[3]) |
| output = DepthModel.objects.first() |
| |
| self.assertEqual(udts[0], output.v_0) |
| self.assertEqual(udts[1], output.v_1) |
| self.assertEqual(udts[2], output.v_2) |
| self.assertEqual(udts[3], output.v_3) |
| |
| def test_can_insert_udts_with_nones(self): |
| """ |
| Test for inserting all column types as empty into a UserType as None's |
| |
| test_can_insert_udts_with_nones tests that each cqlengine column type can be inserted into a UserType as None's. |
| It first creates a UserType that has each cqlengine column type, and a corresponding table/Model. It then creates |
| a UserType instance where all the fields are None's and inserts the UserType as an instance of the Model. Finally, |
| it verifies that each column read from the UserType from Cassandra is None. |
| |
| @since 2.5.0 |
| @jira_ticket PYTHON-251 |
| @expected_result The UserType is inserted with each column type, and the resulting read yields None's for each column. |
| |
| @test_category data_types:udt |
| """ |
| sync_table(AllDatatypesModel) |
| self.addCleanup(drop_table, AllDatatypesModel) |
| |
| input = AllDatatypes(a=None, b=None, c=None, d=None, e=None, f=None, g=None, h=None, i=None, j=None, k=None, |
| l=None, m=None, n=None) |
| AllDatatypesModel.create(id=0, data=input) |
| |
| self.assertEqual(1, AllDatatypesModel.objects.count()) |
| |
| output = AllDatatypesModel.objects.first().data |
| self.assertEqual(input, output) |
| |
| def test_can_insert_udts_with_all_datatypes(self): |
| """ |
| Test for inserting all column types into a UserType |
| |
| test_can_insert_udts_with_all_datatypes tests that each cqlengine column type can be inserted into a UserType. |
| It first creates a UserType that has each cqlengine column type, and a corresponding table/Model. It then creates |
| a UserType instance where all the fields have corresponding data, and inserts the UserType as an instance of the Model. |
| Finally, it verifies that each column read from the UserType from Cassandra is the same as the input parameters. |
| |
| @since 2.5.0 |
| @jira_ticket PYTHON-251 |
| @expected_result The UserType is inserted with each column type, and the resulting read yields proper data for each column. |
| |
| @test_category data_types:udt |
| """ |
| sync_table(AllDatatypesModel) |
| self.addCleanup(drop_table, AllDatatypesModel) |
| |
| input = AllDatatypes(a='ascii', b=2 ** 63 - 1, c=bytearray(b'hello world'), d=True, |
| e=datetime.utcfromtimestamp(872835240), f=Decimal('12.3E+7'), g=2.39, |
| h=3.4028234663852886e+38, i='123.123.123.123', j=2147483647, k='text', |
| l=UUID('FE2B4360-28C6-11E2-81C1-0800200C9A66'), |
| m=UUID('067e6162-3b6f-4ae2-a171-2470b63dff00'), n=int(str(2147483647) + '000')) |
| AllDatatypesModel.create(id=0, data=input) |
| |
| self.assertEqual(1, AllDatatypesModel.objects.count()) |
| output = AllDatatypesModel.objects.first().data |
| |
| for i in range(ord('a'), ord('a') + 14): |
| self.assertEqual(input[chr(i)], output[chr(i)]) |
| |
| def test_can_insert_udts_protocol_v4_datatypes(self): |
| """ |
| Test for inserting all protocol v4 column types into a UserType |
| |
| test_can_insert_udts_protocol_v4_datatypes tests that each protocol v4 cqlengine column type can be inserted |
| into a UserType. It first creates a UserType that has each protocol v4 cqlengine column type, and a corresponding |
| table/Model. It then creates a UserType instance where all the fields have corresponding data, and inserts the |
| UserType as an instance of the Model. Finally, it verifies that each column read from the UserType from Cassandra |
| is the same as the input parameters. |
| |
| @since 2.6.0 |
| @jira_ticket PYTHON-245 |
| @expected_result The UserType is inserted with each protocol v4 column type, and the resulting read yields proper data for each column. |
| |
| @test_category data_types:udt |
| """ |
| |
| if PROTOCOL_VERSION < 4: |
| raise unittest.SkipTest("Protocol v4 datatypes in UDTs require native protocol 4+, currently using: {0}".format(PROTOCOL_VERSION)) |
| |
| class Allv4Datatypes(UserType): |
| a = columns.Date() |
| b = columns.SmallInt() |
| c = columns.Time() |
| d = columns.TinyInt() |
| |
| class Allv4DatatypesModel(Model): |
| id = columns.Integer(primary_key=True) |
| data = columns.UserDefinedType(Allv4Datatypes) |
| |
| sync_table(Allv4DatatypesModel) |
| self.addCleanup(drop_table, Allv4DatatypesModel) |
| |
| input = Allv4Datatypes(a=Date(date(1970, 1, 1)), b=32523, c=Time(time(16, 47, 25, 7)), d=123) |
| Allv4DatatypesModel.create(id=0, data=input) |
| |
| self.assertEqual(1, Allv4DatatypesModel.objects.count()) |
| output = Allv4DatatypesModel.objects.first().data |
| |
| for i in range(ord('a'), ord('a') + 3): |
| self.assertEqual(input[chr(i)], output[chr(i)]) |
| |
| def test_nested_udts_inserts(self): |
| """ |
| Test for inserting collections of user types using cql engine. |
| |
| test_nested_udts_inserts Constructs a model that contains a list of usertypes. It will then attempt to insert |
| them. The expectation is that no exception is thrown during insert. For sanity sake we also validate that our |
| input and output values match. This combination of model, and UT produces a syntax error in 2.5.1 due to |
| improper quoting around the names collection. |
| |
| @since 2.6.0 |
| @jira_ticket PYTHON-311 |
| @expected_result No syntax exception thrown |
| |
| @test_category data_types:udt |
| """ |
| |
| class Name(UserType): |
| type_name__ = "header" |
| |
| name = columns.Text() |
| value = columns.Text() |
| |
| class Container(Model): |
| id = columns.UUID(primary_key=True, default=uuid4) |
| names = columns.List(columns.UserDefinedType(Name)) |
| |
| # Construct the objects and insert them |
| names = [] |
| for i in range(0, 10): |
| names.append(Name(name="name{0}".format(i), value="value{0}".format(i))) |
| |
| # Create table, insert data |
| sync_table(Container) |
| self.addCleanup(drop_table, Container) |
| |
| Container.create(id=UUID('FE2B4360-28C6-11E2-81C1-0800200C9A66'), names=names) |
| |
| # Validate input and output matches |
| self.assertEqual(1, Container.objects.count()) |
| names_output = Container.objects.first().names |
| self.assertEqual(names_output, names) |
| |
| def test_udts_with_unicode(self): |
| """ |
| Test for inserting models with unicode and udt columns. |
| |
| test_udts_with_unicode constructs a model with a user defined type. It then attempts to insert that model with |
| a unicode primary key. It will also attempt to upsert a udt that contains unicode text. |
| |
| @since 3.0.0 |
| @jira_ticket PYTHON-353 |
| @expected_result No exceptions thrown |
| |
| @test_category data_types:udt |
| """ |
| ascii_name = 'normal name' |
| unicode_name = u'Fran\u00E7ois' |
| |
| class UserModelText(Model): |
| id = columns.Text(primary_key=True) |
| info = columns.UserDefinedType(User) |
| |
| sync_table(UserModelText) |
| self.addCleanup(drop_table, UserModelText) |
| |
| # Two udt instances one with a unicode one with ascii |
| user_template_ascii = User(age=25, name=ascii_name) |
| user_template_unicode = User(age=25, name=unicode_name) |
| |
| UserModelText.create(id=ascii_name, info=user_template_unicode) |
| UserModelText.create(id=unicode_name, info=user_template_ascii) |
| UserModelText.create(id=unicode_name, info=user_template_unicode) |
| |
| def test_register_default_keyspace(self): |
| |
| from cassandra.cqlengine import models |
| from cassandra.cqlengine import connection |
| |
| # None emulating no model and no default keyspace before connecting |
| connection.udt_by_keyspace.clear() |
| User.register_for_keyspace(None) |
| self.assertEqual(len(connection.udt_by_keyspace), 1) |
| self.assertIn(None, connection.udt_by_keyspace) |
| |
| # register should be with default keyspace, not None |
| cluster = Mock() |
| connection._register_known_types(cluster) |
| cluster.register_user_type.assert_called_with(models.DEFAULT_KEYSPACE, User.type_name(), User) |
| |
| def test_db_field_override(self): |
| """ |
| Tests for db_field override |
| |
| Tests to ensure that udt's in models can specify db_field for a particular field and that it will be honored. |
| |
| @since 3.1.0 |
| @jira_ticket PYTHON-346 |
| @expected_result The actual cassandra column will use the db_field specified. |
| |
| @test_category data_types:udt |
| """ |
| class db_field_different(UserType): |
| age = columns.Integer(db_field='a') |
| name = columns.Text(db_field='n') |
| |
| class TheModel(Model): |
| id = columns.Integer(primary_key=True) |
| info = columns.UserDefinedType(db_field_different) |
| |
| sync_table(TheModel) |
| self.addCleanup(drop_table, TheModel) |
| |
| cluster = connection.get_cluster() |
| type_meta = cluster.metadata.keyspaces[TheModel._get_keyspace()].user_types[db_field_different.type_name()] |
| |
| type_fields = (db_field_different.age.column, db_field_different.name.column) |
| |
| self.assertEqual(len(type_meta.field_names), len(type_fields)) |
| for f in type_fields: |
| self.assertIn(f.db_field_name, type_meta.field_names) |
| |
| id = 0 |
| age = 42 |
| name = 'John' |
| info = db_field_different(age=age, name=name) |
| TheModel.create(id=id, info=info) |
| |
| self.assertEqual(1, TheModel.objects.count()) |
| |
| john = TheModel.objects.first() |
| self.assertEqual(john.id, id) |
| info = john.info |
| self.assertIsInstance(info, db_field_different) |
| self.assertEqual(info.age, age) |
| self.assertEqual(info.name, name) |
| # also excercise the db_Field mapping |
| self.assertEqual(info.a, age) |
| self.assertEqual(info.n, name) |
| |
| def test_db_field_overload(self): |
| """ |
| Tests for db_field UserTypeDefinitionException |
| |
| Test so that when we override a model's default field witha db_field that it errors appropriately |
| |
| @since 3.1.0 |
| @jira_ticket PYTHON-346 |
| @expected_result Setting a db_field to an existing field causes an exception to occur. |
| |
| @test_category data_types:udt |
| """ |
| |
| with self.assertRaises(UserTypeDefinitionException): |
| class something_silly(UserType): |
| first_col = columns.Integer() |
| second_col = columns.Text(db_field='first_col') |
| |
| with self.assertRaises(UserTypeDefinitionException): |
| class something_silly_2(UserType): |
| first_col = columns.Integer(db_field="second_col") |
| second_col = columns.Text() |
| |
| def test_set_udt_fields(self): |
| # PYTHON-502 |
| |
| u = User() |
| u.age = 20 |
| self.assertEqual(20, u.age) |
| |
| def test_default_values(self): |
| """ |
| Test that default types are set on object creation for UDTs |
| |
| @since 3.7.0 |
| @jira_ticket PYTHON-606 |
| @expected_result Default values should be set. |
| |
| @test_category data_types:udt |
| """ |
| |
| class NestedUdt(UserType): |
| |
| test_id = columns.UUID(default=uuid4) |
| something = columns.Text() |
| default_text = columns.Text(default="default text") |
| |
| class OuterModel(Model): |
| |
| name = columns.Text(primary_key=True) |
| first_name = columns.Text() |
| nested = columns.List(columns.UserDefinedType(NestedUdt)) |
| simple = columns.UserDefinedType(NestedUdt) |
| |
| sync_table(OuterModel) |
| self.addCleanup(drop_table, OuterModel) |
| |
| t = OuterModel.create(name='test1') |
| t.nested = [NestedUdt(something='test')] |
| t.simple = NestedUdt(something="") |
| t.save() |
| self.assertIsNotNone(t.nested[0].test_id) |
| self.assertEqual(t.nested[0].default_text, "default text") |
| self.assertIsNotNone(t.simple.test_id) |
| self.assertEqual(t.simple.default_text, "default text") |
| |
| def test_udt_validate(self): |
| """ |
| Test to verify restrictions are honored and that validate is called |
| for each member of the UDT when an updated is attempted |
| |
| @since 3.10 |
| @jira_ticket PYTHON-505 |
| @expected_result a validation error is arisen due to the name being |
| too long |
| |
| @test_category data_types:object_mapper |
| """ |
| class UserValidate(UserType): |
| age = columns.Integer() |
| name = columns.Text(max_length=2) |
| |
| class UserModelValidate(Model): |
| id = columns.Integer(primary_key=True) |
| info = columns.UserDefinedType(UserValidate) |
| |
| sync_table(UserModelValidate) |
| self.addCleanup(drop_table, UserModelValidate) |
| |
| user = UserValidate(age=1, name="Robert") |
| item = UserModelValidate(id=1, info=user) |
| with self.assertRaises(ValidationError): |
| item.save() |
| |
| def test_udt_validate_with_default(self): |
| """ |
| Test to verify restrictions are honored and that validate is called |
| on the default value |
| |
| @since 3.10 |
| @jira_ticket PYTHON-505 |
| @expected_result a validation error is arisen due to the name being |
| too long |
| |
| @test_category data_types:object_mapper |
| """ |
| class UserValidateDefault(UserType): |
| age = columns.Integer() |
| name = columns.Text(max_length=2, default="Robert") |
| |
| class UserModelValidateDefault(Model): |
| id = columns.Integer(primary_key=True) |
| info = columns.UserDefinedType(UserValidateDefault) |
| |
| sync_table(UserModelValidateDefault) |
| self.addCleanup(drop_table, UserModelValidateDefault) |
| |
| user = UserValidateDefault(age=1) |
| item = UserModelValidateDefault(id=1, info=user) |
| with self.assertRaises(ValidationError): |
| item.save() |
