fix: refactor 1

Author: michiosw

internal/backend/database/database.go

@@ -24,7 +24,7 @@ type DBInterface interface {
 	UpdateServerDetails(id int64, name, url string) error
 
 	// Tool Management
-	UpsertTool(tool models.Tool) (added bool, err error)
+	UpsertTool(tool models.Tool) error
 	ListTools() ([]models.Tool, error)
 	ListToolsByServerID(serverID int64) ([]models.Tool, error)
 	RemoveToolsByServerID(serverID int64) error
@@ -80,7 +80,7 @@ func (db *DB) ensureSchema() error {
 		name TEXT NOT NULL,
 		description TEXT,
 		created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-		updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, -- Added this column
+		updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
 		FOREIGN KEY (source_server_id) REFERENCES mcp_servers(id) ON DELETE CASCADE,
 		UNIQUE (source_server_id, external_id)
 	);
@@ -156,7 +156,7 @@ func (db *DB) RemoveServer(id int64) error {
 		return fmt.Errorf("failed to get rows affected for remove server ID %d: %w", id, err)
 	}
 	if rowsAffected == 0 {
-		return fmt.Errorf("server with ID %d not found for removal", id)
+		return nil
 	}
 	return nil
 }
@@ -173,7 +173,6 @@ func (db *DB) UpdateServerStatus(id int64, state models.ConnectionState, lastErr
 
 // UpdateServerDetails updates the name and URL for a server.
 func (db *DB) UpdateServerDetails(id int64, name, url string) error {
-	// Add validation if needed (e.g., check for URL uniqueness if changing)
 	query := `UPDATE mcp_servers SET name = ?, url = ? WHERE id = ?`
 	_, err := db.Exec(query, name, url, id)
 	if err != nil {
@@ -185,46 +184,36 @@ func (db *DB) UpdateServerDetails(id int64, name, url string) error {
 // --- Tool CRUD ---
 
 // UpsertTool inserts a new tool or updates an existing one based on external_id and source_server_id.
-// It returns true if a new row was inserted, false if an existing row was updated.
-func (db *DB) UpsertTool(tool models.Tool) (added bool, err error) {
+func (db *DB) UpsertTool(tool models.Tool) (err error) {
 	tx, err := db.Beginx()
 	if err != nil {
-		return false, fmt.Errorf("failed to begin transaction for upsert tool: %w", err)
+		return fmt.Errorf("failed to begin transaction for upsert tool: %w", err)
 	}
-	// Ensure rollback on error
+	// Ensure rollback on error, commit on success
 	defer func() {
-		if p := recover(); p != nil {
-			_ = tx.Rollback()
-			panic(p) // Re-panic after rollback
-		} else if err != nil {
-			_ = tx.Rollback() // Rollback on normal error
-		} else {
-			err = tx.Commit() // Commit on success
-			if err != nil {
-				err = fmt.Errorf("failed to commit transaction for upsert tool: %w", err)
+		if err != nil {
+			// Rollback only if the error is not nil on exit
+			if rbErr := tx.Rollback(); rbErr != nil {
+				// Log or wrap the rollback error if necessary, but prioritize the original error
+				fmt.Printf("ERROR: Failed to rollback transaction after error: %v (original error: %v)\n", rbErr, err)
 			}
+			return // Keep the original error
+		}
+		// Commit if err is nil
+		err = tx.Commit()
+		if err != nil {
+			err = fmt.Errorf("failed to commit transaction for upsert tool: %w", err)
 		}
 	}()
 
-	// 1. Check if the tool already exists
-	var exists int
-	checkQuery := `SELECT COUNT(*) FROM tools WHERE external_id = ? AND source_server_id = ?`
-	err = tx.Get(&exists, checkQuery, tool.ExternalID, tool.SourceServerID)
-	if err != nil && err != sql.ErrNoRows { // Allow ErrNoRows here, though count should return 0
-		err = fmt.Errorf("failed to check tool existence: %w", err)
-		return // Defer will rollback
-	}
-
-	added = (exists == 0)
-
-	// 2. Perform the UPSERT
+	// Perform the UPSERT
 	upsertQuery := `
 		INSERT INTO tools (external_id, source_server_id, name, description, created_at, updated_at)
-        VALUES (:external_id, :source_server_id, :name, :description, :created_at, :updated_at)
-        ON CONFLICT(external_id, source_server_id) DO UPDATE SET
-            name = excluded.name,
-            description = excluded.description,
-            updated_at = excluded.updated_at
+		VALUES (:external_id, :source_server_id, :name, :description, :created_at, :updated_at)
+		ON CONFLICT(external_id, source_server_id) DO UPDATE SET
+			name = excluded.name,
+			description = excluded.description,
+			updated_at = excluded.updated_at
 	`
 	now := time.Now().UTC()
 	tool.CreatedAt = now // Set creation time (only used on INSERT)
@@ -245,8 +234,8 @@ func (db *DB) UpsertTool(tool models.Tool) (added bool, err error) {
 		return // Defer will rollback
 	}
 
-	// Defer will commit if err is nil
-	return added, err
+	// Defer handles commit/rollback
+	return err
 }
 
 // ListTools retrieves all tools from the database.

internal/backend/database/database_test.go

@@ -288,14 +288,10 @@ func TestUpsertTool(t *testing.T) {
 	}
 
 	// 1. Insert
-	// Capture both return values
-	added, err := db.UpsertTool(tool1)
+	err = db.UpsertTool(tool1)
 	if err != nil {
 		t.Fatalf("UpsertTool (insert) failed: %v", err)
 	}
-	if !added {
-		t.Error("UpsertTool (insert) should have returned added=true")
-	}
 
 	// Verify insert by listing
 	tools, err := db.ListToolsByServerID(serverID)
@@ -334,14 +330,10 @@ func TestUpsertTool(t *testing.T) {
 		Description:    "Does awesome things better", // Changed description
 	}
 
-	// Capture both return values
-	added, err = db.UpsertTool(updatedTool1)
+	err = db.UpsertTool(updatedTool1)
 	if err != nil {
 		t.Fatalf("UpsertTool (update) failed: %v", err)
 	}
-	if added {
-		t.Error("UpsertTool (update) should have returned added=false")
-	}
 
 	// Verify update
 	tools, err = db.ListToolsByServerID(serverID)
@@ -380,17 +372,15 @@ func TestUpsertTool_UniqueConstraint(t *testing.T) {
 	}
 
 	// Add tool for server 1
-	// Capture both return values (using _)
-	_, err := db.UpsertTool(tool)
+	err := db.UpsertTool(tool)
 	if err != nil {
 		t.Fatalf("UpsertTool for server 1 failed: %v", err)
 	}
 
 	// Add tool with same ExternalID but for server 2 (should succeed)
 	tool.SourceServerID = serverID2
 	tool.Description = "From Server 2"
-	// Capture both return values (using _)
-	_, err = db.UpsertTool(tool)
+	err = db.UpsertTool(tool)
 	if err != nil {
 		t.Fatalf("UpsertTool for server 2 with same external ID failed: %v", err)
 	}
@@ -423,15 +413,13 @@ func TestListTools(t *testing.T) {
 	}
 
 	for _, tool := range toolsS1 {
-		// Capture both return values (using _)
-		_, err := db.UpsertTool(tool)
+		err := db.UpsertTool(tool)
 		if err != nil {
 			t.Fatalf("UpsertTool failed: %v", err)
 		}
 	}
 	for _, tool := range toolsS2 {
-		// Capture both return values (using _)
-		_, err := db.UpsertTool(tool)
+		err := db.UpsertTool(tool)
 		if err != nil {
 			t.Fatalf("UpsertTool failed: %v", err)
 		}
@@ -483,14 +471,13 @@ func TestRemoveToolsByServerID(t *testing.T) {
 	servID2, _ := db.AddServer("S2", "http://s2.keep")
 
 	// Capture both return values (using _)
-	_, err := db.UpsertTool(models.Tool{ExternalID: "t1a", SourceServerID: servID1, Name: "Tool A1"})
+	err := db.UpsertTool(models.Tool{ExternalID: "t1a", SourceServerID: servID1, Name: "Tool A1"})
 	if err != nil {
 		t.Fatalf("UpsertTool failed: %v", err)
 	}
 
 	// Add a tool for server 2 as well
-	// Capture both return values (using _)
-	_, err = db.UpsertTool(models.Tool{ExternalID: "t2b", SourceServerID: servID2, Name: "Tool B2"})
+	err = db.UpsertTool(models.Tool{ExternalID: "t2b", SourceServerID: servID2, Name: "Tool B2"})
 	if err != nil {
 		t.Fatalf("UpsertTool failed: %v", err)
 	}
@@ -524,17 +511,17 @@ func TestForeignKey_CascadeDelete(t *testing.T) {
 	servID2, _ := db.AddServer("ServerToKeep", "http://s2.keep.me")
 
 	// Capture both return values (using _)
-	_, err := db.UpsertTool(models.Tool{ExternalID: "t1a", SourceServerID: servID1, Name: "Tool A1"})
+	err := db.UpsertTool(models.Tool{ExternalID: "t1a", SourceServerID: servID1, Name: "Tool A1"})
 	if err != nil {
 		t.Fatalf("UpsertTool failed: %v", err)
 	}
 	// Capture both return values (using _)
-	_, err = db.UpsertTool(models.Tool{ExternalID: "t1b", SourceServerID: servID1, Name: "Tool B1"})
+	err = db.UpsertTool(models.Tool{ExternalID: "t1b", SourceServerID: servID1, Name: "Tool B1"})
 	if err != nil {
 		t.Fatalf("UpsertTool failed: %v", err)
 	}
 	// Capture both return values (using _)
-	_, err = db.UpsertTool(models.Tool{ExternalID: "t2a", SourceServerID: servID2, Name: "Tool A2"})
+	err = db.UpsertTool(models.Tool{ExternalID: "t2a", SourceServerID: servID2, Name: "Tool A2"})
 	if err != nil {
 		t.Fatalf("UpsertTool failed: %v", err)
 	}

internal/backend/models/models.go

@@ -33,16 +33,13 @@ type MCPServer struct {
 
 // Tool represents tool metadata fetched from an MCP server and stored in the database.
 type Tool struct {
-	ID             int64  `db:"id" json:"id"`
-	ExternalID     string `db:"external_id" json:"external_id"`           // Tool ID from the source MCP server
-	SourceServerID int64  `db:"source_server_id" json:"source_server_id"` // Foreign key to mcp_servers
-	Name           string `db:"name" json:"name"`
-	Description    string `db:"description" json:"description,omitempty"`
-	// InputSchema    json.RawMessage `db:"input_schema" json:"input_schema,omitempty"` // Field removed from DB schema for now
-	// OutputSchema   json.RawMessage `db:"output_schema" json:"output_schema,omitempty"` // Field removed from DB schema for now
-	// Capabilities   json.RawMessage `db:"capabilities" json:"capabilities,omitempty"` // Field removed from DB schema for now
-	UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
-	CreatedAt time.Time `db:"created_at" json:"created_at"`
+	ID             int64     `db:"id" json:"id"`
+	ExternalID     string    `db:"external_id" json:"external_id"`           // Tool ID from the source MCP server
+	SourceServerID int64     `db:"source_server_id" json:"source_server_id"` // Foreign key to mcp_servers
+	Name           string    `db:"name" json:"name"`
+	Description    string    `db:"description" json:"description,omitempty"`
+	UpdatedAt      time.Time `db:"updated_at" json:"updated_at"`
+	CreatedAt      time.Time `db:"created_at" json:"created_at"`
 }
 
 // FetchedTool represents the structure of tool data as fetched from an external MCP server's API.
@@ -51,5 +48,4 @@ type FetchedTool struct {
 	ID          string `json:"id"` // External ID
 	Name        string `json:"name"`
 	Description string `json:"description"`
-	// Add other fields if the MCP protocol includes them (e.g., InputSchema, OutputSchema)
 }