fix(query): Make query serialization thread-safe

[snigdhas]

Serializing multiple queries in multiple threads is not thread safe. The PRINTER global is reused by the Query serializer across these requests and maintains a local Translator in the translator classvar.

This would be fine if the translator were never mutated. However, the visit function temporarily mutates the translator using entity_aliases().

If multiple queries are being serialized concurrently, there are bad inverleavings of this code that result in invalid queries.

This can be reproduced with this diff and some long running queries issued in a batch (example):

diff --git a/snuba_sdk/query.py b/snuba_sdk/query.py
index 35d349a..eea1835 100644
--- a/snuba_sdk/query.py
+++ b/snuba_sdk/query.py
@@ -175,7 +175,9 @@ class Query(BaseQuery):
     def serialize(self) -> str:
         self.validate()
         optimized = self._optimize()
+        print(f"query visiting #{id(self)=}", flush=True)
         return PRINTER.visit(optimized)
+        print(f"query done visiting #{id(self)=}", flush=True)
 
     def print(self) -> str:
         self.validate()
diff --git a/snuba_sdk/query_visitors.py b/snuba_sdk/query_visitors.py
index 3bd8a26..e5f9922 100644
--- a/snuba_sdk/query_visitors.py
+++ b/snuba_sdk/query_visitors.py
@@ -138,7 +138,7 @@ class Printer(QueryVisitor[str]):
 
     def visit(self, query: main.Query) -> str:
         if isinstance(query.match, Join):
-            with entity_aliases(self.translator):
+            with entity_aliases(self.translator, id(query)):
                 return super().visit(query)
 
         return super().visit(query)
diff --git a/snuba_sdk/visitors.py b/snuba_sdk/visitors.py
index c787731..0035e54 100644
--- a/snuba_sdk/visitors.py
+++ b/snuba_sdk/visitors.py
@@ -329,10 +329,12 @@ class Translation(ExpressionVisitor[str]):
 
 
 @contextmanager
-def entity_aliases(translator: Translation) -> Generator[None, None, None]:
+def entity_aliases(translator: Translation, id) -> Generator[None, None, None]:
+    print(f"translator visiting #{id=}", flush=True)
     translator.use_entity_aliases = True
     yield
     translator.use_entity_aliases = False
+    print(f"done visiting #{id=}", flush=True)
 
 
 class ExpressionFinder(ExpressionVisitor[Set[Expression]]):

When we're lucky, the output shows that the translator's use_entity_aliases field is not mutated by multiple threads. When we're unlucky, it is.

Good

query visiting #id(self)=6155267360
translator visiting #id=6155267360
done visiting #id=6155267360
query visiting #id(self)=6155270864
translator visiting #id=6155270864
done visiting #id=6155270864

Bad

query visiting #id(self)=6155296096
translator visiting #id=6155296096
query visiting #id(self)=6155292592
translator visiting #id=6155292592
done visiting #id=6155296096
done visiting #id=6155292592

[armenzg]

Nice!