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)
|