From f4cf25cf45048e61e3eed1e3ca26fa759d51c7e2 Mon Sep 17 00:00:00 2001
From: Will Sheldon <wshel@netflix.com>
Date: Wed, 21 May 2025 10:45:51 -0700
Subject: [PATCH 1/2] fix(db): sqlalchemy 2.0 _entities attr

---
 src/dispatch/database/service.py         | 52 ++++++++++++++++--------
 src/dispatch/search/fulltext/__init__.py |  8 +++-
 2 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/src/dispatch/database/service.py b/src/dispatch/database/service.py
index dab8d1d9127f..ed127c38ebe1 100644
--- a/src/dispatch/database/service.py
+++ b/src/dispatch/database/service.py
@@ -268,23 +268,43 @@ def get_query_models(query):
     """
     models = [col_desc["entity"] for col_desc in query.column_descriptions]
 
-    # account joined entities
+    # In SQLAlchemy 2.x, we need to use a different approach to get joined entities
     try:
-        models.extend(mapper.class_ for mapper in query._compile_state()._join_entities)
-    except InvalidRequestError:
-        # query might not contain columns yet, hence cannot be compiled
-        # try to infer the models from various internals
-        for table_tuple in query._setup_joins + query._legacy_setup_joins:
-            model_class = get_model_from_table(table_tuple[0])
-            if model_class:
-                models.append(model_class)
-
-    # account also query.select_from entities
-    model_class = None
-    if query._from_obj:
-        model_class = get_model_from_table(query._from_obj[0])
-    if model_class and (model_class not in models):
-        models.append(model_class)
+        # Try to get the statement from the query
+        stmt = query.statement
+
+        # Extract entities from the statement's froms
+        for from_obj in stmt.froms:
+            if hasattr(from_obj, "entity"):
+                # For select statements with an entity
+                if from_obj.entity not in models:
+                    models.append(from_obj.entity)
+            elif hasattr(from_obj, "left") and hasattr(from_obj, "right"):
+                # For join objects
+                for side in [from_obj.left, from_obj.right]:
+                    if hasattr(side, "entity"):
+                        if side.entity not in models:
+                            models.append(side.entity)
+                    elif hasattr(side, "table"):
+                        model_class = get_model_from_table(side.table)
+                        if model_class and model_class not in models:
+                            models.append(model_class)
+
+        # Try to extract joined entities from the query's _join_entities
+        # This is for SQLAlchemy 2.x's internal structure
+        if hasattr(query, "_compile_state"):
+            try:
+                compile_state = query._compile_state()
+                if hasattr(compile_state, "_join_entities"):
+                    for mapper in compile_state._join_entities:
+                        if hasattr(mapper, "class_"):
+                            if mapper.class_ not in models:
+                                models.append(mapper.class_)
+            except Exception:
+                pass
+    except (AttributeError, InvalidRequestError):
+        # If we can't get the statement or process it, fall back to simpler approach
+        pass
 
     return {model.__name__: model for model in models}
 
diff --git a/src/dispatch/search/fulltext/__init__.py b/src/dispatch/search/fulltext/__init__.py
index 322c72b06844..bea460c2864d 100644
--- a/src/dispatch/search/fulltext/__init__.py
+++ b/src/dispatch/search/fulltext/__init__.py
@@ -52,7 +52,13 @@ def search(query, search_query, vector=None, regconfig=None, sort=False):
         return query
 
     if vector is None:
-        entity = query._entities[0].entity_zero.class_
+        # Get the entity class from the query in a SQLAlchemy 2.x compatible way
+        try:
+            # For SQLAlchemy 2.x
+            entity = query.column_descriptions[0]['entity']
+        except (AttributeError, IndexError, KeyError):
+            raise ValueError("Could not determine entity class from query. Please provide vector explicitly.")
+
         search_vectors = inspect_search_vectors(entity)
         vector = search_vectors[0]
 

From ee07c99e311cec44726a06b728c78cb461f4e308 Mon Sep 17 00:00:00 2001
From: Will Sheldon <wshel@netflix.com>
Date: Wed, 21 May 2025 10:53:01 -0700
Subject: [PATCH 2/2] lint(ruff): raise none

---
 src/dispatch/search/fulltext/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dispatch/search/fulltext/__init__.py b/src/dispatch/search/fulltext/__init__.py
index bea460c2864d..4619b5d89a37 100644
--- a/src/dispatch/search/fulltext/__init__.py
+++ b/src/dispatch/search/fulltext/__init__.py
@@ -57,7 +57,7 @@ def search(query, search_query, vector=None, regconfig=None, sort=False):
             # For SQLAlchemy 2.x
             entity = query.column_descriptions[0]['entity']
         except (AttributeError, IndexError, KeyError):
-            raise ValueError("Could not determine entity class from query. Please provide vector explicitly.")
+            raise ValueError("Could not determine entity class from query. Please provide vector explicitly.") from None
 
         search_vectors = inspect_search_vectors(entity)
         vector = search_vectors[0]