Files
temp-file-trans/database.py

62 lines
1.8 KiB
Python
Raw Normal View History

import sqlite3
import os
from datetime import datetime, timedelta
from config import DATABASE
def get_db():
conn = sqlite3.connect(DATABASE)
conn.row_factory = sqlite3.Row
return conn
def init_db():
os.makedirs(os.path.dirname(DATABASE) if os.path.dirname(DATABASE) else '.', exist_ok=True)
conn = get_db()
conn.execute('''
CREATE TABLE IF NOT EXISTS files (
id TEXT PRIMARY KEY,
filename TEXT NOT NULL,
filepath TEXT NOT NULL,
filesize INTEGER NOT NULL,
expiry_hours INTEGER NOT NULL,
created_at TIMESTAMP NOT NULL,
expires_at TIMESTAMP NOT NULL
)
''')
conn.commit()
conn.close()
def add_file(file_id, filename, filepath, filesize, expiry_hours):
now = datetime.utcnow()
expires = now + timedelta(hours=expiry_hours)
conn = get_db()
conn.execute(
'INSERT INTO files (id, filename, filepath, filesize, expiry_hours, created_at, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)',
(file_id, filename, filepath, filesize, expiry_hours, now, expires)
)
conn.commit()
conn.close()
def get_file(file_id):
conn = get_db()
row = conn.execute('SELECT * FROM files WHERE id = ?', (file_id,)).fetchone()
conn.close()
return row
def delete_file(file_id):
conn = get_db()
conn.execute('DELETE FROM files WHERE id = ?', (file_id,))
conn.commit()
conn.close()
def cleanup_expired():
now = datetime.utcnow()
conn = get_db()
expired = conn.execute('SELECT * FROM files WHERE expires_at < ?', (now,)).fetchall()
for row in expired:
if os.path.exists(row['filepath']):
os.remove(row['filepath'])
conn.execute('DELETE FROM files WHERE expires_at < ?', (now,))
conn.commit()
conn.close()
return len(expired)