#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Script to generate and save coupons to the database """ import os import sys import random from datetime import datetime, timedelta # Add the backend directory to the path so we can import modules sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) # Import required modules from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from sqlalchemy.exc import SQLAlchemyError from backend.app.admin.model.coupon import Coupon from backend.utils.snowflake import snowflake from backend.core.conf import settings, get_db_uri def generate_coupon_codes(prefix: str, quantity: int): """ Generate coupon codes with specified prefix and quantity. Format: [PREFIX][NUMBER] - Total 6 characters Example: A12345, TEST0, XYZ999 Args: prefix (str): The letter prefix for the coupon codes (should be uppercase) quantity (int): Number of coupon codes to generate Returns: list: List of generated coupon codes """ if not prefix.isalpha() or not prefix.isupper(): raise ValueError("Prefix must be uppercase letters only") if len(prefix) == 0 or len(prefix) > 5: raise ValueError("Prefix must be 1-5 characters long") if quantity <= 0: raise ValueError("Quantity must be greater than 0") # Calculate number of digits based on prefix length (total 6 characters) num_digits = 6 - len(prefix) # Maximum possible combinations max_combinations = 10 ** num_digits if quantity > max_combinations: raise ValueError(f"With prefix '{prefix}' (length {len(prefix)}), can only generate {max_combinations} unique codes (0 to {max_combinations - 1})") codes = [] # Generate incremental numbers starting from 0 for i in range(quantity): # Format with leading zeros to make it the required number of digits formatted_number = f"{i:0{num_digits}d}" # Combine prefix with formatted number coupon_code = f"{prefix}{formatted_number}" codes.append(coupon_code) return codes def save_coupons_to_db(prefix: str, quantity: int, coupon_type: str, points: int, expire_days: int = None): """ Generate and save coupons to the database. Coupon codes are always 6 characters total: - 1-letter prefix: 5 digits (up to 100000 codes: A00000-A99999) - 4-letter prefix: 2 digits (up to 100 codes: TEST00-TEST99) - 5-letter prefix: 1 digit (up to 10 codes: ABCDE0-ABCDE9) Args: prefix (str): The letter prefix for the coupon codes quantity (int): Number of coupon codes to generate coupon_type (str): Type of the coupons points (int): Points value of the coupons expire_days (int, optional): Days until expiration. If None, no expiration. """ # Create database engine and session db_url = get_db_uri(settings) # Replace asyncmy with mysql+mysqlconnector for synchronous connection sync_db_url = db_url.replace('mysql+asyncmy', 'mysql+mysqlconnector') try: engine = create_engine(sync_db_url, echo=False) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) db = SessionLocal() # Generate coupon codes codes = generate_coupon_codes(prefix, quantity) # Create coupon objects coupons = [] for code in codes: # Generate snowflake ID coupon_id = snowflake.generate() # Calculate expiration date if needed expires_at = None if expire_days is not None and expire_days > 0: expires_at = datetime.now() + timedelta(days=expire_days) # Create coupon object # Note: id is auto-generated by snowflake, but we want to use our own snowflake generator coupon = Coupon( code=code, type=coupon_type, points=points, expires_at=expires_at ) # Set the id manually after creation coupon.id = coupon_id coupons.append(coupon) # Bulk insert coupons db.add_all(coupons) db.commit() print(f"Successfully saved {len(coupons)} coupons to the database.") print(f"Prefix: {prefix}, Type: {coupon_type}, Points: {points}") if expire_days: print(f"Expires in: {expire_days} days") # Display first 5 coupons as examples print("\nSample coupons generated:") for coupon in coupons[:5]: print(f" ID: {coupon.id}, Code: {coupon.code}") db.close() except SQLAlchemyError as e: print(f"Database error: {e}") if 'db' in locals(): db.rollback() db.close() except Exception as e: print(f"Error: {e}") if 'db' in locals(): db.close() def main(): """Main function to demonstrate usage""" print("Coupon Generator and Database Saver") print("=" * 40) # Example: Generate and save coupons with different prefixes try: # Single character prefix (5 digits, incremental from 00000) # print("Generating coupons with single character prefix 'A'...") # save_coupons_to_db('A', 5, 'NORMAL', 100, 30) # print("\n" + "-" * 40 + "\n") # 4-character prefix (2 digits, incremental from 00) print("Generating coupons with 4-character prefix 'TEST'...") save_coupons_to_db('VIP', 5, 'test', 1000, 60) print("\n" + "-" * 40 + "\n") # 3-character prefix (3 digits, incremental from 000) # print("Generating coupons with 3-character prefix 'XYZ'...") # save_coupons_to_db('XYZ', 3, 'SPECIAL', 500, 15) # print("\n" + "-" * 40 + "\n") # 5-character prefix (1 digit, incremental from 0) # print("Generating coupons with 5-character prefix 'ABCDE'...") # save_coupons_to_db('ABCDE', 5, 'PREMIUM', 2000, 90) except Exception as e: print(f"Error in main: {e}") if __name__ == "__main__": main()