from contextlib import contextmanager

from everyclass.server.config import get_config

_config = get_config()


def init_pool() -> None:
    """创建连接池"""
    from everyclass.common.postgres import init_pool as init
    init(_config.POSTGRES_SCHEMA, _config.POSTGRES_CONNECTION)


MAX_TRIALS = 2


@contextmanager
def pg_conn_context():
    from everyclass.common.postgres import conn_context as context
    success = False
    trials = 0

    while not success and trials < MAX_TRIALS:
        try:
            with context() as conn:
                yield conn
        except RuntimeError:
            # 连接池没有被初始化
            init_pool()
        else:
            success = True
        finally:
            trials += 1
    if not success:
        raise RuntimeError(f"DB connection context failed after {trials} trials")

@@ -384,16 +384,27 @@
 _context
+_with_retry
  as cont
@@ -411,112 +411,8 @@
 ext%0A
-    success = False%0A    trials = 0%0A%0A    while not success and trials %3C MAX_TRIALS:%0A        try:%0A        
     
@@ -443,24 +443,16 @@
         
-     
-   
 yield co
@@ -458,265 +458,4 @@
 onn%0A
-        except RuntimeError:%0A            # %E8%BF%9E%E6%8E%A5%E6%B1%A0%E6%B2%A1%E6%9C%89%E8%A2%AB%E5%88%9D%E5%A7%8B%E5%8C%96%0A            init_pool()%0A        else:%0A            success = True%0A        finally:%0A            trials += 1%0A    if not success:%0A        raise RuntimeError(f%22DB connection context failed after %7Btrials%7D trials%22)%0A

use conn_context_with_retry