From 8a1b3eeab71f42a93d15ce6468bb37d6c027db41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E1=84=87=E1=85=A1=E1=86=A8=E1=84=80=E1=85=A7=E1=86=BC?= =?UTF-8?q?=E1=84=90=E1=85=A2?= Date: Sun, 18 May 2025 13:55:48 +0900 Subject: [PATCH] refactor: move ClickHouse connection check to runtime and validate env vars --- llm_utils/connect_db.py | 64 +++++++++++++++++++++++++++-------------- 1 file changed, 42 insertions(+), 22 deletions(-) diff --git a/llm_utils/connect_db.py b/llm_utils/connect_db.py index 447c7b2..b101a62 100644 --- a/llm_utils/connect_db.py +++ b/llm_utils/connect_db.py @@ -1,8 +1,9 @@ """ -이 모듈은 ClickHouse 데이터베이스에 연결하고 SQL 쿼리를 실행하여 결과를 pandas DataFrame으로 반환하는 기능을 제공합니다. +ClickHouse 데이터베이스에 연결하고 SQL 쿼리를 실행하여 결과를 pandas DataFrame으로 반환하는 모듈입니다. 구성 요소: -- 환경 변수에서 접속 정보를 불러와 ClickHouse 서버에 연결합니다. +- 환경 변수에서 접속 정보를 불러옵니다. +- ClickHouse에 지연 연결(lazy connection)을 수행합니다. - SQL 쿼리를 실행하고 결과를 pandas DataFrame으로 반환합니다. - 연결 실패 및 쿼리 오류에 대해 로깅을 통해 디버깅을 지원합니다. """ @@ -29,32 +30,51 @@ class ConnectDB: """ ClickHouse 데이터베이스에 연결하고 SQL 쿼리를 실행하는 클래스입니다. - 환경 변수에서 접속 정보를 로드하여 ClickHouse 서버에 연결하며, - SQL 쿼리 실행 결과를 pandas DataFrame으로 반환합니다. + 환경 변수에서 접속 정보를 읽어들이며, 실제 연결은 SQL 실행 시점에 수행됩니다. + 연결된 클라이언트를 통해 SQL 쿼리를 실행하고, 결과를 pandas DataFrame으로 반환합니다. """ def __init__(self): """ ConnectDB 클래스의 인스턴스를 초기화합니다. - 환경 변수에서 ClickHouse 접속 정보를 읽고, 즉시 서버에 연결을 시도합니다. + 환경 변수에서 ClickHouse 접속 정보를 불러오며, 연결은 즉시 수행하지 않습니다. """ self.client: Optional[Client] = None - self.host = os.getenv("CLICKHOUSE_HOST") - self.dbname = os.getenv("CLICKHOUSE_DATABASE") - self.user = os.getenv("CLICKHOUSE_USER") - self.password = os.getenv("CLICKHOUSE_PASSWORD") - self.port = os.getenv("CLICKHOUSE_PORT") + self.is_connected: bool = False - self.connect_to_clickhouse() + self.host = self._get_env_or_raise("CLICKHOUSE_HOST") + self.dbname = self._get_env_or_raise("CLICKHOUSE_DATABASE") + self.user = self._get_env_or_raise("CLICKHOUSE_USER") + self.password = os.getenv("CLICKHOUSE_PASSWORD", "") + self.port = int(self._get_env_or_raise("CLICKHOUSE_PORT")) + + def _get_env_or_raise(self, var_name: str) -> str: + """ + 주어진 환경변수를 읽고, 값이 없으면 예외를 발생시킵니다. + + Parameters: + var_name (str): 환경 변수 이름 + + Returns: + str: 환경 변수의 값 + + Raises: + ValueError: 값이 없을 경우 + """ + + value = os.getenv(var_name) + if not value: + logger.error("Environment variable '%s' is not set.", var_name) + raise ValueError(f"Environment variable '{var_name}' is not set.") + return value def connect_to_clickhouse(self) -> None: """ ClickHouse 서버에 연결을 시도합니다. - 연결에 성공하면 client 객체가 설정되며, 실패 시 예외를 발생시킵니다. - 연결 상태는 로깅을 통해 출력됩니다. + 연결 성공 시 client가 초기화되며, 실패 시 로그를 남기고 예외를 발생시킵니다. """ try: @@ -65,30 +85,31 @@ def connect_to_clickhouse(self) -> None: password=self.password, database=self.dbname, ) + self.is_connected = True logger.info("Successfully connected to ClickHouse.") except Exception as e: + self.is_connected = False logger.error("Failed to connect to ClickHouse: %s", e) raise def run_sql(self, sql: str) -> pd.DataFrame: """ - SQL 쿼리를 실행하고 결과를 pandas DataFrame으로 반환합니다. - 내부적으로 ClickHouse 클라이언트가 없으면 자동으로 재연결을 시도합니다. + 주어진 SQL 쿼리를 실행하고 결과를 pandas DataFrame으로 반환합니다. + + 연결이 설정되지 않은 경우 자동으로 연결을 시도합니다. Parameters: sql (str): 실행할 SQL 쿼리 문자열 Returns: - pd.DataFrame: 쿼리 실행 결과를 담은 DataFrame 객체 + pd.DataFrame: 쿼리 결과를 포함한 DataFrame Raises: - Exception: SQL 실행 중 오류 발생 시 예외를 발생시킵니다. + Exception: SQL 실행 중 오류가 발생한 경우 예외를 발생시킵니다. """ - if not self.client: - logger.warning( - "ClickHouse client is not initialized. Attempting to reconnect..." - ) + if not self.is_connected or not self.client: + logger.warning("ClickHouse client not connected. Attempting to connect...") self.connect_to_clickhouse() try: @@ -97,7 +118,6 @@ def run_sql(self, sql: str) -> pd.DataFrame: column_names = [col[0] for col in columns] df = pd.DataFrame(rows, columns=column_names) return df - except Exception as e: logger.exception("An error occurred while executing SQL: %s", e) raise