mirror of
https://github.com/warmbly/warmbly.git
synced 2026-09-06 08:01:24 +00:00
599 lines
22 KiB
Go
599 lines
22 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"github.com/warmbly/warmbly/internal/models"
|
|
)
|
|
|
|
// WorkerLivenessWindow is how stale a worker's heartbeat may be before
|
|
// placement stops considering it. Workers heartbeat every 30s, so this
|
|
// tolerates a handful of missed beats.
|
|
//
|
|
// Without it, rows for workers that no longer exist stay active and healthy
|
|
// forever and get selected, and the command sent to them times out with no
|
|
// consumer. That is easy to hit: a worker with no WORKER_ID set generates a
|
|
// fresh UUID on every boot, so each container recreate leaves another dead row
|
|
// behind.
|
|
const WorkerLivenessWindow = "10 minutes"
|
|
|
|
// EmailAccountWorkerInfo contains worker info for an email account
|
|
type EmailAccountWorkerInfo struct {
|
|
EmailAccountID uuid.UUID
|
|
WorkerID *uuid.UUID
|
|
UserID uuid.UUID
|
|
FreeTier *bool
|
|
}
|
|
|
|
// EmailAccountPlacementHint carries the small slice of mailbox metadata the
|
|
// assignment service needs to compute a mailbox weight (provider +
|
|
// warmup-only flag). Fetched once per assignment so we don't have to thread
|
|
// the data through every caller.
|
|
type EmailAccountPlacementHint struct {
|
|
Provider string
|
|
IsWarmup bool
|
|
}
|
|
|
|
type WorkerRepository interface {
|
|
// Worker queries
|
|
GetByID(ctx context.Context, id uuid.UUID) (*models.Worker, error)
|
|
GetSharedWorkersByTier(ctx context.Context, freeTier bool) ([]models.Worker, error)
|
|
// IsWorkerLive reports whether a worker is active and heartbeating inside
|
|
// WorkerLivenessWindow, i.e. whether it can still receive commands.
|
|
IsWorkerLive(ctx context.Context, id uuid.UUID) (bool, error)
|
|
GetAllActiveWorkers(ctx context.Context) ([]models.Worker, error)
|
|
GetAvailableDedicatedWorker(ctx context.Context) (*models.Worker, error)
|
|
PromoteIdlePremiumWorkerToDedicated(ctx context.Context) (*models.Worker, error)
|
|
IncrementAccountCount(ctx context.Context, workerID uuid.UUID) error
|
|
DecrementAccountCount(ctx context.Context, workerID uuid.UUID) error
|
|
SetWorkerType(ctx context.Context, workerID uuid.UUID, workerType models.WorkerType) error
|
|
|
|
// Dedicated worker assignments
|
|
CreateDedicatedAssignment(ctx context.Context, assignment *models.DedicatedWorkerAssignment) error
|
|
CreateDedicatedAssignmentIfNotExists(ctx context.Context, assignment *models.DedicatedWorkerAssignment) (bool, error)
|
|
GetActiveDedicatedAssignment(ctx context.Context, userID uuid.UUID) (*models.DedicatedWorkerAssignment, error)
|
|
GetDedicatedWorkerByOrgID(ctx context.Context, orgID uuid.UUID) (*models.Worker, error)
|
|
ReleaseDedicatedAssignment(ctx context.Context, userID uuid.UUID) error
|
|
|
|
// Email account worker queries
|
|
GetEmailAccountsByWorkerID(ctx context.Context, workerID uuid.UUID) ([]uuid.UUID, error)
|
|
GetEmailAccountsByUserID(ctx context.Context, userID uuid.UUID) ([]uuid.UUID, error)
|
|
GetEmailAccountsByOrganizationID(ctx context.Context, orgID uuid.UUID) ([]uuid.UUID, error)
|
|
GetEmailAccountWorkerInfo(ctx context.Context, emailAccountID uuid.UUID) (*EmailAccountWorkerInfo, error)
|
|
UpdateEmailAccountWorker(ctx context.Context, emailAccountID, workerID uuid.UUID) error
|
|
ClearEmailAccountWorker(ctx context.Context, emailAccountID uuid.UUID) error
|
|
UpdateEmailAccountWarmupPoolType(ctx context.Context, emailAccountID uuid.UUID, poolType string) error
|
|
|
|
// SSH-managed workers (admin-driven lifecycle)
|
|
CreateWorker(ctx context.Context, in CreateWorkerInput) error
|
|
GetWorkerDetail(ctx context.Context, id uuid.UUID) (*models.Worker, error)
|
|
ListWorkersDetail(ctx context.Context) ([]models.Worker, error)
|
|
GetWorkerSSHCredentials(ctx context.Context, id uuid.UUID) (*models.WorkerSSHCredentials, error)
|
|
UpdateInstallState(ctx context.Context, id uuid.UUID, state models.WorkerInstallState, lastError string) error
|
|
UpdateLastSeen(ctx context.Context, id uuid.UUID, at time.Time) error
|
|
UpdateHostFingerprint(ctx context.Context, id uuid.UUID, fingerprint string) error
|
|
RotateSSHKey(ctx context.Context, id uuid.UUID, publicKey, privateKeyEncrypted string) error
|
|
DeleteWorker(ctx context.Context, id uuid.UUID) error
|
|
ConsumeEnrollmentToken(ctx context.Context, tokenHash string) (*models.Worker, error)
|
|
RecordEnrolledIP(ctx context.Context, id uuid.UUID, ip string) error
|
|
AssignWorkerProfile(ctx context.Context, workerID uuid.UUID, profileID *uuid.UUID) error
|
|
MarkConfigApplied(ctx context.Context, workerID uuid.UUID, at time.Time) error
|
|
MarkImageVersion(ctx context.Context, workerID uuid.UUID, version string) error
|
|
ListWorkersByProfile(ctx context.Context, profileID uuid.UUID) ([]models.Worker, error)
|
|
|
|
// Threat-level segregation
|
|
SetWorkerRiskPool(ctx context.Context, workerID uuid.UUID, pool models.WorkerRiskPool) error
|
|
SetEmailAccountRiskBand(ctx context.Context, emailAccountID uuid.UUID, band models.EmailRiskBand) error
|
|
GetEmailAccountRiskBand(ctx context.Context, emailAccountID uuid.UUID) (models.EmailRiskBand, error)
|
|
GetSharedWorkersByTierAndPool(ctx context.Context, freeTier bool, pool models.WorkerRiskPool) ([]models.Worker, error)
|
|
PromoteWorkerToPool(ctx context.Context, freeTier bool, target models.WorkerRiskPool) (*models.Worker, error)
|
|
ListRiskCandidates(ctx context.Context, limit int) ([]RiskCandidate, error)
|
|
|
|
// Tags
|
|
GetWorkerTags(ctx context.Context, workerID uuid.UUID) ([]string, error)
|
|
SetWorkerTags(ctx context.Context, workerID uuid.UUID, tags []string) error
|
|
ListAllWorkerTags(ctx context.Context) ([]string, error)
|
|
HydrateWorkerTags(ctx context.Context, workers []*models.Worker) error
|
|
|
|
// Heartbeat (auto-registers new workers on first contact)
|
|
UpsertOnHeartbeat(ctx context.Context, id uuid.UUID, ipAddr, tier, egressKind string) error
|
|
DeactivateWorker(ctx context.Context, id uuid.UUID) error
|
|
|
|
// Health and capacity
|
|
InsertWorkerHealthSample(ctx context.Context, sample *models.WorkerHealthSample) error
|
|
ListCapacityCandidates(ctx context.Context, freeTier bool, allowedStates []models.WorkerHealthState) ([]WorkerCapacityRowDB, error)
|
|
GetCapacityRow(ctx context.Context, workerID uuid.UUID) (*WorkerCapacityRowDB, error)
|
|
AddLoadScore(ctx context.Context, workerID uuid.UUID, delta float64) error
|
|
SetWorkerHealthState(ctx context.Context, workerID uuid.UUID, state models.WorkerHealthState) error
|
|
SetWorkerEgressKind(ctx context.Context, workerID uuid.UUID, kind models.WorkerEgressKind) error
|
|
RefreshWorkerCapacityView(ctx context.Context) error
|
|
GetEmailAccountPlacementHint(ctx context.Context, emailAccountID uuid.UUID) (*EmailAccountPlacementHint, error)
|
|
}
|
|
|
|
type workerRepository struct {
|
|
db *pgxpool.Pool
|
|
}
|
|
|
|
func NewWorkerRepository(db *pgxpool.Pool) WorkerRepository {
|
|
return &workerRepository{db: db}
|
|
}
|
|
|
|
func (r *workerRepository) GetByID(ctx context.Context, id uuid.UUID) (*models.Worker, error) {
|
|
query := `
|
|
SELECT id, ip_addr, active, free_tier, worker_type, account_count, created_at, updated_at
|
|
FROM workers
|
|
WHERE id = $1
|
|
`
|
|
|
|
var w models.Worker
|
|
err := r.db.QueryRow(ctx, query, id).Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
)
|
|
if err == pgx.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &w, nil
|
|
}
|
|
|
|
// GetSharedWorkersByTier retrieves all shared workers matching the specified tier
|
|
// Results are sorted by account_count ASC (least loaded first) for even distribution
|
|
// IsWorkerLive reports whether a worker is still a valid placement target,
|
|
// using the same predicate as selection: active, and heartbeating inside
|
|
// WorkerLivenessWindow. A mailbox already assigned to a worker that fails this
|
|
// is orphaned, because its send commands go to a topic with no consumer.
|
|
func (r *workerRepository) IsWorkerLive(ctx context.Context, id uuid.UUID) (bool, error) {
|
|
query := `
|
|
SELECT EXISTS(
|
|
SELECT 1 FROM workers
|
|
WHERE id = $1
|
|
AND active = true
|
|
AND last_seen_at > now() - $2::interval
|
|
)
|
|
`
|
|
var live bool
|
|
if err := r.db.QueryRow(ctx, query, id, WorkerLivenessWindow).Scan(&live); err != nil {
|
|
return false, err
|
|
}
|
|
return live, nil
|
|
}
|
|
|
|
func (r *workerRepository) GetSharedWorkersByTier(ctx context.Context, freeTier bool) ([]models.Worker, error) {
|
|
query := `
|
|
SELECT id, ip_addr, active, free_tier, worker_type, account_count, created_at, updated_at
|
|
FROM workers
|
|
WHERE worker_type = 'shared'
|
|
AND active = true
|
|
AND free_tier = $1
|
|
AND last_seen_at > now() - $2::interval
|
|
ORDER BY account_count ASC
|
|
`
|
|
|
|
rows, err := r.db.Query(ctx, query, freeTier, WorkerLivenessWindow)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var workers []models.Worker
|
|
for rows.Next() {
|
|
var w models.Worker
|
|
if err := rows.Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
workers = append(workers, w)
|
|
}
|
|
|
|
return workers, rows.Err()
|
|
}
|
|
|
|
// GetAllActiveWorkers retrieves all active workers regardless of tier
|
|
func (r *workerRepository) GetAllActiveWorkers(ctx context.Context) ([]models.Worker, error) {
|
|
query := `
|
|
SELECT id, ip_addr, active, free_tier, worker_type, account_count, created_at, updated_at
|
|
FROM workers
|
|
WHERE active = true
|
|
ORDER BY created_at
|
|
`
|
|
rows, err := r.db.Query(ctx, query)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var workers []models.Worker
|
|
for rows.Next() {
|
|
var w models.Worker
|
|
if err := rows.Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
workers = append(workers, w)
|
|
}
|
|
return workers, rows.Err()
|
|
}
|
|
|
|
// GetAvailableDedicatedWorker finds an available dedicated worker that is not assigned
|
|
func (r *workerRepository) GetAvailableDedicatedWorker(ctx context.Context) (*models.Worker, error) {
|
|
query := `
|
|
SELECT w.id, w.ip_addr, w.active, w.free_tier, w.worker_type, w.account_count, w.created_at, w.updated_at
|
|
FROM workers w
|
|
WHERE w.worker_type = 'dedicated'
|
|
AND w.active = true
|
|
AND w.free_tier = false
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM dedicated_worker_assignments dwa
|
|
WHERE dwa.worker_id = w.id AND dwa.released_at IS NULL
|
|
)
|
|
ORDER BY w.created_at ASC
|
|
LIMIT 1
|
|
`
|
|
|
|
var w models.Worker
|
|
err := r.db.QueryRow(ctx, query).Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
)
|
|
if err == pgx.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &w, nil
|
|
}
|
|
|
|
// PromoteIdlePremiumWorkerToDedicated flips the oldest idle premium shared
|
|
// worker to dedicated and returns it, so the control plane can seed dedicated
|
|
// capacity on demand when GetAvailableDedicatedWorker finds none free. Only
|
|
// premium (free_tier = false), shared, active workers with account_count = 0
|
|
// are eligible — promoting a loaded worker would strand its mailboxes on a box
|
|
// that suddenly belongs to one org. Returns nil if no idle premium worker
|
|
// exists.
|
|
//
|
|
// The candidate is selected and flipped in a single statement
|
|
// (UPDATE ... WHERE id = (SELECT ... FOR UPDATE SKIP LOCKED)), so two
|
|
// concurrent promotions never pick the same row. free/premium capacity
|
|
// separation is preserved: free-tier workers are never promoted.
|
|
func (r *workerRepository) PromoteIdlePremiumWorkerToDedicated(ctx context.Context) (*models.Worker, error) {
|
|
query := `
|
|
UPDATE workers SET worker_type = 'dedicated', updated_at = NOW()
|
|
WHERE id = (
|
|
SELECT id FROM workers
|
|
WHERE worker_type = 'shared'
|
|
AND active = true
|
|
AND free_tier = false
|
|
AND account_count = 0
|
|
ORDER BY created_at ASC
|
|
LIMIT 1
|
|
FOR UPDATE SKIP LOCKED
|
|
)
|
|
RETURNING id, ip_addr, active, free_tier, worker_type, account_count, risk_pool, created_at, updated_at
|
|
`
|
|
var w models.Worker
|
|
err := r.db.QueryRow(ctx, query).Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount, &w.RiskPool,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
)
|
|
if err == pgx.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &w, nil
|
|
}
|
|
|
|
func (r *workerRepository) IncrementAccountCount(ctx context.Context, workerID uuid.UUID) error {
|
|
query := `UPDATE workers SET account_count = account_count + 1, updated_at = NOW() WHERE id = $1`
|
|
_, err := r.db.Exec(ctx, query, workerID)
|
|
return err
|
|
}
|
|
|
|
func (r *workerRepository) DecrementAccountCount(ctx context.Context, workerID uuid.UUID) error {
|
|
query := `UPDATE workers SET account_count = GREATEST(account_count - 1, 0), updated_at = NOW() WHERE id = $1`
|
|
_, err := r.db.Exec(ctx, query, workerID)
|
|
return err
|
|
}
|
|
|
|
func (r *workerRepository) SetWorkerType(ctx context.Context, workerID uuid.UUID, workerType models.WorkerType) error {
|
|
query := `UPDATE workers SET worker_type = $1, updated_at = NOW() WHERE id = $2`
|
|
_, err := r.db.Exec(ctx, query, workerType, workerID)
|
|
return err
|
|
}
|
|
|
|
// CreateDedicatedAssignment creates a new dedicated worker assignment
|
|
func (r *workerRepository) CreateDedicatedAssignment(ctx context.Context, assignment *models.DedicatedWorkerAssignment) error {
|
|
query := `
|
|
INSERT INTO dedicated_worker_assignments (id, worker_id, organization_id, subscription_id, assigned_at)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
`
|
|
|
|
_, err := r.db.Exec(ctx, query,
|
|
assignment.ID,
|
|
assignment.WorkerID,
|
|
assignment.OrganizationID,
|
|
assignment.SubscriptionID,
|
|
assignment.AssignedAt,
|
|
)
|
|
return err
|
|
}
|
|
|
|
// CreateDedicatedAssignmentIfNotExists atomically creates a dedicated worker assignment
|
|
// only if no active (released_at IS NULL) assignment exists for the organization.
|
|
// Returns (true, nil) if created, (false, nil) if already exists.
|
|
func (r *workerRepository) CreateDedicatedAssignmentIfNotExists(ctx context.Context, assignment *models.DedicatedWorkerAssignment) (bool, error) {
|
|
query := `
|
|
INSERT INTO dedicated_worker_assignments (id, worker_id, organization_id, subscription_id, assigned_at)
|
|
SELECT $1, $2, $3, $4, $5
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM dedicated_worker_assignments
|
|
WHERE organization_id = $3 AND released_at IS NULL
|
|
)
|
|
`
|
|
result, err := r.db.Exec(ctx, query,
|
|
assignment.ID,
|
|
assignment.WorkerID,
|
|
assignment.OrganizationID,
|
|
assignment.SubscriptionID,
|
|
assignment.AssignedAt,
|
|
)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return result.RowsAffected() > 0, nil
|
|
}
|
|
|
|
// GetActiveDedicatedAssignment retrieves the active dedicated assignment for an organization
|
|
func (r *workerRepository) GetActiveDedicatedAssignment(ctx context.Context, userID uuid.UUID) (*models.DedicatedWorkerAssignment, error) {
|
|
query := `
|
|
SELECT id, worker_id, organization_id, subscription_id, assigned_at, released_at
|
|
FROM dedicated_worker_assignments
|
|
WHERE organization_id = $1 AND released_at IS NULL
|
|
`
|
|
|
|
var a models.DedicatedWorkerAssignment
|
|
err := r.db.QueryRow(ctx, query, userID).Scan(
|
|
&a.ID, &a.WorkerID, &a.OrganizationID, &a.SubscriptionID, &a.AssignedAt, &a.ReleasedAt,
|
|
)
|
|
if err == pgx.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &a, nil
|
|
}
|
|
|
|
// GetDedicatedWorkerByOrgID retrieves the dedicated worker assigned to an organization
|
|
func (r *workerRepository) GetDedicatedWorkerByOrgID(ctx context.Context, orgID uuid.UUID) (*models.Worker, error) {
|
|
query := `
|
|
SELECT w.id, w.ip_addr, w.active, w.free_tier, w.worker_type, w.account_count, w.created_at, w.updated_at
|
|
FROM workers w
|
|
JOIN dedicated_worker_assignments dwa ON w.id = dwa.worker_id
|
|
WHERE dwa.organization_id = $1 AND dwa.released_at IS NULL
|
|
`
|
|
|
|
var w models.Worker
|
|
err := r.db.QueryRow(ctx, query, orgID).Scan(
|
|
&w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount,
|
|
&w.CreatedAt, &w.UpdatedAt,
|
|
)
|
|
if err == pgx.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &w, nil
|
|
}
|
|
|
|
// ReleaseDedicatedAssignment releases a dedicated worker assignment
|
|
func (r *workerRepository) ReleaseDedicatedAssignment(ctx context.Context, userID uuid.UUID) error {
|
|
query := `
|
|
UPDATE dedicated_worker_assignments
|
|
SET released_at = $1
|
|
WHERE organization_id = $2 AND released_at IS NULL
|
|
`
|
|
|
|
_, err := r.db.Exec(ctx, query, time.Now(), userID)
|
|
return err
|
|
}
|
|
|
|
// GetEmailAccountsByWorkerID retrieves all email account IDs assigned to a worker
|
|
func (r *workerRepository) GetEmailAccountsByWorkerID(ctx context.Context, workerID uuid.UUID) ([]uuid.UUID, error) {
|
|
query := `SELECT id FROM email_accounts WHERE worker_id = $1`
|
|
|
|
rows, err := r.db.Query(ctx, query, workerID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var ids []uuid.UUID
|
|
for rows.Next() {
|
|
var id uuid.UUID
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
|
|
return ids, rows.Err()
|
|
}
|
|
|
|
// GetEmailAccountsByUserID retrieves all email account IDs for a user
|
|
func (r *workerRepository) GetEmailAccountsByUserID(ctx context.Context, userID uuid.UUID) ([]uuid.UUID, error) {
|
|
query := `SELECT id FROM email_accounts WHERE user_id = $1`
|
|
|
|
rows, err := r.db.Query(ctx, query, userID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var ids []uuid.UUID
|
|
for rows.Next() {
|
|
var id uuid.UUID
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
|
|
return ids, rows.Err()
|
|
}
|
|
|
|
// GetEmailAccountsByOrganizationID retrieves all email account IDs for an organization
|
|
func (r *workerRepository) GetEmailAccountsByOrganizationID(ctx context.Context, orgID uuid.UUID) ([]uuid.UUID, error) {
|
|
query := `SELECT id FROM email_accounts WHERE organization_id = $1`
|
|
|
|
rows, err := r.db.Query(ctx, query, orgID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var ids []uuid.UUID
|
|
for rows.Next() {
|
|
var id uuid.UUID
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
|
|
return ids, rows.Err()
|
|
}
|
|
|
|
// UpdateEmailAccountWorker assigns a worker to an email account
|
|
func (r *workerRepository) UpdateEmailAccountWorker(ctx context.Context, emailAccountID, workerID uuid.UUID) error {
|
|
query := `UPDATE email_accounts SET worker_id = $1, updated_at = NOW() WHERE id = $2`
|
|
_, err := r.db.Exec(ctx, query, workerID, emailAccountID)
|
|
return err
|
|
}
|
|
|
|
// ClearEmailAccountWorker removes worker assignment from an email account
|
|
func (r *workerRepository) ClearEmailAccountWorker(ctx context.Context, emailAccountID uuid.UUID) error {
|
|
query := `UPDATE email_accounts SET worker_id = NULL, updated_at = NOW() WHERE id = $1`
|
|
_, err := r.db.Exec(ctx, query, emailAccountID)
|
|
return err
|
|
}
|
|
|
|
// UpdateEmailAccountWarmupPoolType updates the warmup pool type for an email account
|
|
func (r *workerRepository) UpdateEmailAccountWarmupPoolType(ctx context.Context, emailAccountID uuid.UUID, poolType string) error {
|
|
query := `UPDATE email_accounts SET warmup_pool_type = $1, updated_at = NOW() WHERE id = $2`
|
|
_, err := r.db.Exec(ctx, query, poolType, emailAccountID)
|
|
return err
|
|
}
|
|
|
|
// GetEmailAccountWorkerInfo retrieves worker info for an email account
|
|
func (r *workerRepository) GetEmailAccountWorkerInfo(ctx context.Context, emailAccountID uuid.UUID) (*EmailAccountWorkerInfo, error) {
|
|
query := `
|
|
SELECT ea.id, ea.worker_id, ea.user_id, w.free_tier
|
|
FROM email_accounts ea
|
|
LEFT JOIN workers w ON ea.worker_id = w.id
|
|
WHERE ea.id = $1
|
|
`
|
|
|
|
var info EmailAccountWorkerInfo
|
|
err := r.db.QueryRow(ctx, query, emailAccountID).Scan(
|
|
&info.EmailAccountID, &info.WorkerID, &info.UserID, &info.FreeTier,
|
|
)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &info, nil
|
|
}
|
|
|
|
// DisableWarmupByUserID disables warmup for all email accounts belonging to a user
|
|
func DisableWarmupByUserID(ctx context.Context, db *pgxpool.Pool, userID uuid.UUID) error {
|
|
query := `UPDATE email_accounts SET warmup = NULL, updated_at = NOW() WHERE user_id = $1`
|
|
_, err := db.Exec(ctx, query, userID)
|
|
return err
|
|
}
|
|
|
|
// PauseCampaignsByUserID pauses all active campaigns for a user
|
|
func PauseCampaignsByUserID(ctx context.Context, db *pgxpool.Pool, userID uuid.UUID, reason string) error {
|
|
query := `UPDATE campaigns SET status = $1, updated_at = NOW() WHERE user_id = $2 AND status = 'active'`
|
|
_, err := db.Exec(ctx, query, reason, userID)
|
|
return err
|
|
}
|
|
|
|
// GetExpiredTrialsWithoutPayment retrieves subscriptions with expired free trials and no paid subscription
|
|
func GetExpiredTrialsWithoutPayment(ctx context.Context, db *pgxpool.Pool) ([]models.Subscription, error) {
|
|
query := `
|
|
SELECT s.id, s.user_id, s.organization_id, s.plan_id, s.stripe_customer_id, s.stripe_subscription_id,
|
|
s.stripe_price_id, s.status, s.current_period_start, s.current_period_end,
|
|
s.cancel_at_period_end, s.canceled_at, s.trial_start, s.trial_end,
|
|
s.free_trial_started_at, s.free_trial_ends_at, s.is_enterprise, s.created_at, s.updated_at,
|
|
u.email
|
|
FROM subscriptions s
|
|
LEFT JOIN users u ON s.user_id = u.id
|
|
WHERE s.free_trial_ends_at IS NOT NULL
|
|
AND s.free_trial_ends_at < NOW()
|
|
AND s.stripe_subscription_id IS NULL
|
|
AND s.status NOT IN ('canceled', 'incomplete_expired')
|
|
`
|
|
|
|
rows, err := db.Query(ctx, query)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var subs []models.Subscription
|
|
for rows.Next() {
|
|
var s models.Subscription
|
|
if err := rows.Scan(
|
|
&s.ID, &s.UserID, &s.OrganizationID, &s.PlanID, &s.StripeCustomerID, &s.StripeSubscriptionID,
|
|
&s.StripePriceID, &s.Status, &s.CurrentPeriodStart, &s.CurrentPeriodEnd,
|
|
&s.CancelAtPeriodEnd, &s.CanceledAt, &s.TrialStart, &s.TrialEnd,
|
|
&s.FreeTrialStartedAt, &s.FreeTrialEndsAt, &s.IsEnterprise, &s.CreatedAt, &s.UpdatedAt,
|
|
&s.UserEmail,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
subs = append(subs, s)
|
|
}
|
|
|
|
return subs, rows.Err()
|
|
}
|
|
|
|
// PauseCampaignsByOrganizationID pauses all active campaigns for an organization
|
|
func PauseCampaignsByOrganizationID(ctx context.Context, db *pgxpool.Pool, orgID uuid.UUID, reason string) error {
|
|
query := `UPDATE campaigns SET status = $1, updated_at = NOW() WHERE organization_id = $2 AND status = 'active'`
|
|
_, err := db.Exec(ctx, query, reason, orgID)
|
|
return err
|
|
}
|
|
|
|
// DisableWarmupByOrganizationID disables warmup for all email accounts belonging to an organization
|
|
func DisableWarmupByOrganizationID(ctx context.Context, db *pgxpool.Pool, orgID uuid.UUID) error {
|
|
query := `UPDATE email_accounts SET warmup = NULL, updated_at = NOW() WHERE organization_id = $1`
|
|
_, err := db.Exec(ctx, query, orgID)
|
|
return err
|
|
}
|
|
|
|
// MarkSubscriptionTrialExpired marks a subscription as expired trial
|
|
func MarkSubscriptionTrialExpired(ctx context.Context, db *pgxpool.Pool, subID uuid.UUID) error {
|
|
query := `UPDATE subscriptions SET status = 'incomplete_expired', updated_at = NOW() WHERE id = $1`
|
|
_, err := db.Exec(ctx, query, subID)
|
|
return err
|
|
}
|