Python MySQL: with OPEN_DB 만들기

 

VSCode

 

commit은 항상 하고, cursor를 주로 사용하기 때문에

알아서 close 시키는 with문을 만들고 싶었다.

 

custom with문 만들기

import os
from contextlib import contextmanager


@contextmanager
def OPEN_DB():
    # connect to my local MySQL instance using connection string
    db = mysql.connector.connect(
        host="localhost",
        user=os.environ["MYSQL_USER"],
        password=os.environ["MYSQL_PASSWORD"],
        database="mydb",
        charset="utf8",
    )
    cursor = db.cursor(buffered=True)
    yield cursor  # cursor를 반환
    db.commit()  # always commits
    cursor.close()
    db.close()

 

사용법

# pytest
def test_delete():
    with OPEN_DB() as cursor:
        cursor.execute("DELETE FROM notices WHERE id = 12245")

# pytest
def test_lastparsed():
    """Test LAST_PARSED update"""
    from datetime import datetime

    now = datetime.now()
    now = now.strftime("%Y-%m-%d %H:%M:%S.%f")
    with OPEN_DB() as cursor:
        cursor.execute(UPDATE_COMMAND, {"date": now})

        cursor.execute("SELECT date FROM notices WHERE id = 1")

        fetch = cursor.fetchone()[0]
        assert fetch == now, "Updating LAST_PARSED date failed."
        print(fetch)  # 2020-12-14 20:29:24.222095