package repository import ( "fmt" "math/rand" "time" "devone.aplikasi.web.id/gitea/mario/go-ohif-proxy/internal/database" "github.com/jmoiron/sqlx" ) // DBShortCode represents a shortcode from the database type DBShortCode struct { ID int `db:"id"` Shortcode string `db:"Shortcode"` ShortcodeIsUsed string `db:"ShortcodeIsUsed"` ShortcodeIsUsedBy_ShortlinkID *int `db:"ShortcodeIsUsedBy_ShortlinkID"` ShortcodeCreatedAt time.Time `db:"ShortcodeCreatedAt"` ShortcodeUpdatedAt time.Time `db:"ShortcodeUpdatedAt"` } // ShortCodeRepository handles database operations related to shortcodes type ShortCodeRepository struct { *Repository } // NewShortCodeRepository creates a new shortcode repository func NewShortCodeRepository() *ShortCodeRepository { return &ShortCodeRepository{ Repository: NewRepository(), } } // GetUnusedShortCode retrieves an unused shortcode from the database func (r *ShortCodeRepository) GetUnusedShortCode() (*string, error) { var shortcode string query := `SELECT Shortcode FROM shortcodes WHERE ShortcodeIsUsed = 'N' ORDER BY RAND() LIMIT 1` err := database.DB.Get(&shortcode, query) if err != nil { return nil, fmt.Errorf("database error getting unused shortcode: %w", err) } return &shortcode, nil } // MarkShortCodeAsUsed marks a shortcode as used by a specific shortlink func (r *ShortCodeRepository) MarkShortCodeAsUsed(tx *sqlx.Tx, shortcode string, shortlinkID int) error { query := `UPDATE shortcodes SET ShortcodeIsUsed = 'Y', ShortcodeIsUsedBy_ShortlinkID = ? WHERE Shortcode = ?` _, err := tx.Exec(query, shortlinkID, shortcode) if err != nil { return fmt.Errorf("database error marking shortcode as used: %w", err) } return nil } // MarkShortCodeAsUnused marks a shortcode as unused when shortlink expires or is deleted func (r *ShortCodeRepository) MarkShortCodeAsUnused(shortcode string) error { query := `UPDATE shortcodes SET ShortcodeIsUsed = 'N', ShortcodeIsUsedBy_ShortlinkID = NULL WHERE Shortcode = ?` _, err := database.DB.Exec(query, shortcode) if err != nil { return fmt.Errorf("database error marking shortcode as unused: %w", err) } return nil } // MarkShortCodeAsUnusedByShortlinkID marks a shortcode as unused by shortlink ID func (r *ShortCodeRepository) MarkShortCodeAsUnusedByShortlinkID(shortlinkID int) error { query := `UPDATE shortcodes SET ShortcodeIsUsed = 'N', ShortcodeIsUsedBy_ShortlinkID = NULL WHERE ShortcodeIsUsedBy_ShortlinkID = ?` _, err := database.DB.Exec(query, shortlinkID) if err != nil { return fmt.Errorf("database error marking shortcode as unused by shortlink ID: %w", err) } return nil } // MarkExpiredShortlinkShortcodesAsUnused finds all shortcodes associated with expired shortlinks // and marks them as unused func (r *ShortCodeRepository) MarkExpiredShortlinkShortcodesAsUnused() (int, error) { query := ` UPDATE shortcodes sc INNER JOIN shortlink sl ON sc.ShortcodeIsUsedBy_ShortlinkID = sl.ShortlinkID SET sc.ShortcodeIsUsed = 'N', sc.ShortcodeIsUsedBy_ShortlinkID = NULL WHERE (sl.ShortlinkExpiredAt < NOW() OR sl.ShortlinkIsRevoked = TRUE) AND sc.ShortcodeIsUsed = 'Y' ` result, err := database.DB.Exec(query) if err != nil { return 0, fmt.Errorf("database error freeing expired shortlink shortcodes: %w", err) } count, err := result.RowsAffected() if err != nil { return 0, fmt.Errorf("error getting affected rows count: %w", err) } return int(count), nil } // GenerateUniqueShortCode generates a unique 5-character capital letter shortcode func GenerateUniqueShortCode() string { const charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" const codeLength = 5 code := make([]byte, codeLength) for i := range code { code[i] = charset[rand.Intn(len(charset))] } return string(code) } // SeedShortCodes seeds the shortcodes table with n unique 5-character codes func (r *ShortCodeRepository) SeedShortCodes(n int) error { // Initialize random seed rand.Seed(time.Now().UnixNano()) // Use a transaction for better performance tx, err := database.DB.Beginx() if err != nil { return fmt.Errorf("failed to start transaction: %w", err) } // Set up deferred rollback that will be canceled if we commit defer func() { if tx != nil { tx.Rollback() } }() // Check how many shortcodes already exist var count int err = tx.Get(&count, "SELECT COUNT(*) FROM shortcodes") if err != nil { return fmt.Errorf("failed to count existing shortcodes: %w", err) } fmt.Printf("Found %d existing shortcodes, generating %d more\n", count, n) // Track already generated codes to avoid duplicates generatedCodes := make(map[string]bool) // Get existing codes to avoid duplicates existingCodes := []string{} err = tx.Select(&existingCodes, "SELECT Shortcode FROM shortcodes") if err != nil { return fmt.Errorf("failed to retrieve existing shortcodes: %w", err) } // Add existing codes to the map for _, code := range existingCodes { generatedCodes[code] = true } // Prepare the insert statement stmt, err := tx.Prepare("INSERT INTO shortcodes (Shortcode, ShortcodeIsUsed) VALUES (?, 'N')") if err != nil { return fmt.Errorf("failed to prepare statement: %w", err) } defer stmt.Close() // Generate and insert the shortcodes inserted := 0 attempts := 0 maxAttempts := n * 10 // Limiting attempts to avoid infinite loop for inserted < n && attempts < maxAttempts { attempts++ code := GenerateUniqueShortCode() if !generatedCodes[code] { _, err := stmt.Exec(code) if err != nil { continue // Skip if insertion fails and try another code } generatedCodes[code] = true inserted++ // Print progress every 100 codes if inserted%100 == 0 { fmt.Printf("Generated %d/%d codes\n", inserted, n) } } } // Commit the transaction if err = tx.Commit(); err != nil { return fmt.Errorf("failed to commit transaction: %w", err) } // Clear the tx to prevent the deferred rollback tx = nil fmt.Printf("Successfully generated %d unique shortcodes\n", inserted) return nil }