Dask_bot/utils/database.py
2026-03-30 16:10:15 +02:00

34 lines
1.1 KiB
Python

import sqlite3
class UserLocation:
def __init__(self, db_name='user_location.db'):
self.db_name = db_name
self.create_table()
def create_table(self):
conn = sqlite3.connect(self.db_name)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS user_locations (
user_id INTEGER PRIMARY KEY,
location TEXT NOT NULL
)
''')
conn.commit()
conn.close()
def set_user_location(self, user_id: int, location: str):
conn = sqlite3.connect(self.db_name)
cursor = conn.cursor()
cursor.execute('''
INSERT INTO user_locations (user_id, location) VALUES (?, ?) ON CONFLICT(user_id) DO UPDATE SET location=excluded.location''', (user_id, location))
conn.commit()
conn.close()
def get_user_location(self, user_id: int) -> str:
conn = sqlite3.connect(self.db_name)
cursor = conn.cursor()
cursor.execute('SELECT location FROM user_locations WHERE user_id = ?', (user_id,))
row = cursor.fetchone()
conn.close()
return row[0] if row else None