import pymysql
import pymysql.cursors
import pg8000
import oss2
from .config import SOURCE_DB, HK_SONGS_DB, CRAWLER_DB, OSS_CONFIG, OSS_CONNECTION_POOL_SIZE


def get_source_conn() -> pymysql.Connection:
    return pymysql.connect(**SOURCE_DB, cursorclass=pymysql.cursors.DictCursor)


def get_spider_conn() -> pymysql.Connection:
    cfg = SOURCE_DB.copy()
    cfg['database'] = 'hikoon-data-spider'
    return pymysql.connect(**cfg, cursorclass=pymysql.cursors.DictCursor)


def get_hk_songs_conn() -> pymysql.Connection:
    return pymysql.connect(**HK_SONGS_DB, cursorclass=pymysql.cursors.DictCursor)


def get_pg_conn() -> pg8000.Connection:
    conn = pg8000.connect(**CRAWLER_DB)
    conn.run("SET TimeZone = 'Asia/Shanghai'")
    return conn


def get_oss_bucket() -> oss2.Bucket:
    oss2.defaults.connection_pool_size = OSS_CONNECTION_POOL_SIZE
    auth = oss2.Auth(OSS_CONFIG['access_key_id'], OSS_CONFIG['access_key_secret'])
    return oss2.Bucket(auth, OSS_CONFIG['endpoint'], OSS_CONFIG['bucket_name'])