OAuth 인증과 커넥션 풀링

import os
from databricks.sdk import WorkspaceClient

# Databricks Apps 환경에서는 자동으로 서비스 프린시펄 인증이 설정됩니다
w = WorkspaceClient()

# Lakebase 크레덴셜 생성
credential = w.lakebase.generate_credential(
    database_name="my_catalog.my_schema.shop_db"
)

# 생성된 크레덴셜로 연결
conn = psycopg2.connect(
    host=credential.host,
    port=credential.port,
    dbname=credential.database,
    user=credential.username,
    password=credential.password,
    sslmode="require"
)

from psycopg2 import pool

# 앱 시작 시 풀 초기화
connection_pool = pool.ThreadedConnectionPool(
    minconn=2,      # 최소 연결 수
    maxconn=20,     # 최대 연결 수
    host="app-db-xxxx.lakebase.databricks.com",
    port=5432,
    dbname="shop_db",
    user="token",
    password="dapi_your_token",
    sslmode="require"
)

# 요청 처리 시
def handle_request():
    conn = connection_pool.getconn()
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM orders LIMIT 10")
        results = cursor.fetchall()
        return results
    finally:
        connection_pool.putconn(conn)  # 풀에 연결 반환

def transfer_order(conn, order_id, new_status):
    """주문 상태 변경과 재고 업데이트를 트랜잭션으로 처리합니다."""
    try:
        cursor = conn.cursor()
        
        # 트랜잭션 시작 (autocommit=False가 기본)
        cursor.execute(
            "UPDATE orders SET status = %s WHERE id = %s RETURNING product",
            (new_status, order_id)
        )
        result = cursor.fetchone()
        if result is None:
            raise ValueError(f"주문 {order_id}을 찾을 수 없습니다")
        
        product = result[0]
        
        # 상태가 'shipped'이면 재고 차감
        if new_status == 'shipped':
            cursor.execute(
                "UPDATE inventory SET stock = stock - 1 "
                "WHERE product = %s AND stock > 0 RETURNING stock",
                (product,)
            )
            if cursor.fetchone() is None:
                raise ValueError(f"재고 부족: {product}")
        
        conn.commit()  # 모든 변경사항 확정
        return {"success": True, "message": f"주문 {order_id} → {new_status}"}
    
    except Exception as e:
        conn.rollback()  # 오류 시 모든 변경사항 취소
        return {"success": False, "error": str(e)}

import psycopg2
from psycopg2 import errors
import logging

logger = logging.getLogger(__name__)

def safe_db_operation(pool, query, params=None):
    """안전한 데이터베이스 작업 패턴"""
    conn = None
    try:
        conn = pool.getconn()
        cursor = conn.cursor()
        cursor.execute(query, params)
        
        if query.strip().upper().startswith("SELECT"):
            result = cursor.fetchall()
        else:
            conn.commit()
            result = cursor.rowcount
        
        return {"success": True, "data": result}
    
    except errors.UniqueViolation as e:
        # 중복 키 오류 (예: 이메일 중복)
        if conn: conn.rollback()
        logger.warning(f"중복 데이터: {e}")
        return {"success": False, "error": "이미 존재하는 데이터입니다"}
    
    except errors.ForeignKeyViolation as e:
        # 외래 키 참조 오류
        if conn: conn.rollback()
        logger.warning(f"참조 무결성 위반: {e}")
        return {"success": False, "error": "참조하는 데이터가 존재하지 않습니다"}
    
    except errors.OperationalError as e:
        # 연결 오류 (네트워크, 타임아웃 등)
        logger.error(f"DB 연결 오류: {e}")
        if conn:
            pool.putconn(conn, close=True)  # 불량 연결 폐기
            conn = None
        return {"success": False, "error": "데이터베이스 연결 오류가 발생했습니다"}
    
    except Exception as e:
        if conn: conn.rollback()
        logger.error(f"예상치 못한 오류: {e}")
        return {"success": False, "error": "처리 중 오류가 발생했습니다"}
    
    finally:
        if conn:
            pool.putconn(conn)