diff --git a/django_spanner/schema.py b/django_spanner/schema.py index 6d71f31673..d28dcc4f6e 100644 --- a/django_spanner/schema.py +++ b/django_spanner/schema.py @@ -6,6 +6,7 @@ from django.db import NotSupportedError from django.db.backends.base.schema import BaseDatabaseSchemaEditor +from django_spanner._opentelemetry_tracing import trace_call class DatabaseSchemaEditor(BaseDatabaseSchemaEditor): @@ -119,7 +120,15 @@ def create_model(self, model): sql += " " + tablespace_sql # Prevent using [] as params, in the case a literal '%' is used in the # definition - self.execute(sql, params or None) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table) + } + with trace_call( + "CloudSpannerDjango.create_model", + self.connection, + trace_attributes, + ): + self.execute(sql, params or None) # Add any field index and index_together's (deferred as SQLite # _remake_table needs it) @@ -144,8 +153,25 @@ def delete_model(self, model): model, index=True, primary_key=False ) for index_name in index_names: - self.execute(self._delete_index_sql(model, index_name)) - super().delete_model(model) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "index_name": index_name, + } + with trace_call( + "CloudSpannerDjango.delete_model.delete_index", + self.connection, + trace_attributes, + ): + self.execute(self._delete_index_sql(model, index_name)) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table) + } + with trace_call( + "CloudSpannerDjango.delete_model", + self.connection, + trace_attributes, + ): + super().delete_model(model) def add_field(self, model, field): """ @@ -250,8 +276,28 @@ def remove_field(self, model, field): # column. index_names = self._constraint_names(model, [field.column], index=True) for index_name in index_names: - self.execute(self._delete_index_sql(model, index_name)) - super().remove_field(model, field) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "field": field.column, + "index_name": index_name, + } + with trace_call( + "CloudSpannerDjango.remove_field.delete_index", + self.connection, + trace_attributes, + ): + self.execute(self._delete_index_sql(model, index_name)) + + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "field": field.column, + } + with trace_call( + "CloudSpannerDjango.remove_field", + self.connection, + trace_attributes, + ): + super().remove_field(model, field) def column_sql( self, model, field, include_default=False, exclude_not_null=False @@ -320,7 +366,14 @@ def add_index(self, model, index): (field_name, " DESC" if order == "DESC" else "") for field_name, order in index.fields_orders ] - super().add_index(model, index) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "index": "|".join(index.fields), + } + with trace_call( + "CloudSpannerDjango.add_index", self.connection, trace_attributes, + ): + super().add_index(model, index) def quote_value(self, value): # A more complete implementation isn't currently required. @@ -355,20 +408,48 @@ def _alter_field( "index isn't yet supported." ) for index_name in index_names: - self.execute(self._delete_index_sql(model, index_name)) - super()._alter_field( - model, - old_field, - new_field, - old_type, - new_type, - old_db_params, - new_db_params, - strict=False, - ) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "alter_field": old_field.column, + "index_name": index_name, + } + with trace_call( + "CloudSpannerDjango.alter_field.delete_index", + self.connection, + trace_attributes, + ): + self.execute(self._delete_index_sql(model, index_name)) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "alter_field": old_field.column, + } + with trace_call( + "CloudSpannerDjango.alter_field", + self.connection, + trace_attributes, + ): + super()._alter_field( + model, + old_field, + new_field, + old_type, + new_type, + old_db_params, + new_db_params, + strict=False, + ) # Recreate the index that was dropped earlier. if nullability_changed and new_field.db_index: - self.execute(self._create_index_sql(model, [new_field])) + trace_attributes = { + "model_name": self.quote_name(model._meta.db_table), + "alter_field": new_field.column, + } + with trace_call( + "CloudSpannerDjango.alter_field.recreate_index", + self.connection, + trace_attributes, + ): + self.execute(self._create_index_sql(model, [new_field])) def _alter_column_type_sql(self, model, old_field, new_field, new_type): # Spanner needs to use sql_alter_column_not_null if the field is diff --git a/tests/unit/django_spanner/test_schema.py b/tests/unit/django_spanner/test_schema.py index 8c8ee35f0d..570bcc572a 100644 --- a/tests/unit/django_spanner/test_schema.py +++ b/tests/unit/django_spanner/test_schema.py @@ -5,12 +5,28 @@ # https://developers.google.com/open-source/licenses/bsd +from .models import Author +from django.db import NotSupportedError from django.db.models import Index from django.db.models.fields import IntegerField from django_spanner.schema import DatabaseSchemaEditor +from tests._helpers import HAS_OPENTELEMETRY_INSTALLED from tests.unit.django_spanner.simple_test import SpannerSimpleTestClass from unittest import mock -from .models import Author +from tests.unit.django_spanner.test__opentelemetry_tracing import ( + PROJECT, + INSTANCE_ID, + DATABASE_ID, +) + + +BASE_ATTRIBUTES = { + "db.type": "spanner", + "db.engine": "django_spanner", + "db.project": PROJECT, + "db.instance": INSTANCE_ID, + "db.name": DATABASE_ID, +} class TestUtils(SpannerSimpleTestClass): @@ -43,6 +59,14 @@ def test_create_model(self): + "PRIMARY KEY(id)", None, ) + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 1) + self.assertSpanAttributes( + "CloudSpannerDjango.create_model", + attributes=dict(BASE_ATTRIBUTES, model_name="tests_author"), + span=span_list[0], + ) def test_delete_model(self): """ @@ -57,6 +81,58 @@ def test_delete_model(self): "DROP TABLE tests_author", ) + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 1) + self.assertSpanAttributes( + "CloudSpannerDjango.delete_model", + attributes=dict(BASE_ATTRIBUTES, model_name="tests_author"), + span=span_list[0], + ) + + def test_delete_model_with_index(self): + """ + Tests deleting a model with index + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + + def delete_index_sql(*args, **kwargs): + # Overriding Statement creation with sql string. + return "DROP INDEX num_unique" + + def constraint_names(*args, **kwargs): + return ["num_unique"] + + schema_editor._delete_index_sql = delete_index_sql + schema_editor._constraint_names = constraint_names + schema_editor.delete_model(Author) + + calls = [ + mock.call("DROP INDEX num_unique"), + mock.call("DROP TABLE tests_author"), + ] + + schema_editor.execute.assert_has_calls(calls) + + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 2) + self.assertSpanAttributes( + "CloudSpannerDjango.delete_model.delete_index", + attributes=dict( + BASE_ATTRIBUTES, + model_name="tests_author", + index_name="num_unique", + ), + span=span_list[0], + ) + self.assertSpanAttributes( + "CloudSpannerDjango.delete_model", + attributes=dict(BASE_ATTRIBUTES, model_name="tests_author",), + span=span_list[1], + ) + def test_add_field(self): """ Tests adding fields to models @@ -71,6 +147,80 @@ def test_add_field(self): "ALTER TABLE tests_author ADD COLUMN age INT64", [] ) + def test_remove_field(self): + """ + Tests remove fields from models + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + schema_editor._constraint_names = mock.MagicMock() + remove_field = IntegerField(unique=True) + remove_field.set_attributes_from_name("num") + schema_editor.remove_field(Author, remove_field) + + schema_editor.execute.assert_called_once_with( + "ALTER TABLE tests_author DROP COLUMN num" + ) + + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 1) + self.assertSpanAttributes( + "CloudSpannerDjango.remove_field", + attributes=dict( + BASE_ATTRIBUTES, model_name="tests_author", field="num", + ), + span=span_list[0], + ) + + def test_remove_field_with_index(self): + """ + Tests remove fields from models + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + + def delete_index_sql(*args, **kwargs): + # Overriding Statement creation with sql string. + return "DROP INDEX num_unique" + + def constraint_names(*args, **kwargs): + return ["num_unique"] + + schema_editor._delete_index_sql = delete_index_sql + schema_editor._constraint_names = constraint_names + + remove_field = IntegerField(unique=True) + remove_field.set_attributes_from_name("num") + schema_editor.remove_field(Author, remove_field) + + calls = [ + mock.call("DROP INDEX num_unique"), + mock.call("ALTER TABLE tests_author DROP COLUMN num"), + ] + schema_editor.execute.assert_has_calls(calls) + + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 2) + self.assertSpanAttributes( + "CloudSpannerDjango.remove_field.delete_index", + attributes=dict( + BASE_ATTRIBUTES, + model_name="tests_author", + field="num", + index_name="num_unique", + ), + span=span_list[0], + ) + self.assertSpanAttributes( + "CloudSpannerDjango.remove_field", + attributes=dict( + BASE_ATTRIBUTES, model_name="tests_author", field="num", + ), + span=span_list[1], + ) + def test_column_sql_not_null_field(self): """ Tests column sql for not null field @@ -81,6 +231,7 @@ def test_column_sql_not_null_field(self): new_field.set_attributes_from_name("num") sql, params = schema_editor.column_sql(Author, new_field) self.assertEqual(sql, "INT64 NOT NULL") + self.assertEqual(params, []) def test_column_sql_nullable_field(self): """ @@ -92,6 +243,7 @@ def test_column_sql_nullable_field(self): new_field.set_attributes_from_name("num") sql, params = schema_editor.column_sql(Author, new_field) self.assertEqual(sql, "INT64") + self.assertEqual(params, []) def test_column_add_index(self): """ @@ -102,12 +254,22 @@ def test_column_add_index(self): index = Index(name="test_author_index_num", fields=["num"]) schema_editor.add_index(Author, index) name, args, kwargs = schema_editor.execute.mock_calls[0] - self.assertEqual( str(args[0]), "CREATE INDEX test_author_index_num ON tests_author (num)", ) self.assertEqual(kwargs["params"], None) + self.assertEqual(name, "") + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 1) + self.assertSpanAttributes( + "CloudSpannerDjango.add_index", + attributes=dict( + BASE_ATTRIBUTES, model_name="tests_author", index="num", + ), + span=span_list[0], + ) def test_alter_field(self): """ @@ -124,3 +286,112 @@ def test_alter_field(self): schema_editor.execute.assert_called_once_with( "ALTER TABLE tests_author RENAME COLUMN num TO author_num" ) + + def test_alter_field_change_null_with_single_index(self): + """ + Tests altering nullability of field with single index + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + + def delete_index_sql(*args, **kwargs): + # Overriding Statement creation with sql string. + return "DROP INDEX num_unique" + + def create_index_sql(*args, **kwargs): + # Overriding Statement creation with sql string. + return "CREATE INDEX tests_author ON tests_author (author_num)" + + def constraint_names(*args, **kwargs): + return ["num_unique"] + + schema_editor._delete_index_sql = delete_index_sql + schema_editor._create_index_sql = create_index_sql + schema_editor._constraint_names = constraint_names + old_field = IntegerField(null=True, db_index=True) + old_field.set_attributes_from_name("num") + new_field = IntegerField(db_index=True) + new_field.set_attributes_from_name("author_num") + schema_editor.alter_field(Author, old_field, new_field) + + calls = [ + mock.call("DROP INDEX num_unique"), + mock.call( + "ALTER TABLE tests_author RENAME COLUMN num TO author_num" + ), + mock.call( + "ALTER TABLE tests_author ALTER COLUMN author_num INT64 NOT NULL", + [], + ), + mock.call( + "CREATE INDEX tests_author ON tests_author (author_num)" + ), + ] + schema_editor.execute.assert_has_calls(calls) + if HAS_OPENTELEMETRY_INSTALLED: + span_list = self.ot_exporter.get_finished_spans() + self.assertEqual(len(span_list), 3) + self.assertSpanAttributes( + "CloudSpannerDjango.alter_field.delete_index", + attributes=dict( + BASE_ATTRIBUTES, + model_name="tests_author", + index_name="num_unique", + alter_field="num", + ), + span=span_list[0], + ) + self.assertSpanAttributes( + "CloudSpannerDjango.alter_field", + attributes=dict( + BASE_ATTRIBUTES, + model_name="tests_author", + alter_field="num", + ), + span=span_list[1], + ) + self.assertSpanAttributes( + "CloudSpannerDjango.alter_field.recreate_index", + attributes=dict( + BASE_ATTRIBUTES, + model_name="tests_author", + alter_field="author_num", + ), + span=span_list[2], + ) + + def test_alter_field_nullability_change_raise_not_support_error(self): + """ + Tests altering nullability of existing field in table + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + + def constraint_names(*args, **kwargs): + return ["num_unique"] + + schema_editor._constraint_names = constraint_names + old_field = IntegerField(null=True) + old_field.set_attributes_from_name("num") + new_field = IntegerField() + new_field.set_attributes_from_name("author_num") + with self.assertRaises(NotSupportedError): + schema_editor.alter_field(Author, old_field, new_field) + + def test_alter_field_change_null_with_multiple_index_error(self): + """ + Tests altering nullability of field with multiple index not supported + """ + with DatabaseSchemaEditor(self.connection) as schema_editor: + schema_editor.execute = mock.MagicMock() + + def constraint_names(*args, **kwargs): + return ["num_unique", "dummy_index"] + + schema_editor._constraint_names = constraint_names + old_field = IntegerField(null=True, db_index=True) + old_field.set_attributes_from_name("num") + new_field = IntegerField() + new_field.set_attributes_from_name("author_num") + with self.assertRaises(NotSupportedError): + schema_editor.alter_field(Author, old_field, new_field)