refactor: Improve GMS server configuration and fetcher management

ehddnr301 · ehddnr301 · commit cbe19d42058e · 2025-02-02T02:46:43.000Z
- Replace global fetcher with environment-based configuration
- Add get_fetcher() method to dynamically retrieve DatahubMetadataFetcher
- Update CLI option name from gms_server to datahub_server
- Modify tools module to use environment variable for server configuration
diff --git a/cli/__init__.py b/cli/__init__.py
@@ -7,13 +7,13 @@
 @click.version_option(version="0.1.4")
 @click.pass_context
 @click.option(
-    "--gms-server", default="http://localhost:8080", help="Datahub GMS 서버 URL"
+    "--datahub_server", default="http://localhost:8080", help="Datahub GMS 서버 URL"
 )
 @click.option("--run-streamlit", is_flag=True, help="Run the Streamlit app.")
 @click.option("-p", "--port", type=int, default=8501, help="Streamlit port")
-def cli(ctx, gms_server, run_streamlit, port):
+def cli(ctx, datahub_server, run_streamlit, port):
     try:
-        set_gms_server(gms_server)
+        set_gms_server(datahub_server)
     except ValueError as e:
         click.echo(str(e))
         ctx.exit(1)
diff --git a/llm_utils/tools.py b/llm_utils/tools.py
@@ -1,25 +1,31 @@
+import os
 from typing import List, Dict
 
 from langchain.tools import tool
 
 from data_utils.datahub_source import DatahubMetadataFetcher
 
-fetcher = None
-
 
 def set_gms_server(gms_server: str):
-    global fetcher
-
     try:
+        os.environ["DATAHUB_SERVER"] = gms_server
         fetcher = DatahubMetadataFetcher(gms_server=gms_server)
     except ValueError as e:
         # 유효하지 않은 GMS 서버 주소일 경우 예외를 발생시킴
         raise ValueError(f"GMS 서버 설정 실패: {str(e)}")
 
 
+def get_fetcher():
+    gms_server = os.getenv("DATAHUB_SERVER")
+    if not gms_server:
+        raise ValueError("GMS 서버가 설정되지 않았습니다.")
+    return DatahubMetadataFetcher(gms_server=gms_server)
+
+
 @tool
 def get_table_info() -> Dict[str, str]:
     """전체 테이블 이름과 설명을 가져오는 함수"""
+    fetcher = get_fetcher()
     urns = fetcher.get_urns()
     table_info = {}
     for urn in urns:
@@ -33,6 +39,7 @@ def get_table_info() -> Dict[str, str]:
 @tool
 def get_column_info(table_name: str) -> List[Dict[str, str]]:
     """table_name에 해당하는 컬럼 이름과 설명을 가져오는 함수"""
+    fetcher = get_fetcher()
     urns = fetcher.get_urns()
     for urn in urns:
         if fetcher.get_table_name(urn) == table_name:
