Merge branch 'main' into refactor/cloud-integration-handlers

* fix: adding migration for fixing wrong cloud integration unique index

* refactor: removing std errors pkg

* refactor: normalising account_id if empty

* feat: adding integration test

ee/licensing/httplicensing/provider.go

@@ -198,7 +198,10 @@ func (provider *provider) Checkout(ctx context.Context, organizationID valuer.UU
 
 	response, err := provider.zeus.GetCheckoutURL(ctx, activeLicense.Key, body)
 	if err != nil {
-		return nil, errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to generate checkout session")
+		if errors.Ast(err, errors.TypeAlreadyExists) {
+			return nil, errors.WithAdditionalf(err, "checkout has already been completed for this account. Please click 'Refresh Status' to sync your subscription")
+		}
+		return nil, err
 	}
 
 	return &licensetypes.GettableSubscription{RedirectURL: gjson.GetBytes(response, "url").String()}, nil
@@ -217,7 +220,7 @@ func (provider *provider) Portal(ctx context.Context, organizationID valuer.UUID
 
 	response, err := provider.zeus.GetPortalURL(ctx, activeLicense.Key, body)
 	if err != nil {
-		return nil, errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to generate portal session")
+		return nil, err
 	}
 
 	return &licensetypes.GettableSubscription{RedirectURL: gjson.GetBytes(response, "url").String()}, nil

pkg/query-service/app/cloudintegrations/accountsRepo.go

@@ -177,7 +177,7 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 	onConflictClause := ""
 	if len(onConflictSetStmts) > 0 {
 		onConflictClause = fmt.Sprintf(
-			"conflict(id, provider, org_id) do update SET\n%s",
+			"conflict(id) do update SET\n%s",
 			strings.Join(onConflictSetStmts, ",\n"),
 		)
 	}
@@ -202,6 +202,8 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 		Exec(ctx)
 
 	if dbErr != nil {
+		// for now returning internal error even if there is a conflict,
+		// will be handled better in the future iteration
 		return nil, model.InternalError(fmt.Errorf(
 			"could not upsert cloud account record: %w", dbErr,
 		))

pkg/signoz/provider.go

@@ -175,6 +175,7 @@ func NewSQLMigrationProviderFactories(
 		sqlmigration.NewMigrateRulesV4ToV5Factory(sqlstore, telemetryStore),
 		sqlmigration.NewAddStatusUserFactory(sqlstore, sqlschema),
 		sqlmigration.NewDeprecateUserInviteFactory(sqlstore, sqlschema),
+		sqlmigration.NewUpdateCloudIntegrationUniqueIndexFactory(sqlstore, sqlschema),
 	)
 }
 

pkg/sqlmigration/069_update_cloud_integration_index.go

@@ -0,0 +1,255 @@
+package sqlmigration
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"time"
+
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/SigNoz/signoz/pkg/factory"
+	"github.com/SigNoz/signoz/pkg/sqlschema"
+	"github.com/SigNoz/signoz/pkg/sqlstore"
+	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/migrate"
+)
+
+type updateCloudIntegrationUniqueIndex struct {
+	sqlstore  sqlstore.SQLStore
+	sqlschema sqlschema.SQLSchema
+}
+
+func NewUpdateCloudIntegrationUniqueIndexFactory(sqlstore sqlstore.SQLStore, sqlschema sqlschema.SQLSchema) factory.ProviderFactory[SQLMigration, Config] {
+	return factory.NewProviderFactory(
+		factory.MustNewName("update_cloud_integration_index"),
+		func(ctx context.Context, ps factory.ProviderSettings, c Config) (SQLMigration, error) {
+			return &updateCloudIntegrationUniqueIndex{
+				sqlstore:  sqlstore,
+				sqlschema: sqlschema,
+			}, nil
+		},
+	)
+}
+
+func (migration *updateCloudIntegrationUniqueIndex) Register(migrations *migrate.Migrations) error {
+	if err := migrations.Register(migration.Up, migration.Down); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+type cloudIntegrationRow struct {
+	bun.BaseModel `bun:"table:cloud_integration"`
+
+	ID        string    `bun:"id"`
+	AccountID string    `bun:"account_id"`
+	Provider  string    `bun:"provider"`
+	OrgID     string    `bun:"org_id"`
+	Config    string    `bun:"config"`
+	UpdatedAt time.Time `bun:"updated_at"`
+}
+
+type cloudIntegrationAccountConfig struct {
+	Regions []string `json:"regions"`
+}
+
+// duplicateGroup holds the keeper (first element) and losers (rest) for a duplicate (account_id, provider, org_id) group.
+type duplicateGroup struct {
+	keeper *cloudIntegrationRow
+	losers []*cloudIntegrationRow
+}
+
+func (migration *updateCloudIntegrationUniqueIndex) Up(ctx context.Context, db *bun.DB) error {
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = tx.Rollback()
+	}()
+
+	sqls := [][]byte{}
+
+	// Step 1: Drop the wrong index on (id, provider, org_id)
+	dropSqls := migration.sqlschema.Operator().DropIndex(
+		(&sqlschema.UniqueIndex{
+			TableName:   "cloud_integration",
+			ColumnNames: []sqlschema.ColumnName{"id", "provider", "org_id"},
+		}).Named("unique_cloud_integration"),
+	)
+	sqls = append(sqls, dropSqls...)
+
+	// Step 2: Normalize empty-string account_id to NULL
+	// Older table structure could store "" instead of NULL for unconnected accounts.
+	// Empty strings would violate the partial unique index since '' = '' (unlike NULL != NULL).
+	_, err = tx.NewUpdate().
+		TableExpr("cloud_integration").
+		Set("account_id = NULL").
+		Where("account_id = ''").
+		Exec(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Step 3: Fetch all active rows with non-null account_id, ordered for grouping
+	var activeRows []*cloudIntegrationRow
+	err = tx.NewSelect().
+		Model(&activeRows).
+		Where("removed_at IS NULL").
+		Where("account_id IS NOT NULL").
+		OrderExpr("account_id, provider, org_id, updated_at DESC").
+		Scan(ctx)
+	if err != nil && !errors.Is(err, sql.ErrNoRows) {
+		return err
+	}
+
+	// Group by (account_id, provider, org_id)
+	groups := groupCloudIntegrationRows(activeRows)
+
+	now := time.Now()
+	var loserIDs []string
+
+	for _, group := range groups {
+		if len(group.losers) == 0 {
+			continue
+		}
+
+		// Step 4: Merge config from losers into keeper
+		if err = mergeCloudIntegrationConfigs(ctx, tx, group); err != nil {
+			return err
+		}
+
+		// Step 5: Reassign non-conflicting cloud_integration_service rows to keeper
+		for _, loser := range group.losers {
+			_, err = tx.NewUpdate().
+				TableExpr("cloud_integration_service").
+				Set("cloud_integration_id = ?", group.keeper.ID).
+				Where("cloud_integration_id = ?", loser.ID).
+				Where("type NOT IN (?)",
+					tx.NewSelect().
+						TableExpr("cloud_integration_service").
+						Column("type").
+						Where("cloud_integration_id = ?", group.keeper.ID),
+				).
+				Exec(ctx)
+			if err != nil {
+				return err
+			}
+
+			loserIDs = append(loserIDs, loser.ID)
+		}
+	}
+
+	// Step 6: Soft-delete all loser rows
+	if len(loserIDs) > 0 {
+		_, err = tx.NewUpdate().
+			TableExpr("cloud_integration").
+			Set("removed_at = ?", now).
+			Set("updated_at = ?", now).
+			Where("id IN (?)", bun.In(loserIDs)).
+			Exec(ctx)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Step 7: Create the correct partial unique index on (account_id, provider, org_id) WHERE removed_at IS NULL
+	createSqls := migration.sqlschema.Operator().CreateIndex(
+		&sqlschema.PartialUniqueIndex{
+			TableName:   "cloud_integration",
+			ColumnNames: []sqlschema.ColumnName{"account_id", "provider", "org_id"},
+			Where:       "removed_at IS NULL",
+		},
+	)
+	sqls = append(sqls, createSqls...)
+
+	for _, sql := range sqls {
+		if _, err = tx.ExecContext(ctx, string(sql)); err != nil {
+			return err
+		}
+	}
+
+	return tx.Commit()
+}
+
+func (migration *updateCloudIntegrationUniqueIndex) Down(ctx context.Context, db *bun.DB) error {
+	return nil
+}
+
+// groupCloudIntegrationRows groups rows by (account_id, provider, org_id).
+// Rows must be pre-sorted by account_id, provider, org_id, updated_at DESC
+// so the first row in each group is the keeper (most recently updated).
+func groupCloudIntegrationRows(rows []*cloudIntegrationRow) []duplicateGroup {
+	if len(rows) == 0 {
+		return nil
+	}
+
+	var groups []duplicateGroup
+	var current duplicateGroup
+	current.keeper = rows[0]
+
+	for i := 1; i < len(rows); i++ {
+		row := rows[i]
+		if row.AccountID == current.keeper.AccountID &&
+			row.Provider == current.keeper.Provider &&
+			row.OrgID == current.keeper.OrgID {
+			current.losers = append(current.losers, row)
+		} else {
+			groups = append(groups, current)
+			current = duplicateGroup{keeper: row}
+		}
+	}
+	groups = append(groups, current)
+
+	return groups
+}
+
+// mergeCloudIntegrationConfigs unions the EnabledRegions from all rows in the group into the keeper's config and updates
+func mergeCloudIntegrationConfigs(ctx context.Context, tx bun.Tx, group duplicateGroup) error {
+	regionSet := make(map[string]struct{})
+
+	// Parse keeper's config
+	parseRegions(group.keeper.Config, regionSet)
+
+	// Parse each loser's config
+	for _, loser := range group.losers {
+		parseRegions(loser.Config, regionSet)
+	}
+
+	// Build merged config
+	mergedRegions := make([]string, 0, len(regionSet))
+	for region := range regionSet {
+		mergedRegions = append(mergedRegions, region)
+	}
+
+	merged := cloudIntegrationAccountConfig{Regions: mergedRegions}
+	mergedJSON, err := json.Marshal(merged)
+	if err != nil {
+		return err
+	}
+
+	// Update keeper's config
+	_, err = tx.NewUpdate().
+		TableExpr("cloud_integration").
+		Set("config = ?", string(mergedJSON)).
+		Where("id = ?", group.keeper.ID).
+		Exec(ctx)
+	return err
+}
+
+// parseRegions unmarshals a config JSON string and adds its regions to the set.
+func parseRegions(configJSON string, regionSet map[string]struct{}) {
+	if configJSON == "" {
+		return
+	}
+
+	var config cloudIntegrationAccountConfig
+	if err := json.Unmarshal([]byte(configJSON), &config); err != nil {
+		return
+	}
+
+	for _, region := range config.Regions {
+		regionSet[region] = struct{}{}
+	}
+}

tests/integration/fixtures/cloudintegrations.py

@@ -1,7 +1,7 @@
 """Fixtures for cloud integration tests."""
 
 from http import HTTPStatus
-from typing import Callable, Optional
+from typing import Callable
 
 import pytest
 import requests
@@ -18,14 +18,12 @@ def create_cloud_integration_account(
     request: pytest.FixtureRequest,
     signoz: types.SigNoz,
 ) -> Callable[[str, str], dict]:
-    created_account_id: Optional[str] = None
-    cloud_provider_used: Optional[str] = None
+    created_accounts: list[tuple[str, str]] = []
 
     def _create(
         admin_token: str,
         cloud_provider: str = "aws",
     ) -> dict:
-        nonlocal created_account_id, cloud_provider_used
         endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/accounts/generate-connection-url"
 
         request_payload = {
@@ -52,35 +50,31 @@ def create_cloud_integration_account(
         ), f"Failed to create test account: {response.status_code}"
 
         data = response.json().get("data", response.json())
-        created_account_id = data.get("account_id")
-        cloud_provider_used = cloud_provider
+        created_accounts.append((data.get("account_id"), cloud_provider))
 
         return data
 
-    def _disconnect(admin_token: str, cloud_provider: str) -> requests.Response:
-        assert created_account_id
-        disconnect_endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/accounts/{created_account_id}/disconnect"
-        return requests.post(
-            signoz.self.host_configs["8080"].get(disconnect_endpoint),
-            headers={"Authorization": f"Bearer {admin_token}"},
-            timeout=10,
-        )
-
     # Yield factory to the test
     yield _create
 
-    # Post-test cleanup: generate admin token and disconnect the created account
-    if created_account_id and cloud_provider_used:
+    # Post-test cleanup: disconnect all created accounts
+    if created_accounts:
         get_token = request.getfixturevalue("get_token")
         try:
             admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
-            r = _disconnect(admin_token, cloud_provider_used)
-            if r.status_code != HTTPStatus.OK:
-                logger.info(
-                    "Disconnect cleanup returned %s for account %s",
-                    r.status_code,
-                    created_account_id,
+            for account_id, cloud_provider in created_accounts:
+                disconnect_endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/accounts/{account_id}/disconnect"
+                r = requests.post(
+                    signoz.self.host_configs["8080"].get(disconnect_endpoint),
+                    headers={"Authorization": f"Bearer {admin_token}"},
+                    timeout=10,
                 )
-            logger.info("Cleaned up test account: %s", created_account_id)
+                if r.status_code != HTTPStatus.OK:
+                    logger.info(
+                        "Disconnect cleanup returned %s for account %s",
+                        r.status_code,
+                        account_id,
+                    )
+                logger.info("Cleaned up test account: %s", account_id)
         except Exception as exc:  # pylint: disable=broad-except
             logger.info("Post-test disconnect cleanup failed: %s", exc)

tests/integration/fixtures/cloudintegrationsutils.py

@@ -1,7 +1,5 @@
 """Fixtures for cloud integration tests."""
 
-from http import HTTPStatus
-
 import requests
 
 from fixtures import types
@@ -16,7 +14,7 @@ def simulate_agent_checkin(
     cloud_provider: str,
     account_id: str,
     cloud_account_id: str,
-) -> dict:
+) -> requests.Response:
     endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/agent-check-in"
 
     checkin_payload = {
@@ -32,16 +30,11 @@ def simulate_agent_checkin(
         timeout=10,
     )
 
-    if response.status_code != HTTPStatus.OK:
+    if not response.ok:
         logger.error(
             "Agent check-in failed: %s, response: %s",
             response.status_code,
             response.text,
         )
 
-    assert (
-        response.status_code == HTTPStatus.OK
-    ), f"Agent check-in failed: {response.status_code}"
-
-    response_data = response.json()
-    return response_data.get("data", response_data)
+    return response

tests/integration/src/cloudintegrations/02_generate_connection_url.py

@@ -1,3 +1,4 @@
+import uuid
 from http import HTTPStatus
 from typing import Callable
 
@@ -5,6 +6,8 @@ import requests
 
 from fixtures import types
 from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD
+from fixtures.cloudintegrations import create_cloud_integration_account
+from fixtures.cloudintegrationsutils import simulate_agent_checkin
 from fixtures.logger import setup_logger
 
 logger = setup_logger(__name__)
@@ -142,3 +145,42 @@ def test_generate_connection_url_unsupported_provider(
     assert (
         "unsupported cloud provider" in response_data["error"].lower()
     ), "Error message should indicate unsupported provider"
+
+
+def test_duplicate_cloud_account_checkins(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+) -> None:
+    """Test that two accounts cannot check in with the same cloud_account_id."""
+
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    cloud_provider = "aws"
+    same_cloud_account_id = str(uuid.uuid4())
+
+    # Create two separate cloud integration accounts via generate-connection-url
+    account1 = create_cloud_integration_account(admin_token, cloud_provider)
+    account1_id = account1["account_id"]
+
+    account2 = create_cloud_integration_account(admin_token, cloud_provider)
+    account2_id = account2["account_id"]
+
+    assert account1_id != account2_id, "Two accounts should have different internal IDs"
+
+#     First check-in succeeds: account1 claims cloud_account_id
+    response = simulate_agent_checkin(
+        signoz, admin_token, cloud_provider, account1_id, same_cloud_account_id
+    )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for first check-in, got {response.status_code}: {response.text}"
+#
+    # Second check-in should fail: account2 tries to use the same cloud_account_id
+    response = simulate_agent_checkin(
+        signoz, admin_token, cloud_provider, account2_id, same_cloud_account_id
+    )
+
+    assert (
+        response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
+    ), f"Expected 500 for duplicate cloud_account_id, got {response.status_code}: {response.text}"

tests/integration/src/cloudintegrations/03_accounts.py

@@ -57,9 +57,12 @@ def test_list_connected_accounts_with_account(
 
     # Simulate agent check-in to mark as connected
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # List accounts
     endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/accounts"
@@ -161,9 +164,12 @@ def test_update_account_config(
 
     # Simulate agent check-in to mark as connected
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Update account configuration
     endpoint = (
@@ -226,9 +232,12 @@ def test_disconnect_account(
 
     # Simulate agent check-in to mark as connected
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Disconnect the account
     endpoint = (

tests/integration/src/cloudintegrations/04_services.py

@@ -61,9 +61,12 @@ def test_list_services_with_account(
     account_id = account_data["account_id"]
 
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # List services for the account
     endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/services?cloud_account_id={cloud_account_id}"
@@ -152,9 +155,12 @@ def test_get_service_details_with_account(
     account_id = account_data["account_id"]
 
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Get list of services first
     list_endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/services"
@@ -253,9 +259,12 @@ def test_update_service_config(
     account_id = account_data["account_id"]
 
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Get list of services to pick a valid service ID
     list_endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/services"
@@ -365,9 +374,12 @@ def test_update_service_config_invalid_service(
     account_id = account_data["account_id"]
 
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Try to update config for invalid service
     fake_service_id = "non-existent-service"
@@ -409,9 +421,12 @@ def test_update_service_config_disable_service(
     account_id = account_data["account_id"]
 
     cloud_account_id = str(uuid.uuid4())
-    simulate_agent_checkin(
+    response = simulate_agent_checkin(
         signoz, admin_token, cloud_provider, account_id, cloud_account_id
     )
+    assert (
+        response.status_code == HTTPStatus.OK
+    ), f"Expected 200 for agent check-in, got {response.status_code}: {response.text}"
 
     # Get a valid service
     list_endpoint = f"/api/v1/cloud-integrations/{cloud_provider}/services"