[Python] 죽지않는 웹소켓 만들기

 

https://stickode.tistory.com/1045

 

저번 시간에 이어서, 오늘은 데이터 송수신이 없더라도 소켓 연결이 해제되지 않도록 코드를 수정해보겠습니다.

기존 코드는 특정 시간 동안 웹소켓이 데이터를 주고받지 않는다면 연결이 해제되도록 작성되었습니다.

변경한 사용자가 ctrl + c를 눌러서 직접 종료하지 않는 한 소켓 연결이 해제되지 않습니다.

 

아래 파일을 복사하신 후 실행하시면 됩니다. 파일명은 main.py입니다

import json
import websocket
import pymysql
from websocket import WebSocketApp, WebSocketConnectionClosedException
import logging
import time

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 데이터베이스 연결 설정
db = pymysql.connect(host='localhost', user='root', password='root', db='stickode_231204_upbit', charset='utf8mb4')
cursor = db.cursor()

def on_message(ws, message):
    data = json.loads(message)  # JSON 데이터 파싱

    # 데이터베이스에 데이터 삽입
    sql = '''INSERT INTO trades (type, code, timestamp, trade_date, trade_time, trade_timestamp, trade_price, trade_volume, ask_bid, prev_closing_price, `change`, change_price, sequential_id, stream_type) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)'''
    values = (
        data['type'], data['code'], data['timestamp'], data['trade_date'], data['trade_time'], data['trade_timestamp'],
        data['trade_price'], data['trade_volume'], data['ask_bid'], data['prev_closing_price'], data['change'],
        data['change_price'], data['sequential_id'], data['stream_type'])

    try:
        cursor.execute(sql, values)
        db.commit()
        logging.info("Insert successful!")  # 데이터 삽입 성공 메시지 출력

    except Exception as e:
        logging.error(f"Database Error: {e}")
        db.rollback()

    # logging.info(data)  # 디코딩된 데이터 출력


def on_connect(ws):
    logging.info("connected!")  # 연결됨 메시지 출력
    # 연결 시 서버에 보낼 메시지 전송
    ws.send('[{"ticket": "test" }, {"type": "trade", "codes": ["KRW-BTC" ] }, {"format": "DEFAULT" } ]')


def on_error(ws, err):
    logging.error(err)  # 에러 발생 시 에러 메시지 출력


def on_close(ws, status_code, msg):
    logging.info("closed!")  # 연결 종료 시 메시지 출력


def run_websocket():
    while True:
        ws_app = WebSocketApp("wss://api.upbit.com/websocket/v1",
                              on_message=on_message,
                              on_error=on_error,
                              on_close=on_close)
        ws_app.on_open = on_connect
        ws_app.run_forever(ping_interval=60, ping_timeout=30)
        time.sleep(10)  # 연결이 끊어진 후 재시도 전 대기 시간

try:
    run_websocket()
except KeyboardInterrupt:
    logging.info("프로그램이 사용자에 의해 중단되었습니다.")

# 프로그램 종료 시 데이터베이스 연결 종료
db.close()