324 lines
11 KiB
Python
324 lines
11 KiB
Python
|
"""
|
||
|
|
Ticket Storage Utility
|
||
|
|
======================
|
||
|
|
Modular storage system for ticket data with caching support.
|
||
|
|
"""
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import Dict, Optional, Any, List
|
||
|
|
from threading import Lock
|
||
|
|
|
||
|
|
from .models import TicketData, GuildTicketConfig, TicketStatus
|
||
|
|
|
||
|
|
|
||
|
|
class TicketStorage:
|
||
|
|
"""
|
||
|
|
Thread-safe storage manager for ticket data.
|
||
|
|
|
||
|
|
Features:
|
||
|
|
- Automatic file creation
|
||
|
|
- In-memory caching with write-through
|
||
|
|
- Backup creation before writes
|
||
|
|
- Data validation on load
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, base_path: str = "data/tickets"):
|
||
|
|
self.base_path = base_path
|
||
|
|
self.config_path = os.path.join(base_path, "config.json")
|
||
|
|
self.tickets_path = os.path.join(base_path, "tickets")
|
||
|
|
self.backup_path = os.path.join(base_path, "backups")
|
||
|
|
self.cache: Dict[str, Any] = {}
|
||
|
|
self.cache_lock = Lock()
|
||
|
|
self._ensure_directories()
|
||
|
|
|
||
|
|
def _ensure_directories(self) -> None:
|
||
|
|
"""Create necessary directories"""
|
||
|
|
for path in [self.base_path, self.tickets_path, self.backup_path]:
|
||
|
|
os.makedirs(path, exist_ok=True)
|
||
|
|
|
||
|
|
def _create_backup(self, file_path: str) -> str:
|
||
|
|
"""Create a timestamped backup of a file"""
|
||
|
|
if not os.path.exists(file_path):
|
||
|
|
return ""
|
||
|
|
|
||
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
|
|
filename = os.path.basename(file_path)
|
||
|
|
backup_file = os.path.join(self.backup_path, f"{timestamp}_{filename}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
with open(file_path, 'r', encoding='utf-8') as src:
|
||
|
|
with open(backup_file, 'w', encoding='utf-8') as dst:
|
||
|
|
dst.write(src.read())
|
||
|
|
return backup_file
|
||
|
|
except Exception:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
def _get_cache_key(self, guild_id: int, key_type: str) -> str:
|
||
|
|
"""Generate a cache key"""
|
||
|
|
return f"{guild_id}_{key_type}"
|
||
|
|
|
||
|
|
# --- Config Methods ---
|
||
|
|
|
||
|
|
def load_config(self, guild_id: int) -> GuildTicketConfig:
|
||
|
|
"""Load guild ticket configuration"""
|
||
|
|
cache_key = self._get_cache_key(guild_id, "config")
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
if cache_key in self.cache:
|
||
|
|
return self.cache[cache_key]
|
||
|
|
|
||
|
|
try:
|
||
|
|
if os.path.exists(self.config_path):
|
||
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
||
|
|
all_configs = json.load(f)
|
||
|
|
|
||
|
|
config_data = all_configs.get(str(guild_id), {})
|
||
|
|
config = GuildTicketConfig.from_dict(config_data)
|
||
|
|
config.guild_id = guild_id
|
||
|
|
else:
|
||
|
|
config = GuildTicketConfig(guild_id=guild_id)
|
||
|
|
except (json.JSONDecodeError, KeyError) as e:
|
||
|
|
print(f"[TicketStorage] Error loading config for guild {guild_id}: {e}")
|
||
|
|
config = GuildTicketConfig(guild_id=guild_id)
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
self.cache[cache_key] = config
|
||
|
|
|
||
|
|
return config
|
||
|
|
|
||
|
|
def save_config(self, guild_id: int, config: GuildTicketConfig) -> bool:
|
||
|
|
"""Save guild ticket configuration"""
|
||
|
|
cache_key = self._get_cache_key(guild_id, "config")
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Create backup
|
||
|
|
self._create_backup(self.config_path)
|
||
|
|
|
||
|
|
# Load all configs
|
||
|
|
all_configs = {}
|
||
|
|
if os.path.exists(self.config_path):
|
||
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
||
|
|
all_configs = json.load(f)
|
||
|
|
|
||
|
|
# Update config
|
||
|
|
all_configs[str(guild_id)] = config.to_dict()
|
||
|
|
|
||
|
|
# Save
|
||
|
|
with open(self.config_path, 'w', encoding='utf-8') as f:
|
||
|
|
json.dump(all_configs, f, indent=4, ensure_ascii=False)
|
||
|
|
|
||
|
|
# Update cache
|
||
|
|
with self.cache_lock:
|
||
|
|
self.cache[cache_key] = config
|
||
|
|
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[TicketStorage] Error saving config for guild {guild_id}: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
# --- Ticket Methods ---
|
||
|
|
|
||
|
|
def get_ticket_path(self, channel_id: int) -> str:
|
||
|
|
"""Get the file path for a ticket"""
|
||
|
|
return os.path.join(self.tickets_path, f"{channel_id}.json")
|
||
|
|
|
||
|
|
def load_ticket(self, channel_id: int) -> Optional[TicketData]:
|
||
|
|
"""Load a ticket by channel ID"""
|
||
|
|
cache_key = self._get_cache_key(channel_id, "ticket")
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
if cache_key in self.cache:
|
||
|
|
return self.cache[cache_key]
|
||
|
|
|
||
|
|
path = self.get_ticket_path(channel_id)
|
||
|
|
|
||
|
|
try:
|
||
|
|
if os.path.exists(path):
|
||
|
|
with open(path, 'r', encoding='utf-8') as f:
|
||
|
|
data = json.load(f)
|
||
|
|
ticket = TicketData.from_dict(data)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
except (json.JSONDecodeError, KeyError) as e:
|
||
|
|
print(f"[TicketStorage] Error loading ticket {channel_id}: {e}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
self.cache[cache_key] = ticket
|
||
|
|
|
||
|
|
return ticket
|
||
|
|
|
||
|
|
def save_ticket(self, ticket: TicketData) -> bool:
|
||
|
|
"""Save a ticket"""
|
||
|
|
cache_key = self._get_cache_key(ticket.channel_id, "ticket")
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Create backup
|
||
|
|
self._create_backup(self.get_ticket_path(ticket.channel_id))
|
||
|
|
|
||
|
|
path = self.get_ticket_path(ticket.channel_id)
|
||
|
|
|
||
|
|
with open(path, 'w', encoding='utf-8') as f:
|
||
|
|
json.dump(ticket.to_dict(), f, indent=4, ensure_ascii=False)
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
self.cache[cache_key] = ticket
|
||
|
|
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[TicketStorage] Error saving ticket {ticket.channel_id}: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
def delete_ticket(self, channel_id: int) -> bool:
|
||
|
|
"""Delete a ticket file"""
|
||
|
|
cache_key = self._get_cache_key(channel_id, "ticket")
|
||
|
|
path = self.get_ticket_path(channel_id)
|
||
|
|
|
||
|
|
try:
|
||
|
|
if os.path.exists(path):
|
||
|
|
# Create backup before delete
|
||
|
|
self._create_backup(path)
|
||
|
|
os.remove(path)
|
||
|
|
|
||
|
|
with self.cache_lock:
|
||
|
|
self.cache.pop(cache_key, None)
|
||
|
|
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[TicketStorage] Error deleting ticket {channel_id}: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
def get_guild_tickets(self, guild_id: int, status: Optional[TicketStatus] = None) -> List[TicketData]:
|
||
|
|
"""Get all tickets for a guild, optionally filtered by status"""
|
||
|
|
tickets = []
|
||
|
|
|
||
|
|
if not os.path.exists(self.tickets_path):
|
||
|
|
return tickets
|
||
|
|
|
||
|
|
for filename in os.listdir(self.tickets_path):
|
||
|
|
if not filename.endswith('.json'):
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
channel_id = int(filename.replace('.json', ''))
|
||
|
|
ticket = self.load_ticket(channel_id)
|
||
|
|
|
||
|
|
if ticket and ticket.guild_id == guild_id:
|
||
|
|
if status is None or ticket.status == status:
|
||
|
|
tickets.append(ticket)
|
||
|
|
except (ValueError, TypeError):
|
||
|
|
continue
|
||
|
|
|
||
|
|
return tickets
|
||
|
|
|
||
|
|
def get_user_tickets(self, guild_id: int, user_id: int, open_only: bool = True) -> List[TicketData]:
|
||
|
|
"""Get all tickets for a user"""
|
||
|
|
tickets = self.get_guild_tickets(guild_id)
|
||
|
|
user_tickets = [t for t in tickets if t.user_id == user_id]
|
||
|
|
|
||
|
|
if open_only:
|
||
|
|
user_tickets = [t for t in user_tickets if t.is_open]
|
||
|
|
|
||
|
|
return user_tickets
|
||
|
|
|
||
|
|
def count_user_tickets(self, guild_id: int, user_id: int, category_id: Optional[str] = None) -> int:
|
||
|
|
"""Count tickets for a user, optionally in a specific category"""
|
||
|
|
tickets = self.get_user_tickets(guild_id, user_id, open_only=True)
|
||
|
|
|
||
|
|
if category_id:
|
||
|
|
tickets = [t for t in tickets if t.category_id == category_id]
|
||
|
|
|
||
|
|
return len(tickets)
|
||
|
|
|
||
|
|
# --- Bulk Operations ---
|
||
|
|
|
||
|
|
def get_all_guild_ids(self) -> List[int]:
|
||
|
|
"""Get all guild IDs with ticket configurations"""
|
||
|
|
if not os.path.exists(self.config_path):
|
||
|
|
return []
|
||
|
|
|
||
|
|
try:
|
||
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
||
|
|
configs = json.load(f)
|
||
|
|
return [int(gid) for gid in configs.keys()]
|
||
|
|
except (json.JSONDecodeError, ValueError):
|
||
|
|
return []
|
||
|
|
|
||
|
|
def cleanup_cache(self, max_size: int = 1000) -> int:
|
||
|
|
"""Clear cache if it exceeds max size"""
|
||
|
|
with self.cache_lock:
|
||
|
|
if len(self.cache) <= max_size:
|
||
|
|
return 0
|
||
|
|
|
||
|
|
# Clear all non-config keys (keep recent)
|
||
|
|
removed = 0
|
||
|
|
keys_to_remove = [k for k in self.cache.keys() if not k.endswith("_config")]
|
||
|
|
|
||
|
|
for key in keys_to_remove[:len(keys_to_remove) - max_size]:
|
||
|
|
del self.cache[key]
|
||
|
|
removed += 1
|
||
|
|
|
||
|
|
return removed
|
||
|
|
|
||
|
|
def invalidate_cache(self, guild_id: Optional[int] = None) -> None:
|
||
|
|
"""Invalidate cache for a guild or all"""
|
||
|
|
with self.cache_lock:
|
||
|
|
if guild_id is None:
|
||
|
|
self.cache.clear()
|
||
|
|
else:
|
||
|
|
keys_to_remove = [k for k in self.cache.keys() if k.startswith(str(guild_id))]
|
||
|
|
for key in keys_to_remove:
|
||
|
|
del self.cache[key]
|
||
|
|
|
||
|
|
# --- Statistics ---
|
||
|
|
|
||
|
|
def get_storage_stats(self) -> Dict[str, Any]:
|
||
|
|
"""Get storage statistics"""
|
||
|
|
stats = {
|
||
|
|
"total_tickets": 0,
|
||
|
|
"total_backups": 0,
|
||
|
|
"config_size_bytes": 0,
|
||
|
|
"cache_size": 0,
|
||
|
|
"tickets_by_status": {}
|
||
|
|
}
|
||
|
|
|
||
|
|
# Count tickets
|
||
|
|
if os.path.exists(self.tickets_path):
|
||
|
|
for filename in os.listdir(self.tickets_path):
|
||
|
|
if filename.endswith('.json'):
|
||
|
|
stats["total_tickets"] += 1
|
||
|
|
|
||
|
|
# Count backups
|
||
|
|
if os.path.exists(self.backup_path):
|
||
|
|
stats["total_backups"] = len(os.listdir(self.backup_path))
|
||
|
|
|
||
|
|
# Config size
|
||
|
|
if os.path.exists(self.config_path):
|
||
|
|
stats["config_size_bytes"] = os.path.getsize(self.config_path)
|
||
|
|
|
||
|
|
# Cache size
|
||
|
|
with self.cache_lock:
|
||
|
|
stats["cache_size"] = len(self.cache)
|
||
|
|
|
||
|
|
return stats
|
||
|
|
|
||
|
|
|
||
|
|
# Singleton instance
|
||
|
|
_storage_instance: Optional[TicketStorage] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_storage() -> TicketStorage:
|
||
|
|
"""Get the singleton storage instance"""
|
||
|
|
global _storage_instance
|
||
|
|
if _storage_instance is None:
|
||
|
|
_storage_instance = TicketStorage()
|
||
|
|
return _storage_instance
|
||
|
|
|
||
|
|
|
||
|
|
def reset_storage() -> None:
|
||
|
|
"""Reset the storage instance (for testing)"""
|
||
|
|
global _storage_instance
|
||
|
|
_storage_instance = None
|
||
|
|
|