From 6136a13a114c3f2aa7d59390e3d1254af9dea565 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Wed, 8 Oct 2025 14:15:17 +0530
Subject: [PATCH] fix: update creation date in migration file and enhance
 exception handling in Elasticsearch connector

---
 .../22_add_elasticsearch_connector_enum.py    |  2 +-
 .../app/connectors/elasticsearch_connector.py | 44 +++++++++++++------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/surfsense_backend/alembic/versions/22_add_elasticsearch_connector_enum.py b/surfsense_backend/alembic/versions/22_add_elasticsearch_connector_enum.py
index 6ab4daa54..ad13af1f8 100644
--- a/surfsense_backend/alembic/versions/22_add_elasticsearch_connector_enum.py
+++ b/surfsense_backend/alembic/versions/22_add_elasticsearch_connector_enum.py
@@ -2,7 +2,7 @@
 
 Revision ID: 22
 Revises: 21
-Create Date: 2025-10-4 12:00:00.000000
+Create Date: 2025-10-08 12:00:00.000000
 
 """
 
diff --git a/surfsense_backend/app/connectors/elasticsearch_connector.py b/surfsense_backend/app/connectors/elasticsearch_connector.py
index 8a8c259e1..0ef536d69 100644
--- a/surfsense_backend/app/connectors/elasticsearch_connector.py
+++ b/surfsense_backend/app/connectors/elasticsearch_connector.py
@@ -6,6 +6,19 @@ from typing import Any
 from elasticsearch import AsyncElasticsearch
 from elasticsearch.exceptions import AuthenticationException, ConnectionError
 
+try:
+    from elastic_transport import (
+        ApiError,
+        ConnectionError as TransportConnectionError,
+        TlsError,
+    )
+except ImportError:
+    # Fallback for older versions
+    TransportConnectionError = ConnectionError
+    TlsError = Exception
+    ApiError = Exception
+
+
 from app.db import DocumentType
 
 logger = logging.getLogger(__name__)
@@ -74,9 +87,13 @@ class ElasticsearchConnector:
             )
             return True
 
-        except (ConnectionError, AuthenticationException) as e:
+        # Updated exception handling for better version compatibility
+        except (ConnectionError, TransportConnectionError, TlsError, ApiError) as e:
             logger.error(f"Failed to connect to Elasticsearch: {e}")
             return False
+        except AuthenticationException as e:
+            logger.error(f"Authentication failed for Elasticsearch: {e}")
+            return False
         except Exception as e:
             logger.error(f"Unexpected error connecting to Elasticsearch: {e}")
             return False
@@ -140,23 +157,24 @@ class ElasticsearchConnector:
                 return
 
             # Build search query
-            search_body = {
-                "query": {
-                    "multi_match": {
-                        "query": query,
-                        "fields": fields or ["*"],
-                        "type": "best_fields",
-                        "fuzziness": "AUTO",
-                    }
-                },
-                "size": size,
-                "_source": True,
+            search_query = {
+                "multi_match": {
+                    "query": query,
+                    "fields": fields or ["*"],
+                    "type": "best_fields",
+                    "fuzziness": "AUTO",
+                }
             }
 
             # Search across specified indices or all indices
             index_pattern = ",".join(indices) if indices else "*"
 
-            response = await self.client.search(index=index_pattern, body=search_body)
+            response = await self.client.search(
+                index=index_pattern,
+                query=search_query,
+                size=size,
+                source=True,
+            )
 
             for hit in response["hits"]["hits"]:
                 try: