feat: Enhance CLI and DataHub metadata fetching with GMS server configuration

- Add GMS server URL configuration option to CLI
- Implement GMS server URL validation in DatahubMetadataFetcher
- Create dynamic GMS server setup in tools module
- Update package version to 0.1.4

Author: ehddnr301

cli/__init__.py

@@ -1,22 +1,35 @@
 import click
 import subprocess
+from llm_utils.tools import set_gms_server
 
 
 @click.group()
-@click.version_option(version="2.4.2")
+@click.version_option(version="0.1.4")
 @click.pass_context
-def cli(ctx):
-    pass
-
-
-@click.command()
+@click.option(
+    "--gms-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 run_streamlit(run_streamlit, port):
+def cli(ctx, gms_server, run_streamlit, port):
+    try:
+        set_gms_server(gms_server)
+    except ValueError as e:
+        click.echo(str(e))
+        ctx.exit(1)
+    if run_streamlit:
+        run_streamlit_command(port)
+
+
+def run_streamlit_command(port):
     """Run the Streamlit app."""
     subprocess.run(
         ["streamlit", "run", "interface/streamlit_app.py", "--server.port", str(port)]
     )
 
 
-cli.add_command(run_streamlit)
+@cli.command()
+@click.option("-p", "--port", type=int, default=8501, help="Streamlit port")
+def run_streamlit(port):
+    """Run the Streamlit app."""
+    run_streamlit_command(port)

data_utils/datahub_source.py

@@ -2,15 +2,34 @@
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.metadata.schema_classes import DatasetPropertiesClass, SchemaMetadataClass
 from datahub.emitter.rest_emitter import DatahubRestEmitter
+import requests
 
 
 class DatahubMetadataFetcher:
     def __init__(self, gms_server="http://localhost:8080", extra_headers={}):
+        # gms_server 주소 유효성 검사
+        if not self._is_valid_gms_server(gms_server):
+            raise ValueError(f"유효하지 않은 GMS 서버 주소: {gms_server}")
+
         self.emitter = DatahubRestEmitter(
             gms_server=gms_server, extra_headers=extra_headers
         )
         self.datahub_graph = self.emitter.to_graph()
 
+    def _is_valid_gms_server(self, gms_server):
+        # GMS 서버 주소의 유효성을 검사하는 로직 추가
+        # GraphQL 요청을 사용하여 서버 상태 확인
+        query = {"query": "{ health { status } }"}
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            response = requests.post(
+                f"{gms_server}/api/graphql", json=query, headers=headers
+            )
+            return response.status_code == 200
+        except requests.exceptions.RequestException:
+            return False
+
     def get_urns(self):
         # 필터를 적용하여 데이터셋의 URN 가져오기
         return self.datahub_graph.get_urns_by_filter()

llm_utils/tools.py

@@ -4,7 +4,17 @@
 
 from data_utils.datahub_source import DatahubMetadataFetcher
 
-fetcher = DatahubMetadataFetcher()
+fetcher = None
+
+
+def set_gms_server(gms_server: str):
+    global fetcher
+
+    try:
+        fetcher = DatahubMetadataFetcher(gms_server=gms_server)
+    except ValueError as e:
+        # 유효하지 않은 GMS 서버 주소일 경우 예외를 발생시킴
+        raise ValueError(f"GMS 서버 설정 실패: {str(e)}")
 
 
 @tool

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="lang2sql",  # 패키지 이름
-    version="0.1.2",  # 버전
+    version="0.1.4",  # 버전
     description="Lang2SQL - Query Generator for Data Warehouse",
     long_description=long_description,
     long_description_content_type="text/markdown",