feat: Implement query fetching in DatahubMetadataFetcher

Author: ehddnr301

data_utils/datahub_source.py

@@ -6,7 +6,11 @@
 from datahub.metadata.schema_classes import UpstreamLineageClass
 from collections import defaultdict
 import requests
-from data_utils.queries import ROOT_GLOSSARY_NODES_QUERY, GLOSSARY_NODE_QUERY
+from data_utils.queries import (
+    ROOT_GLOSSARY_NODES_QUERY,
+    GLOSSARY_NODE_QUERY,
+    LIST_QUERIES_QUERY,
+)
 
 
 class DatahubMetadataFetcher:
@@ -480,3 +484,101 @@ def get_glossary_data(self):
                 return {"error": True, "message": f"결과 구조 파싱 중 오류 발생: {e}"}
         else:
             return {"error": True, "message": "용어집 노드를 가져오지 못했습니다."}
+
+    def get_queries(self, start=0, count=10, query="*", filters=None):
+        """
+        DataHub에서 쿼리 목록을 가져오는 함수
+
+        Args:
+            start (int): 시작 인덱스 (기본값=0)
+            count (int): 반환할 쿼리 수 (기본값=10)
+            query (str): 필터링에 사용할 쿼리 문자열 (기본값="*")
+            filters (list): 추가 필터 (기본값=None)
+
+        Returns:
+            dict: 쿼리 목록 정보
+        """
+        # GraphQL 요청용 입력 변수 준비
+        input_params = {"start": start, "count": count, "query": query}
+
+        if filters:
+            input_params["filters"] = filters
+
+        variables = {"input": input_params}
+
+        # GraphQL 요청 보내기
+        headers = {"Content-Type": "application/json"}
+        response = requests.post(
+            f"{self.gms_server}/api/graphql",
+            json={"query": LIST_QUERIES_QUERY, "variables": variables},
+            headers=headers,
+        )
+        print(response.json())
+
+        # 결과 반환
+        if response.status_code == 200:
+            return response.json()
+        else:
+            return {
+                "error": True,
+                "status_code": response.status_code,
+                "message": response.text,
+            }
+
+    def process_queries(self, result):
+        """
+        쿼리 목록 결과를 처리하고 간소화된 형태로 반환하는 함수
+
+        Args:
+            result (dict): API 응답 결과
+
+        Returns:
+            dict: 처리된 쿼리 목록 데이터 (urn, name, description, statement만 포함)
+        """
+        if "error" in result:
+            return result
+
+        processed_result = {"total_queries": 0, "count": 0, "start": 0, "queries": []}
+
+        if "data" in result and "listQueries" in result["data"]:
+            list_queries = result["data"]["listQueries"]
+            processed_result["total_queries"] = list_queries.get("total", 0)
+            processed_result["count"] = list_queries.get("count", 0)
+            processed_result["start"] = list_queries.get("start", 0)
+
+            for query in list_queries.get("queries", []):
+                query_info = {"urn": query.get("urn")}
+
+                props = query.get("properties", {})
+                query_info["name"] = props.get("name")
+                query_info["description"] = props.get("description")
+                query_info["statement"] = props.get("statement", {}).get("value")
+
+                processed_result["queries"].append(query_info)
+
+        return processed_result
+
+    def get_query_data(self, start=0, count=10, query="*", filters=None):
+        """
+        DataHub에서 쿼리 목록을 가져와 처리하는 함수
+
+        Args:
+            start (int): 시작 인덱스 (기본값=0)
+            count (int): 반환할 쿼리 수 (기본값=10)
+            query (str): 필터링에 사용할 쿼리 문자열 (기본값="*")
+            filters (list): 추가 필터 (기본값=None)
+
+        Returns:
+            dict: 처리된 쿼리 목록 데이터
+        """
+        # DataHub 서버에 연결하여 쿼리 목록 가져오기
+        result = self.get_queries(start, count, query, filters)
+
+        # 결과 처리
+        if result:
+            try:
+                return self.process_queries(result)
+            except KeyError as e:
+                return {"error": True, "message": f"결과 구조 파싱 중 오류 발생: {e}"}
+        else:
+            return {"error": True, "message": "쿼리 목록을 가져오지 못했습니다."}