tidy up exports cmds

.gitignore

@@ -1,7 +1,6 @@
 .idea
 *.iml
 .env
-export/
-export_user/
+exports/
 /new
-/done.json
+/done.json

cmd/exportguild/main.go

@@ -7,16 +7,18 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"strconv"
-	"strings"
+	"sync/atomic"
+	"time"
 
+	"github.com/TicketsBot-cloud/common/encryption"
 	"github.com/TicketsBot-cloud/gdl/objects/channel/message"
+	"github.com/TicketsBot-cloud/logarchiver/internal/util"
 	"github.com/TicketsBot-cloud/logarchiver/pkg/config"
 	"github.com/TicketsBot-cloud/logarchiver/pkg/model"
 	v1 "github.com/TicketsBot-cloud/logarchiver/pkg/model/v1"
-	v22 "github.com/TicketsBot-cloud/logarchiver/pkg/model/v2"
+	v2 "github.com/TicketsBot-cloud/logarchiver/pkg/model/v2"
 	"github.com/TicketsBot-cloud/logarchiver/pkg/s3client"
-	"github.com/TicketsBot/common/encryption"
+	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
 	"golang.org/x/sync/errgroup"
@@ -27,17 +29,32 @@ const workers = 15
 var (
 	guildId       = flag.Uint64("guildid", 0, "guild id to export")
 	key           = flag.String("key", "", "aes key")
+	dbUri         = flag.String("dburi", "", "database uri")
 	ticketId      = flag.Int("ticketid", 0, "set to export a single ticket")
 	convert       = flag.Bool("convert", false, "convert to v2 if necessary")
 	userWhitelist = flag.Uint64("userwhitelist", 0, "only export tickets from this user")
 	after         = flag.Int("after", 0, "export ticket IDs above this value (inclusive)")
+	dateFromStr   = flag.String("date-from", "", "only export tickets created from this date (YYYY-MM-DD, inclusive)")
+	dateToStr     = flag.String("date-to", "", "only export tickets created up to this date (YYYY-MM-DD, inclusive)")
 )
 
 func main() {
 	flag.Parse()
 	conf := config.Parse[config.CliConfig]()
 
-	// create minio client
+	if *dbUri == "" {
+		panic("-dburi is required")
+	}
+
+	fmt.Printf("Exporting guild data for %d\n", *guildId)
+
+	exportDir := fmt.Sprintf("exports/guild/%d", *guildId)
+
+	if err := os.MkdirAll(exportDir, 0700); err != nil {
+		panic(fmt.Sprintf("failed to create export directory: %v", err))
+	}
+
+	fmt.Println("[1/3] Connecting to S3...")
 	m, err := minio.New(conf.Endpoint, &minio.Options{
 		Creds:  credentials.NewStaticV4(conf.AccessKey, conf.SecretKey, ""),
 		Secure: conf.Secure,
@@ -48,39 +65,42 @@ func main() {
 
 	client := s3client.NewS3Client(m, conf.Bucket)
 
-	// likely to be file exists
-	_ = os.Mkdir(fmt.Sprintf("export/%d", *guildId), 0)
-
+	fmt.Println("[2/3] Exporting transcripts...")
 	if ticketId != nil && *ticketId > 0 {
-		export(*ticketId, client)
+		exportTicket(*ticketId, client, exportDir)
+		printProgress(1, 1, 0, 0)
 	} else {
-		keys, err := client.GetAllKeysForGuild(context.Background(), *guildId)
-		if err != nil {
-			panic(err)
-		}
+		ticketIds := getTicketIds(*guildId, *dateFromStr, *dateToStr)
+		fmt.Printf("  Found %d tickets\n", len(ticketIds))
+
+		total := int64(len(ticketIds))
+		var processed, skipped, failed atomic.Int64
 
-		keyCh := make(chan string)
+		idCh := make(chan int)
 		go func() {
-			for _, key := range keys {
-				keyCh <- key
+			for _, id := range ticketIds {
+				idCh <- id
 			}
-
-			close(keyCh)
+			close(idCh)
 		}()
 
 		group, _ := errgroup.WithContext(context.Background())
 		for i := 0; i < workers; i++ {
 			group.Go(func() error {
-				for key := range keyCh {
-					id := key[strings.LastIndex(key, "/")+1:]
-					parsed, err := strconv.Atoi(id)
-					must(err)
-
-					if after != nil && *after > 0 && parsed < *after {
+				for id := range idCh {
+					if after != nil && *after > 0 && id < *after {
+						skipped.Add(1)
+						processed.Add(1)
+						printProgress(int(processed.Load()), int(total), int(skipped.Load()), int(failed.Load()))
 						continue
 					}
 
-					export(parsed, client)
+					if result := exportTicket(id, client, exportDir); result < 0 {
+						failed.Add(1)
+					}
+
+					processed.Add(1)
+					printProgress(int(processed.Load()), int(total), int(skipped.Load()), int(failed.Load()))
 				}
 
 				return nil
@@ -91,83 +111,162 @@ func main() {
 			panic(err)
 		}
 	}
+
+	fmt.Println("[3/3] Creating zip archive...")
+	zipPath := fmt.Sprintf("exports/guild/%d.zip", *guildId)
+	if err := util.ZipFiles(exportDir, zipPath); err != nil {
+		panic(fmt.Sprintf("could not zip files: %v", err))
+	}
+
+	os.RemoveAll(exportDir)
+	fmt.Printf("Done! Export saved to %s\n", zipPath)
 }
 
-func export(id int, client *s3client.S3Client) {
-	data, err := client.GetTicket(context.Background(), *guildId, id)
-	must(err)
+func getTicketIds(guildId uint64, fromStr, toStr string) []int {
+	pool, err := pgxpool.Connect(context.Background(), *dbUri)
+	if err != nil {
+		panic(fmt.Sprintf("failed to connect to database: %v", err))
+	}
+	defer pool.Close()
+
+	query := `SELECT id FROM tickets WHERE guild_id = $1 AND has_transcript = 't' AND open = 'f'`
+	args := []interface{}{guildId}
+	argIdx := 2
+
+	if fromStr != "" {
+		dateFrom, err := time.Parse("2006-01-02", fromStr)
+		if err != nil {
+			panic(fmt.Sprintf("invalid -date-from: %v", err))
+		}
+		query += fmt.Sprintf(` AND open_time >= $%d`, argIdx)
+		args = append(args, dateFrom)
+		argIdx++
+	}
+
+	if toStr != "" {
+		dateTo, err := time.Parse("2006-01-02", toStr)
+		if err != nil {
+			panic(fmt.Sprintf("invalid -date-to: %v", err))
+		}
+		// End of day so the date is inclusive
+		dateTo = dateTo.Add(24*time.Hour - time.Nanosecond)
+		query += fmt.Sprintf(` AND open_time <= $%d`, argIdx)
+		args = append(args, dateTo)
+	}
+
+	query += ` ORDER BY id;`
+
+	rows, err := pool.Query(context.Background(), query, args...)
+	if err != nil {
+		panic(fmt.Sprintf("failed to query tickets: %v", err))
+	}
+	defer rows.Close()
 
-	data, err = encryption.Decompress(data)
-	must(err)
+	var ids []int
+	for rows.Next() {
+		var id int
+		if err := rows.Scan(&id); err != nil {
+			panic(fmt.Sprintf("failed to scan ticket id: %v", err))
+		}
+		ids = append(ids, id)
+	}
+
+	return ids
+}
+
+func printProgress(current, total, skipped, failed int) {
+	const barWidth = 40
+	filled := barWidth * current / total
+	bar := make([]byte, barWidth)
+	for i := range bar {
+		if i < filled {
+			bar[i] = '#'
+		} else {
+			bar[i] = '-'
+		}
+	}
+	fmt.Printf("\r[%s] %d/%d transcripts (skipped: %d, failed: %d)", bar, current, total, skipped, failed)
+	if current == total {
+		fmt.Println()
+	}
+}
+
+// exportTicket exports a single ticket. Returns 1 for success, -1 for failure.
+func exportTicket(id int, client *s3client.S3Client, exportDir string) int {
+	data, err := client.GetTicket(context.Background(), *guildId, id)
+	if err != nil {
+		return -1
+	}
 
 	data, err = encryption.Decrypt([]byte(*key), data)
-	must(err)
+	if err != nil {
+		return -1
+	}
 
 	if *convert || (userWhitelist != nil && *userWhitelist > 0) {
-		var transcript v22.Transcript
+		var transcript v2.Transcript
 
 		version := model.GetVersion(data)
 		switch version {
 		case model.V1:
 			var messages []message.Message
 			if err := json.Unmarshal(data, &messages); err != nil {
-				panic(err)
+				return -1
 			}
 
 			transcript = v1.ConvertToV2(messages)
 		case model.V2:
 			if err := json.Unmarshal(data, &transcript); err != nil {
-				panic(err)
+				return -1
 			}
 		default:
-			panic(fmt.Sprintf("Unknown version %d", version))
-		}
-
-		data, err = json.Marshal(transcript)
-		must(err)
-	}
-
-	if userWhitelist != nil && *userWhitelist > 0 {
-		var transcript v22.Transcript
-		if err := json.Unmarshal(data, &transcript); err != nil {
-			panic(err)
+			return -1
 		}
 
-		transcript.Entities.Channels = nil
-		transcript.Entities.Roles = nil
+		if userWhitelist != nil && *userWhitelist > 0 {
+			transcript.Entities.Channels = nil
+			transcript.Entities.Roles = nil
 
-		user, ok := transcript.Entities.Users[*userWhitelist]
-		if !ok {
-			transcript.Entities.Users = nil
-		} else {
-			transcript.Entities.Users = map[uint64]v22.User{
-				user.Id: user,
+			user, ok := transcript.Entities.Users[*userWhitelist]
+			if !ok {
+				transcript.Entities.Users = nil
+			} else {
+				transcript.Entities.Users = map[uint64]v2.User{
+					user.Id: user,
+				}
 			}
-		}
 
-		var messages []v22.Message
-		for _, message := range transcript.Messages {
-			if message.AuthorId == *userWhitelist {
-				messages = append(messages, message)
+			var messages []v2.Message
+			for _, message := range transcript.Messages {
+				if message.AuthorId == *userWhitelist {
+					messages = append(messages, message)
+				}
 			}
-		}
 
-		transcript.Messages = messages
+			transcript.Messages = messages
+		}
 
 		data, err = json.Marshal(transcript)
-		must(err)
+		if err != nil {
+			return -1
+		}
 	}
 
 	var encoded bytes.Buffer
-	must(json.Indent(&encoded, data, "", "  "))
+	if err := json.Indent(&encoded, data, "", "  "); err != nil {
+		return -1
+	}
 
-	f, err := os.Create(fmt.Sprintf("export/%d/%d.json", *guildId, id))
-	must(err)
+	f, err := os.Create(fmt.Sprintf("%s/%d.json", exportDir, id))
+	if err != nil {
+		return -1
+	}
 
-	_, err = encoded.WriteTo(f)
-	must(err)
+	if _, err := encoded.WriteTo(f); err != nil {
+		return -1
+	}
 
-	fmt.Printf("exported %d\n", id)
+	return 1
 }
 
 func must(err error) {