feat: instrument repositories

* add SQLAlchemyInstrumentor
* load tagging details and identifier eagerly when getting TaggingResult from DB
* increase batch_size for similarity_pairs

Change-Id: I4720c109dee14c0d019f1dbba422d6c5cdf0a60b

Author: AndreyMarkinPPC

libs/media_similarity/media_similarity/media_similarity_service.py

@@ -576,7 +576,7 @@ def _calculate_cluster_assignments(
       if final_clusters.get(media_1) == final_clusters.get(media_2):
         trimmed_similar_media.add(similar_medium)
 
-    trim_ratio = ((1 - len(trimmed_similar_media) / len(similar_media)),)
+    trim_ratio = 1 - len(trimmed_similar_media) / len(similar_media)
     logger.info('trimmed graph edges by %.2f', trim_ratio)
     span.set_attribute('trim_ratio', trim_ratio)
   return ClusteringResults(

libs/media_similarity/media_similarity/repositories.py

@@ -22,12 +22,14 @@
 
 import sqlalchemy
 from media_tagging.repositories import SqlAlchemyRepository
+from opentelemetry import trace
 from sqlalchemy.orm import declarative_base
 from typing_extensions import override
 
 from media_similarity import media_pair
+from media_similarity.telemetry import tracer
 
-DEFAULT_CHUNK_SIZE: Final[int] = 100
+DEFAULT_CHUNK_SIZE: Final[int] = 10_000
 
 
 def _batched(iterable: Iterable[Any], chunk_size: int):
@@ -39,6 +41,7 @@ def _batched(iterable: Iterable[Any], chunk_size: int):
 class BaseSimilarityPairsRepository(abc.ABC):
   """Interface for defining repositories."""
 
+  @tracer.start_as_current_span('get_pairs')
   def get(
     self, pairs: str | Sequence[str], tagger: str | None = None
   ) -> list[media_pair.SimilarityPair]:
@@ -55,14 +58,17 @@ def get(
       return list(itertools.chain.from_iterable(results))
     return self._get(pairs, tagger)
 
+  @tracer.start_as_current_span('add_pairs')
   def add(
     self,
     pairs: media_pair.SimilarityPair | Sequence[media_pair.SimilarityPair],
   ) -> None:
     """Specifies add operations."""
+    span = trace.get_current_span()
     if not isinstance(pairs, MutableSequence):
       pairs = [pairs]
     self._add(pairs)
+    span.set_attribute('pairs_added', len(pairs))
 
   @abc.abstractmethod
   def _get(

libs/media_tagging/media_tagging/repositories.py

@@ -23,7 +23,8 @@
 from collections.abc import Sequence
 
 import sqlalchemy
-from sqlalchemy.orm import declarative_base, relationship
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from sqlalchemy.orm import declarative_base, joinedload, relationship
 from sqlalchemy.pool import StaticPool
 from typing_extensions import override
 
@@ -188,6 +189,7 @@ def engine(self) -> sqlalchemy.engine.Engine:
       )
     else:
       self._engine = sqlalchemy.create_engine(self.db_url)
+    SQLAlchemyInstrumentor().instrument(engine=self._engine)
     return self._engine
 
 
@@ -242,7 +244,12 @@ def get(
           TaggingResults.tagging_details_id == tagging_details_hash
         )
 
-      if not (results := query.all()):
+      if not (
+        results := query.options(
+          joinedload(TaggingResults.identifier),
+          joinedload(TaggingResults.tagging_details),
+        ).all()
+      ):
         return []
       tagging_results = []
       for res in results:

libs/media_tagging/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
   "opentelemetry-api",
   "opentelemetry-sdk",
   "opentelemetry-exporter-otlp",
+  "opentelemetry-instrumentation-sqlalchemy",
   "google-api-python-client",
 ]
 authors = [