Connection pool 사용하기

김무성·2022년 3월 2일
1

API

목록 보기
4/4

사용방법

클래스 생성

  • util 폴더 내 connection pool 클래스 생성
from psycopg2 import pool
from contextlib import contextmanager
import atexit

class DBHelper:
    def __init__(self):
        self._connection_pool = None

    def initialize_connection_pool(self):
        self._connection_pool = pool.ThreadedConnectionPool(1, 5, user = "user",
                                              password = "password",
                                              host = "0.0.0.0",
                                              port = "0000",
                                              database = "db")

    @contextmanager
    def get_resource_rdb(self, autocommit=True):
        if self._connection_pool is None:
            self.initialize_connection_pool()

        conn = self._connection_pool.getconn()
        conn.autocommit = autocommit
        cursor = conn.cursor()
        
        try:
            yield cursor, conn
        finally:
            cursor.close()
            self._connection_pool.putconn(conn)

    def shutdown_connection_pool(self):
        if self._connection_pool is not None:
            self._connection_pool.closeall()

db_helper = DBHelper()

@atexit.register
def shutdown_connection_pool():
    db_helper.shutdown_connection_pool()

클래스 import

  • connection pool을 사용하는 파일에서 클래스 import
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
from util.db_helper import db_helper
  • 다른 폴더의 클래스를 import하기 위해 경로지정이 필요함

클래스 내 함수 사용

  • 파이썬 with 함수를 이용하여 클래스를 불러와 connection pool을 생성하고 불러오고 반납 함
with db_helper.get_resource_rdb() as (cursor, _):
        cursor.execute(query)

db접근

  • get_resource_rdb 함수 사용
    @contextmanager
    def get_resource_rdb(self, autocommit=True):
        if self._connection_pool is None:
            self.initialize_connection_pool()

        conn = self._connection_pool.getconn()
        conn.autocommit = autocommit
        cursor = conn.cursor()
        
        try:
            yield cursor, conn
        finally:
            cursor.close()
            self._connection_pool.putconn(conn)

데이터를 적재하는 경우

  • auto commit이 되어있으므로 쿼리를 excute
with db_helper.get_resource_rdb() as (cursor, _):
        cursor.execute(query)

데이터를 불러오는 경우

  • fetchall을 통해 데이터 불러 옴
with db_helper.get_resource_rdb() as (cursor, _):
        cursor.execute(query)
        result = cursor.fetchall()
profile
graph data scientist

0개의 댓글