50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import yaml
|
||
import psycopg2
|
||
|
||
def read_sql_config(yaml_name):
|
||
"""
|
||
读取 SQL 配置
|
||
"""
|
||
yaml_path = f"{yaml_name}.yaml"
|
||
with open(yaml_path, 'r', encoding='utf-8') as f:
|
||
config = yaml.safe_load(f)
|
||
sql_config = config.get('sql')
|
||
if not sql_config:
|
||
raise ValueError("未找到 'sql' 配置块")
|
||
return sql_config
|
||
|
||
def get_user_power(user_id: str, yaml_name: str) -> int | None:
|
||
"""
|
||
根据 user_id 查询数据库中对应的 power，找不到返回 None
|
||
"""
|
||
conn = None
|
||
try:
|
||
sql_config = read_sql_config(yaml_name)
|
||
conn = psycopg2.connect(
|
||
dbname=sql_config['dbname'],
|
||
user=sql_config['user'],
|
||
password=sql_config['password'],
|
||
host=sql_config['host'],
|
||
port=sql_config['port']
|
||
)
|
||
|
||
conn.autocommit = True # 只读查询推荐设置
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT * FROM public.threduser where user_id=%s;", (user_id,))
|
||
row = cur.fetchone()
|
||
if row:
|
||
return row[3]
|
||
else:
|
||
return None
|
||
except Exception as e:
|
||
print(f"数据库操作异常: {e}")
|
||
return None
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
|
||
if __name__ == '__main__':
|
||
user_id = '20250801'
|
||
yaml_name = 'config'
|
||
power = get_user_power(user_id, yaml_name)
|
||
print(power) |