From cfb785006d61acf9a039c7b9b0cabf18af29bd32 Mon Sep 17 00:00:00 2001
From: lrangine <19699092+lokeshrangineni@users.noreply.github.com>
Date: Mon, 21 Oct 2024 15:58:43 -0400
Subject: [PATCH] Printing more verbose logs when we start the offline server
 to address the github issue https://github.com/feast-dev/feast/issues/4639

Signed-off-by: lrangine <19699092+lokeshrangineni@users.noreply.github.com>
---
 sdk/python/feast/offline_server.py | 32 ++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/sdk/python/feast/offline_server.py b/sdk/python/feast/offline_server.py
index 0cb40ad934c..cec043129e7 100644
--- a/sdk/python/feast/offline_server.py
+++ b/sdk/python/feast/offline_server.py
@@ -1,10 +1,13 @@
 import ast
 import json
 import logging
+import os
+import sys
 import traceback
 from datetime import datetime
 from typing import Any, Dict, List, cast
 
+import click
 import pyarrow as pa
 import pyarrow.flight as fl
 from google.protobuf.json_format import Parse
@@ -503,6 +506,24 @@ def get_table_column_names_and_types_from_data_source(self, command: dict):
         )
         return pa.table({"name": column_names, "type": types})
 
+    def serve(self):
+        message = "offline server starting with pid: "
+        logger.info(
+            message + "[%d]",
+            os.getpid(),
+            extra={"color_message": message + "[" + click.style("%d", fg="cyan") + "]"},
+        )
+        super().serve()
+
+    def shutdown(self):
+        message = "Sending a shutdown signal to the offline server running with pid:: "
+        logger.info(
+            message + "[%d]",
+            os.getpid(),
+            extra={"color_message": message + "[" + click.style("%d", fg="cyan") + "]"},
+        )
+        super().shutdown()
+
 
 def remove_dummies(fv: FeatureView) -> FeatureView:
     """
@@ -533,5 +554,12 @@ def start_server(
 
     location = "grpc+tcp://{}:{}".format(host, port)
     server = OfflineServer(store, location)
-    logger.info(f"Offline store server serving on {location}")
-    server.serve()
+    try:
+        logger.info(f"Offline store server serving at: {location}")
+        server.serve()
+    except KeyboardInterrupt:
+        logger.info("KeyboardInterrupt received, stopping the offline server.")
+    finally:
+        server.shutdown()
+        logger.info("offline server stopped.")
+        sys.exit(0)