fix auth token lookup logic and add request timeouts

Author: RinZ27

x2text-service/app/authentication_middleware.py

@@ -8,6 +8,8 @@
 
 
 def authentication_middleware(func: Any) -> Any:
+    """Decorator to enforce bearer token authentication on flask routes."""
+
     def wrapper(*args: Any, **kwargs: Any) -> Any:
         token = AuthenticationMiddleware.get_token_from_auth_header(request)
         # Check if bearer token exists and validate it
@@ -23,32 +25,26 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 class AuthenticationMiddleware:
     @classmethod
     def validate_bearer_token(cls, token: str | None) -> bool:
+        """Validate the provided bearer token against the database."""
         try:
             if token is None:
                 current_app.logger.error("Authentication failed. Empty bearer token")
                 return False
             platform_key_table = f'"{Env.DB_SCHEMA}".{DBTable.PLATFORM_KEY}'
-            query = f"SELECT * FROM {platform_key_table} WHERE key = '{token}'"
-            cursor = be_db.execute_sql(query)
+            query = f"SELECT * FROM {platform_key_table} WHERE key = %s"
+            cursor = be_db.execute_sql(query, (token,))
             result_row = cursor.fetchone()
             cursor.close()
             if not result_row or len(result_row) == 0:
-                current_app.logger.error(
-                    f"Authentication failed. bearer token not found {token}"
-                )
+                current_app.logger.error("Authentication failed. bearer token not found")
                 return False
             platform_key = str(result_row[1])
             is_active = bool(result_row[2])
             if not is_active:
-                current_app.logger.error(
-                    f"Token is not active. Activate \
-                        before using it. token {token}"
-                )
+                current_app.logger.error("Token is not active. Activate before using it.")
                 return False
             if platform_key != token:
-                current_app.logger.error(
-                    f"Authentication failed. Invalid bearer token: {token}"
-                )
+                current_app.logger.error("Authentication failed. Invalid bearer token")
                 return False
 
         except Exception as e:
@@ -62,6 +58,7 @@ def validate_bearer_token(cls, token: str | None) -> bool:
 
     @classmethod
     def get_token_from_auth_header(cls, request: Request) -> str | None:
+        """Extract the bearer token from the Authorization header."""
         try:
             bearer_token = request.headers.get("Authorization")
             if not bearer_token:
@@ -99,6 +96,7 @@ def get_organization_from_bearer_token(cls, token: str) -> tuple[int | None, str
 
     @classmethod
     def execute_query(cls, query: str, params: tuple = ()) -> Any:
+        """Execute a SQL query and return the first result."""
         cursor = be_db.execute_sql(query, params)
         result_row = cursor.fetchone()
         cursor.close()

x2text-service/app/controllers/controller.py

@@ -26,13 +26,15 @@
 
 @basic.route("/health", methods=["GET"])
 def health() -> str:
+    """Check the health status of the service."""
     logging.info("Checking health from : %s", request.remote_addr)
     return "OK"
 
 
 @basic.route("/test-connection", methods=["POST"])
 @authentication_middleware
 def test_connection() -> Any:
+    """Test the connection to the Unstructured API."""
     logging.info("Received a test connection request from %s", request.remote_addr)
     form_data = dict(request.form)
     unstructured_api_key = X2TextUtil.get_value_for_key(UNSTRUCTURED_API_KEY, form_data)
@@ -54,7 +56,7 @@ def test_connection() -> Any:
             headers=headers,
             data=None,
             files=files,
-            timeout=None,
+            timeout=60,
         )
 
         if response.status_code == 400:
@@ -76,6 +78,7 @@ def test_connection() -> Any:
 @basic.route("/process", methods=["POST"])
 @authentication_middleware
 def process() -> Any:
+    """Process a document for text extraction."""
     logging.info("Received a doc processing request from %s", request.remote_addr)
     form_data = dict(request.form)
     url = X2TextUtil.get_value_for_key(UNSTRUCTURED_URL, form_data)
@@ -116,14 +119,21 @@ def process() -> Any:
     }
     payload = form_data
 
-    response = requests.request(
-        "POST",
-        url,
-        headers=headers,
-        data=payload,
-        files=files,
-        timeout=None,
-    )
+    try:
+        response = requests.request(
+            "POST",
+            url,
+            headers=headers,
+            data=payload,
+            files=files,
+            timeout=60,
+        )
+    except requests.exceptions.RequestException as e:
+        logging.error("Text extraction request failed: %s", e)
+        x2_text_audit.status = "Failed"
+        x2_text_audit.save()
+        return {"message": "Text extraction request failed"}, 502
+
     if response.ok:
         json_response = response.json()
         response_text = X2TextUtil.get_text_content(json_response)