Updated owner management

bot.go

@@ -63,6 +63,10 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) {
 		}
 
 		if err := db.Create(&owner).Error; err != nil {
+			// If unique constraint is violated, another owner already exists
+			if strings.Contains(err.Error(), "unique index") {
+				return nil, fmt.Errorf("an owner already exists for this bot")
+			}
 			return nil, fmt.Errorf("failed to create owner user: %w", err)
 		}
 	} else if err != nil {
@@ -123,6 +127,10 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
 			}
 
 			if err := b.db.Create(&user).Error; err != nil {
+				// Handle unique constraint for owner
+				if isOwner && strings.Contains(err.Error(), "unique index") {
+					return User{}, fmt.Errorf("an owner already exists for this bot")
+				}
 				return User{}, err
 			}
 		} else {
@@ -139,11 +147,11 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
 				return User{}, err
 			}
 			// Promote to owner
-			user.Role, err = b.getRoleByName("owner")
+			role, err := b.getRoleByName("owner")
 			if err != nil {
 				return User{}, err
 			}
-			user.RoleID = user.Role.ID
+			user.RoleID = role.ID
 			user.IsOwner = true
 			if err := b.db.Save(&user).Error; err != nil {
 				return User{}, err

database.go

@@ -27,14 +27,21 @@ func initDB() (*gorm.DB, error) {
 		return nil, fmt.Errorf("failed to connect to database: %w", err)
 	}
 
-	// AutoMigrate with unique constraint for owners
+	// AutoMigrate the models
 	err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
 	if err != nil {
 		return nil, fmt.Errorf("failed to migrate database schema: %w", err)
 	}
 
-	// Add unique index for owners per bot
+	// Enforce unique owner per bot using raw SQL
-	db.SetupJoinTable(&BotModel{}, "Users", &User{})
+	// Note: SQLite doesn't support partial indexes, but we can simulate it by making a unique index on (BotID, IsOwner)
+	// and ensuring that IsOwner can only be true for one user per BotID.
+	// This approach allows multiple users with IsOwner=false for the same BotID,
+	// but only one user can have IsOwner=true per BotID.
+	err = db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_bot_owner ON users (bot_id, is_owner) WHERE is_owner = 1;`).Error
+	if err != nil {
+		return nil, fmt.Errorf("failed to create unique index for bot owners: %w", err)
+	}
 
 	err = createDefaultRoles(db)
 	if err != nil {

handlers.go

@@ -84,7 +84,10 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
 
 	userMessage := b.createMessage(chatID, userID, username, user.Role.Name, text, true)
 	userMessage.UserRole = string(anthropic.RoleUser) // Convert to string
-	b.storeMessage(userMessage)
+	if err := b.storeMessage(userMessage); err != nil {
+		log.Printf("Error storing user message: %v", err)
+		return
+	}
 
 	chatMemory := b.getOrCreateChatMemory(chatID)
 	b.addMessageToChatMemory(chatMemory, userMessage)
@@ -98,10 +101,16 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
 		response = "I'm sorry, I'm having trouble processing your request right now."
 	}
 
-	b.sendResponse(ctx, chatID, response, businessConnectionID)
+	if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
+		log.Printf("Error sending response: %v", err)
+		return
+	}
 
-	assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false)
+	assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false)
-	b.storeMessage(assistantMessage)
+	if err := b.storeMessage(assistantMessage); err != nil {
+		log.Printf("Error storing assistant message: %v", err)
+		return
+	}
 	b.addMessageToChatMemory(chatMemory, assistantMessage)
 }
 

models.go

@@ -61,3 +61,8 @@ type User struct {
 	Role       Role `gorm:"foreignKey:RoleID"`
 	IsOwner    bool `gorm:"default:false"` // Indicates if the user is the owner
 }
+
+// Compound unique index to ensure only one owner per bot
+func (User) TableName() string {
+	return "users"
+}

rate_limiter_test.go

@@ -1,8 +1,12 @@
 package main
 
 import (
+	"strings"
 	"testing"
 	"time"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
 )
 
 // TestCheckRateLimits tests the checkRateLimits method of the Bot.
@@ -80,6 +84,109 @@ func TestCheckRateLimits(t *testing.T) {
 	}
 }
 
+func TestOwnerAssignment(t *testing.T) {
+	// Initialize in-memory database for testing
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	if err != nil {
+		t.Fatalf("Failed to open in-memory database: %v", err)
+	}
+
+	// Migrate the schema
+	err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
+	if err != nil {
+		t.Fatalf("Failed to migrate database schema: %v", err)
+	}
+
+	// Create default roles
+	err = createDefaultRoles(db)
+	if err != nil {
+		t.Fatalf("Failed to create default roles: %v", err)
+	}
+
+	// Create a bot configuration
+	config := BotConfig{
+		ID:              "test_bot",
+		TelegramToken:   "TEST_TELEGRAM_TOKEN",
+		MemorySize:      10,
+		MessagePerHour:  5,
+		MessagePerDay:   10,
+		TempBanDuration: "1m",
+		SystemPrompts:   make(map[string]string),
+		Active:          true,
+		OwnerTelegramID: 111111111,
+	}
+
+	// Initialize MockClock
+	mockClock := &MockClock{
+		currentTime: time.Now(),
+	}
+
+	// Create the bot
+	bot, err := NewBot(db, config, mockClock)
+	if err != nil {
+		t.Fatalf("Failed to create bot: %v", err)
+	}
+
+	// Verify that the owner exists
+	var owner User
+	err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
+	if err != nil {
+		t.Fatalf("Owner was not created: %v", err)
+	}
+
+	// Attempt to create another owner for the same bot
+	_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
+	if err == nil {
+		t.Fatalf("Expected error when creating a second owner, but got none")
+	}
+
+	// Verify that the error message is appropriate
+	expectedErrorMsg := "an owner already exists for this bot"
+	if err.Error() != expectedErrorMsg && !strings.Contains(err.Error(), "unique index") {
+		t.Fatalf("Unexpected error message: %v", err)
+	}
+
+	// Assign admin role to a new user
+	adminUser, err := bot.getOrCreateUser(333333333, "AdminUser", false)
+	if err != nil {
+		t.Fatalf("Failed to create admin user: %v", err)
+	}
+
+	if adminUser.Role.Name != "admin" {
+		t.Fatalf("Expected role 'admin', got '%s'", adminUser.Role.Name)
+	}
+
+	// Assign owner role to a user from a different bot
+	otherBotConfig := BotConfig{
+		ID:              "other_bot",
+		TelegramToken:   "OTHER_TELEGRAM_TOKEN",
+		MemorySize:      10,
+		MessagePerHour:  5,
+		MessagePerDay:   10,
+		TempBanDuration: "1m",
+		SystemPrompts:   make(map[string]string),
+		Active:          true,
+		OwnerTelegramID: 444444444,
+	}
+
+	otherBot, err := NewBot(db, otherBotConfig, mockClock)
+	if err != nil {
+		t.Fatalf("Failed to create other bot: %v", err)
+	}
+
+	_, err = otherBot.getOrCreateUser(config.OwnerTelegramID, "OwnerOfOtherBot", true)
+	if err != nil {
+		t.Fatalf("Failed to assign existing owner to another bot: %v", err)
+	}
+
+	// Verify multiple bots can have the same owner telegram ID
+	var ownerOfOtherBot User
+	err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, otherBot.botID, true).First(&ownerOfOtherBot).Error
+	if err != nil {
+		t.Fatalf("Owner of other bot was not created: %v", err)
+	}
+}
+
 // To ensure thread safety and avoid race conditions during testing,
 // you can run the tests with the `-race` flag:
 // go test -race -v