Repository: googleapis/python-spanner-django

Issue:
feat: Add support for open telemetry
Add support for open telemetry tracing if open telemetry is installed.
Added the same as extra dependency.

Also added unit tests for the same.

Follow up PR to add tracing in schema with related checks in unit tests. https://github.com/googleapis/python-spanner-django/pull/648

Fixes #633 🦕


Discussion:

Resolving patch:
diff --git a/django_spanner/schema.py b/django_spanner/schema.py
--- 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


Test patch:
diff --git a/tests/unit/django_spanner/test_schema.py b/tests/unit/django_spanner/test_schema.py
--- 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)