add database utilities

Author: Yoonjin Hwang

.env.example

@@ -66,8 +66,69 @@ AWS_BEDROCK_EMBEDDING_MODEL=amazon.titan-embed-text-v2:0
 # HUGGING_FACE_EMBEDDING_API_TOKEN=
 
 DATAHUB_SERVER = 'http://-.-.-.-:-'
-CLICKHOUSE_HOST = '-.-.-.-'
-CLICKHOUSE_DATABASE = 'main'
-CLICKHOUSE_USER = '-'
-CLICKHOUSE_PASSWORD = '-'
-CLICKHOUSE_PORT = 9000
+
+
+###############################################
+######## Database Connector SELECTION #########
+###############################################
+
+# clickhouse
+# DB_TYPE=clickhouse
+# CLICKHOUSE_HOST=_._._._
+# CLICKHOUSE_PORT=9000
+# CLICKHOUSE_USER=_
+# CLICKHOUSE_PASSWORD=_
+# CLICKHOUSE_DATABASE=_
+
+# databricks
+# DB_TYPE=databricks
+# DATABRICKS_HOST=_
+# DATABRICKS_HTTP_PATH=_
+# DATABRICKS_ACCESS_TOKEN=_
+
+# duckdb
+# DB_TYPE=duckdb
+# DUCKDB_PATH=./data/duckdb.db
+
+# mariadb
+# DB_TYPE=mariadb
+# MARIADB_HOST=_
+# MARIADB_PORT=3306
+# MARIADB_USER=_
+# MARIADB_PASSWORD=_
+# MARIADB_DATABASE=_
+
+# mysql
+# DB_TYPE=mysql
+# MYSQL_HOST=_
+# MYSQL_PORT=3306
+# MYSQL_USER=_
+# MYSQL_PASSWORD=_
+# MYSQL_DATABASE=_
+
+# oracle
+# DB_TYPE=oracle
+# ORACLE_HOST=_
+# ORACLE_PORT=1521
+# ORACLE_USER=_
+# ORACLE_PASSWORD=_
+# ORACLE_DATABASE=_
+# ORACLE_SERVICE_NAME=_
+
+# postgres
+# DB_TYPE=postgres
+# POSTGRES_HOST=_
+# POSTGRES_PORT=5432
+# POSTGRES_USER=_
+# POSTGRES_PASSWORD=_
+# POSTGRES_DATABASE=_
+
+# snowflake
+# DB_TYPE=snowflake
+# SNOWFLAKE_USER=_
+# SNOWFLAKE_PASSWORD=_
+# SNOWFLAKE_ACCOUNT=_
+
+# sqlite
+# DB_TYPE=sqlite
+# SQLITE_PATH=./data/sqlite.db

db_utils/__init__.py

@@ -0,0 +1,127 @@
+from typing import Optional
+import os
+from .config import DBConfig
+from .logger import logger
+
+from dotenv import load_dotenv
+
+from .base_connector import BaseConnector
+
+from .clickhouse_connector import ClickHouseConnector
+from .postgres_connector import PostgresConnector
+from .mysql_connector import MySQLConnector
+from .mariadb_connector import MariaDBConnector
+from .oracle_connector import OracleConnector
+from .duckdb_connector import DuckDBConnector
+from .databricks_connector import DatabricksConnector
+from .snowflake_connector import SnowflakeConnector
+
+env_path = os.path.join(os.getcwd(), ".env")
+
+if os.path.exists(env_path):
+    load_dotenv(env_path, override=True)
+    print(f"✅ 환경변수 파일(.env)이 {os.getcwd()}에 로드되었습니다!")
+else:
+    print(f"⚠️  환경변수 파일(.env)이 {os.getcwd()}에 없습니다!")
+
+def get_db_connector(db_type: Optional[str] = None, config: Optional[DBConfig] = None):
+    """
+    Return the appropriate DB connector instance.
+    - If db_type is not provided, loads from environment variable DB_TYPE
+    - If config is not provided, loads from environment using db_type
+
+    Parameters:
+        db_type (Optional[str]): Database type (e.g., 'postgresql', 'mysql')
+        config (Optional[DBConfig]): Connection config
+
+    Returns:
+        BaseConnector: Initialized DB connector instance
+
+    Raises:
+        ValueError: If type/config is missing or invalid
+    """
+    if db_type is None:
+        db_type = os.getenv("DB_TYPE")
+        if not db_type:
+            raise ValueError("DB type must be provided or set in environment as DB_TYPE.")
+
+    db_type = db_type.lower()
+
+    if config is None:
+        config = load_config_from_env(db_type.upper())
+
+        connector_map = {
+        "clickhouse": ClickHouseConnector,
+        "postgresql": PostgresConnector,
+        "mysql": MySQLConnector,
+        "mariadb": MariaDBConnector,
+        "oracle": OracleConnector,
+        "duckdb": DuckDBConnector,
+        "databricks": DatabricksConnector,
+        "snowflake": SnowflakeConnector,
+    }
+
+    if db_type not in connector_map:
+        logger.error(f"Unsupported DB type: {db_type}")
+        raise ValueError(f"Unsupported DB type: {db_type}")
+
+    required_fields = {
+        "oracle": ["extra.service_name"],
+        "databricks": ["extra.http_path", "extra.access_token"],
+        "snowflake": ["extra.account"],
+    }
+
+    missing = []
+    for path in required_fields.get(db_type, []):
+        cur = config
+        for key in path.split("."):
+            cur = cur.get(key) if isinstance(cur, dict) else None
+            if cur is None:
+                missing.append(path)
+                break
+
+    if missing:
+        logger.error(f"Missing required fields for {db_type}: {', '.join(missing)}")
+        raise ValueError(f"Missing required fields for {db_type}: {', '.join(missing)}")
+
+    return connector_map[db_type](config)
+
+
+
+def load_config_from_env(prefix: str) -> DBConfig:
+    """
+    Load DBConfig from environment variables with a given prefix.
+    Standard keys are extracted, all other prefixed keys go to 'extra'.
+
+    Example:
+        If prefix = 'SNOWFLAKE', loads:
+        - SNOWFLAKE_HOST
+        - SNOWFLAKE_USER
+        - SNOWFLAKE_PASSWORD
+        - SNOWFLAKE_PORT
+        - SNOWFLAKE_DATABASE
+        Other keys like SNOWFLAKE_ACCOUNT, SNOWFLAKE_WAREHOUSE -> extra
+    """
+    base_keys = {"HOST", "PORT", "USER", "PASSWORD", "DATABASE"}
+
+    # Extract standard values
+    config = {
+        "host": os.getenv(f"{prefix}_HOST"),
+        "port": int(os.getenv(f"{prefix}_PORT")) if os.getenv(f"{prefix}_PORT") else None,
+        "user": os.getenv(f"{prefix}_USER"),
+        "password": os.getenv(f"{prefix}_PASSWORD"),
+        "database": os.getenv(f"{prefix}_DATABASE"),
+    }
+
+    # Auto-detect extra keys
+    extra = {}
+    for key, value in os.environ.items():
+        if key.startswith(f"{prefix}_"):
+            suffix = key[len(f"{prefix}_"):]
+            if suffix.upper() not in base_keys:
+                extra[suffix.lower()] = value
+
+    if extra:
+        config["extra"] = extra
+
+    return DBConfig(**config)

db_utils/base_connector.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+import pandas as pd
+
+class BaseConnector(ABC):
+    """
+    Abstract base class for database connectors.
+    """
+
+    @abstractmethod
+    def connect(self):
+        """
+        Initialize the database connection.
+        """
+        pass
+
+    @abstractmethod
+    def run_sql(self, sql: str) -> pd.DataFrame:
+        """
+        Returns the result of the SQL query as a pandas DataFrame.
+
+        Parameters:
+            sql (str): SQL query string to be executed.
+
+        Returns:
+            pd.DataFrame: Result of the SQL query as a pandas DataFrame.
+        """
+        pass

db_utils/clickhouse_connector.py

@@ -0,0 +1,70 @@
+from .base_connector import BaseConnector
+from clickhouse_driver import Client
+import pandas as pd
+from db_utils import DBConfig, logger
+
+class ClickHouseConnector(BaseConnector):
+    """
+    Connect to ClickHouse and execute SQL queries.
+    """
+    client = None
+
+    def __init__(self, config: DBConfig):
+        """
+        Initialize the ClickHouseConnector with connection parameters.
+
+        Parameters:
+            config (DBConfig): Configuration object containing connection parameters.
+        """
+        self.host = config["host"]
+        self.port = config["port"]
+        self.user = config["user"]
+        self.password = config["password"]
+        self.database = config["database"]
+        self.connect()
+
+    def connect(self) -> None:
+        """
+        Establish a connection to the ClickHouse server.
+        """
+        try:
+            self.client = Client(
+                host=self.host,
+                port=self.port,
+                user=self.user,
+                password=self.password,
+                database=self.database,
+            )
+            logger.info("Successfully connected to ClickHouse.")
+        except Exception as e:
+            logger.error(f"Failed to connect to ClickHouse: {e}")
+            raise
+
+    def run_sql(self, sql: str) -> pd.DataFrame:
+        """
+        Execute a SQL query and return the result as a pandas DataFrame.
+
+        Parameters:
+            sql (str): SQL query string to be executed.
+
+        Returns:
+            pd.DataFrame: Result of the SQL query as a pandas DataFrame.
+        """
+        if self.client is None:
+            self.connect()
+
+        try:
+            result = self.client.query(sql).result()
+            return pd.DataFrame(result)
+        except Exception as e:
+            logger.error(f"Failed to execute SQL query: {e}")
+            raise
+
+    def close(self) -> None:
+        """
+        Close the connection to the ClickHouse server.
+        """
+        if self.client:
+            self.client.disconnect()
+            logger.error("Connection to ClickHouse closed.")
+        self.client = None

db_utils/config.py

@@ -0,0 +1,10 @@
+from typing import Optional, Dict, TypedDict
+
+class DBConfig(TypedDict):
+	
+    host: str
+    port: Optional[int]
+    user: Optional[str]
+    password: Optional[str]
+    database: Optional[str]
+    extra: Optional[Dict[str, str]]

db_utils/databricks_connector.py

@@ -0,0 +1,78 @@
+from databricks import sql
+import pandas as pd
+from .base_connector import BaseConnector
+from .config import DBConfig
+from .logger import logger
+
+class DatabricksConnector(BaseConnector):
+    """
+    Connect to Databricks SQL Warehouse and execute queries.
+    """
+    connection = None
+
+    def __init__(self, config: DBConfig):
+        """
+        Initialize the DatabricksConnector with connection parameters.
+
+        Parameters:
+            config (DBConfig): Configuration object containing connection parameters.
+                                 Required keys: host, extra.http_path, extra.access_token
+                                 Optional keys: extra.catalog, extra.schema
+        """
+        self.server_hostname = config["host"]
+        self.http_path = config["extra"]["http_path"]
+        self.access_token = config["extra"]["access_token"]
+        self.catalog = config.get("extra", {}).get("catalog")
+        self.schema = config.get("extra", {}).get("schema")
+        self.connect()
+
+    def connect(self) -> None:
+        """
+        Establish a connection to the Databricks SQL endpoint.
+        """
+        try:
+            self.connection = sql.connect(
+                server_hostname=self.server_hostname,
+                http_path=self.http_path,
+                access_token=self.access_token,
+                catalog=self.catalog,
+                schema=self.schema,
+            )
+            logger.info("Successfully connected to Databricks.")
+        except Exception as e:
+            logger.error(f"Failed to connect to Databricks: {e}")
+            raise
+
+    def run_sql(self, sql: str) -> pd.DataFrame:
+        """
+        Execute a SQL query and return result as pandas DataFrame.
+
+        Parameters:
+            sql (str): SQL query string to be executed.
+
+        Returns:
+            pd.DataFrame: Result of the SQL query as a pandas DataFrame.
+        """
+        if self.connection is None:
+            self.connect()
+
+        try:
+            cursor = self.connection.cursor()
+            cursor.execute(sql)
+            columns = [desc[0] for desc in cursor.description]
+            rows = cursor.fetchall()
+            return pd.DataFrame(rows, columns=columns)
+        except Exception as e:
+            logger.error(f"Failed to execute SQL query: {e}")
+            raise
+        finally:
+            cursor.close()
+
+    def close(self) -> None:
+        """
+        Close the Databricks connection.
+        """
+        if self.connection:
+            self.connection.close()
+            logger.error("Connection to Databricks closed.")
+        self.connection = None