From 5d1bfdc5e543c96ab968a0b71035085c773469ba Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 2 Jan 2025 20:18:48 +0530
Subject: [PATCH 01/38] feat: init app/cloud_integrations

---
 pkg/query-service/app/cloud_integrations/Readme.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 pkg/query-service/app/cloud_integrations/Readme.md

diff --git a/pkg/query-service/app/cloud_integrations/Readme.md b/pkg/query-service/app/cloud_integrations/Readme.md
new file mode 100644
index 00000000000..9ac1d5fe4a3
--- /dev/null
+++ b/pkg/query-service/app/cloud_integrations/Readme.md
@@ -0,0 +1,5 @@
+# SigNoz Cloud Integrations
+
+Cloud integrations are unlike the rest of SigNoz integrations.
+They have a different UX and so require a different API.
+They will also be limited in number and are not expected to have community contributed implementations

From c20663827caa653b7cb223bf217fa2576d75c0cd Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 3 Jan 2025 17:44:31 +0530
Subject: [PATCH 02/38] feat: get API test started for cloudintegrations
 account lifecycle

---
 .../signoz_cloud_integrations_test.go         | 141 ++++++++++++++++++
 1 file changed, 141 insertions(+)
 create mode 100644 pkg/query-service/tests/integration/signoz_cloud_integrations_test.go

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
new file mode 100644
index 00000000000..957f306a4cf
--- /dev/null
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -0,0 +1,141 @@
+package tests
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/jmoiron/sqlx"
+	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
+	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/pkg/query-service/app"
+	cloudintegrations "go.signoz.io/signoz/pkg/query-service/app/cloud_integrations"
+	"go.signoz.io/signoz/pkg/query-service/auth"
+	"go.signoz.io/signoz/pkg/query-service/dao"
+	"go.signoz.io/signoz/pkg/query-service/featureManager"
+	"go.signoz.io/signoz/pkg/query-service/model"
+	"go.signoz.io/signoz/pkg/query-service/utils"
+)
+
+func TestCloudIntegrationLifecycle(t *testing.T) {
+	// Test for the happy path of connecting and managing cloud integration accounts
+
+	require := require.New(t)
+	testbed := NewCloudIntegrationsTestBed(t, nil)
+
+	accountsListResp := testbed.GetAccountsListFromQS("aws")
+	require.Equal(len(accountsListResp.Accounts), 0,
+		"No accounts should be connected at the beginning",
+	)
+
+	// Should be able to generate connection url - initializing an account
+
+	// Should be able to poll for account connection status
+
+	// The unconnected account should not show up in accounts list yet
+
+	// An agent should be able to check in to the new account
+	// Should have the settings that were specified while generating connection url
+
+	// Polling for connection status should return latest status now.
+
+	// The account should now show up in list of connected accounts.
+
+	// Should be able to update account settings
+
+	// The agent should now receive latest settings.
+
+	// Should be able to disconnect account.
+
+	// The agent should receive the disconnected status post disconnection
+
+	require.Equal(1, 2)
+}
+
+type CloudIntegrationsTestBed struct {
+	t              *testing.T
+	testUser       *model.User
+	qsHttpHandler  http.Handler
+	mockClickhouse mockhouse.ClickConnMockCommon
+}
+
+// testDB can be injected for sharing a DB across multiple integration testbeds.
+func NewCloudIntegrationsTestBed(t *testing.T, testDB *sqlx.DB) *CloudIntegrationsTestBed {
+	if testDB == nil {
+		testDB = utils.NewQueryServiceDBForTests(t)
+	}
+
+	controller, err := cloudintegrations.NewController(testDB)
+	if err != nil {
+		t.Fatalf("could not create cloud integrations controller: %v", err)
+	}
+
+	fm := featureManager.StartManager()
+	reader, mockClickhouse := NewMockClickhouseReader(t, testDB, fm)
+	mockClickhouse.MatchExpectationsInOrder(false)
+
+	apiHandler, err := app.NewAPIHandler(app.APIHandlerOpts{
+		Reader:                      reader,
+		AppDao:                      dao.DB(),
+		CloudIntegrationsController: controller,
+		FeatureFlags:                fm,
+	})
+	if err != nil {
+		t.Fatalf("could not create a new ApiHandler: %v", err)
+	}
+
+	router := app.NewRouter()
+	am := app.NewAuthMiddleware(auth.GetUserFromRequest)
+	apiHandler.RegisterRoutes(router, am)
+	apiHandler.RegisterIntegrationRoutes(router, am)
+
+	user, apiErr := createTestUser()
+	if apiErr != nil {
+		t.Fatalf("could not create a test user: %v", apiErr)
+	}
+
+	return &CloudIntegrationsTestBed{
+		t:              t,
+		testUser:       user,
+		qsHttpHandler:  router,
+		mockClickhouse: mockClickhouse,
+	}
+}
+
+func (tb *CloudIntegrationsTestBed) GetAccountsListFromQS(
+	cloudProvider string,
+) *cloudintegrations.AccountsListResponse {
+	result := tb.RequestQS(fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts", cloudProvider), nil)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var accountsListResp cloudintegrations.AccountsListResponse
+	err = json.Unmarshal(dataJson, &accountsListResp)
+	if err != nil {
+		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into AccountsListResponse")
+	}
+
+	return &accountsListResp
+}
+
+func (tb *CloudIntegrationsTestBed) RequestQS(
+	path string,
+	postData interface{},
+) *app.ApiResponse {
+	req, err := AuthenticatedRequestForTest(
+		tb.testUser, path, postData,
+	)
+	if err != nil {
+		tb.t.Fatalf("couldn't create authenticated test request: %v", err)
+	}
+
+	result, err := HandleTestRequest(tb.qsHttpHandler, req, 200)
+	if err != nil {
+		tb.t.Fatalf("test request failed: %v", err)
+	}
+	return result
+}

From 9973cadc8eb9714ba15f72e186b09808cc0219d6 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 3 Jan 2025 17:45:03 +0530
Subject: [PATCH 03/38] feat: cloudintegrations: get controller started

---
 .../app/cloud_integrations/controller.go      | 19 +++++++++++++++++++
 .../app/cloud_integrations/manager.go         |  1 +
 .../app/cloud_integrations/model.go           |  5 +++++
 3 files changed, 25 insertions(+)
 create mode 100644 pkg/query-service/app/cloud_integrations/controller.go
 create mode 100644 pkg/query-service/app/cloud_integrations/manager.go
 create mode 100644 pkg/query-service/app/cloud_integrations/model.go

diff --git a/pkg/query-service/app/cloud_integrations/controller.go b/pkg/query-service/app/cloud_integrations/controller.go
new file mode 100644
index 00000000000..8be06bafaee
--- /dev/null
+++ b/pkg/query-service/app/cloud_integrations/controller.go
@@ -0,0 +1,19 @@
+package cloudintegrations
+
+import "github.com/jmoiron/sqlx"
+
+type Controller struct {
+	db *sqlx.DB
+}
+
+func NewController(db *sqlx.DB) (
+	*Controller, error,
+) {
+	return &Controller{
+		db: db,
+	}, nil
+}
+
+type AccountsListResponse struct {
+	Accounts []Account `json:"accounts"`
+}
diff --git a/pkg/query-service/app/cloud_integrations/manager.go b/pkg/query-service/app/cloud_integrations/manager.go
new file mode 100644
index 00000000000..b1b0ea798cb
--- /dev/null
+++ b/pkg/query-service/app/cloud_integrations/manager.go
@@ -0,0 +1 @@
+package cloudintegrations
diff --git a/pkg/query-service/app/cloud_integrations/model.go b/pkg/query-service/app/cloud_integrations/model.go
new file mode 100644
index 00000000000..f600fec9624
--- /dev/null
+++ b/pkg/query-service/app/cloud_integrations/model.go
@@ -0,0 +1,5 @@
+package cloudintegrations
+
+// Represents a cloud provider account for cloud integrations
+type Account struct {
+}

From 45302ce34dfaf3137b70fb7462bb82f40231abc1 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 3 Jan 2025 17:58:32 +0530
Subject: [PATCH 04/38] feat: cloud integrations: add
 cloudintegrations.Controller to APIHandler and servers

---
 ee/query-service/app/api/api.go                          | 3 +++
 ee/query-service/app/server.go                           | 9 +++++++++
 .../{cloud_integrations => cloudintegrations}/Readme.md  | 0
 .../controller.go                                        | 0
 .../{cloud_integrations => cloudintegrations}/manager.go | 0
 .../{cloud_integrations => cloudintegrations}/model.go   | 0
 pkg/query-service/app/http_handler.go                    | 7 +++++++
 pkg/query-service/app/server.go                          | 7 +++++++
 .../tests/integration/signoz_cloud_integrations_test.go  | 2 +-
 9 files changed, 27 insertions(+), 1 deletion(-)
 rename pkg/query-service/app/{cloud_integrations => cloudintegrations}/Readme.md (100%)
 rename pkg/query-service/app/{cloud_integrations => cloudintegrations}/controller.go (100%)
 rename pkg/query-service/app/{cloud_integrations => cloudintegrations}/manager.go (100%)
 rename pkg/query-service/app/{cloud_integrations => cloudintegrations}/model.go (100%)

diff --git a/ee/query-service/app/api/api.go b/ee/query-service/app/api/api.go
index 181186d3230..25787667fac 100644
--- a/ee/query-service/app/api/api.go
+++ b/ee/query-service/app/api/api.go
@@ -12,6 +12,7 @@ import (
 	"go.signoz.io/signoz/ee/query-service/license"
 	"go.signoz.io/signoz/ee/query-service/usage"
 	baseapp "go.signoz.io/signoz/pkg/query-service/app"
+	"go.signoz.io/signoz/pkg/query-service/app/cloudintegrations"
 	"go.signoz.io/signoz/pkg/query-service/app/integrations"
 	"go.signoz.io/signoz/pkg/query-service/app/logparsingpipeline"
 	"go.signoz.io/signoz/pkg/query-service/cache"
@@ -34,6 +35,7 @@ type APIHandlerOptions struct {
 	FeatureFlags                  baseint.FeatureLookup
 	LicenseManager                *license.Manager
 	IntegrationsController        *integrations.Controller
+	CloudIntegrationsController   *cloudintegrations.Controller
 	LogsParsingPipelineController *logparsingpipeline.LogParsingPipelineController
 	Cache                         cache.Cache
 	Gateway                       *httputil.ReverseProxy
@@ -62,6 +64,7 @@ func NewAPIHandler(opts APIHandlerOptions) (*APIHandler, error) {
 		RuleManager:                   opts.RulesManager,
 		FeatureFlags:                  opts.FeatureFlags,
 		IntegrationsController:        opts.IntegrationsController,
+		CloudIntegrationsController:   opts.CloudIntegrationsController,
 		LogsParsingPipelineController: opts.LogsParsingPipelineController,
 		Cache:                         opts.Cache,
 		FluxInterval:                  opts.FluxInterval,
diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go
index 1b17ca43d01..51c320111ae 100644
--- a/ee/query-service/app/server.go
+++ b/ee/query-service/app/server.go
@@ -40,6 +40,7 @@ import (
 
 	"go.signoz.io/signoz/pkg/query-service/agentConf"
 	baseapp "go.signoz.io/signoz/pkg/query-service/app"
+	"go.signoz.io/signoz/pkg/query-service/app/cloudintegrations"
 	"go.signoz.io/signoz/pkg/query-service/app/dashboards"
 	baseexplorer "go.signoz.io/signoz/pkg/query-service/app/explorer"
 	"go.signoz.io/signoz/pkg/query-service/app/integrations"
@@ -221,6 +222,13 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		)
 	}
 
+	cloudIntegrationsController, err := cloudintegrations.NewController(localDB)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"couldn't create cloud provider integrations controller: %w", err,
+		)
+	}
+
 	// ingestion pipelines manager
 	logParsingPipelineController, err := logparsingpipeline.NewLogParsingPipelinesController(
 		localDB, "sqlite", integrationsController.GetPipelinesForInstalledIntegrations,
@@ -271,6 +279,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		FeatureFlags:                  lm,
 		LicenseManager:                lm,
 		IntegrationsController:        integrationsController,
+		CloudIntegrationsController:   cloudIntegrationsController,
 		LogsParsingPipelineController: logParsingPipelineController,
 		Cache:                         c,
 		FluxInterval:                  fluxInterval,
diff --git a/pkg/query-service/app/cloud_integrations/Readme.md b/pkg/query-service/app/cloudintegrations/Readme.md
similarity index 100%
rename from pkg/query-service/app/cloud_integrations/Readme.md
rename to pkg/query-service/app/cloudintegrations/Readme.md
diff --git a/pkg/query-service/app/cloud_integrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
similarity index 100%
rename from pkg/query-service/app/cloud_integrations/controller.go
rename to pkg/query-service/app/cloudintegrations/controller.go
diff --git a/pkg/query-service/app/cloud_integrations/manager.go b/pkg/query-service/app/cloudintegrations/manager.go
similarity index 100%
rename from pkg/query-service/app/cloud_integrations/manager.go
rename to pkg/query-service/app/cloudintegrations/manager.go
diff --git a/pkg/query-service/app/cloud_integrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
similarity index 100%
rename from pkg/query-service/app/cloud_integrations/model.go
rename to pkg/query-service/app/cloudintegrations/model.go
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index 030f715dbeb..e2ce9d83c4c 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -25,6 +25,7 @@ import (
 	"github.com/prometheus/prometheus/promql"
 
 	"go.signoz.io/signoz/pkg/query-service/agentConf"
+	"go.signoz.io/signoz/pkg/query-service/app/cloudintegrations"
 	"go.signoz.io/signoz/pkg/query-service/app/dashboards"
 	"go.signoz.io/signoz/pkg/query-service/app/explorer"
 	"go.signoz.io/signoz/pkg/query-service/app/inframetrics"
@@ -101,6 +102,8 @@ type APIHandler struct {
 
 	IntegrationsController *integrations.Controller
 
+	CloudIntegrationsController *cloudintegrations.Controller
+
 	LogsParsingPipelineController *logparsingpipeline.LogParsingPipelineController
 
 	// SetupCompleted indicates if SigNoz is ready for general use.
@@ -154,6 +157,9 @@ type APIHandlerOpts struct {
 	// Integrations
 	IntegrationsController *integrations.Controller
 
+	// Cloud Provider Integrations
+	CloudIntegrationsController *cloudintegrations.Controller
+
 	// Log parsing pipelines
 	LogsParsingPipelineController *logparsingpipeline.LogParsingPipelineController
 
@@ -225,6 +231,7 @@ func NewAPIHandler(opts APIHandlerOpts) (*APIHandler, error) {
 		ruleManager:                   opts.RuleManager,
 		featureFlags:                  opts.FeatureFlags,
 		IntegrationsController:        opts.IntegrationsController,
+		CloudIntegrationsController:   opts.CloudIntegrationsController,
 		LogsParsingPipelineController: opts.LogsParsingPipelineController,
 		querier:                       querier,
 		querierV2:                     querierv2,
diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go
index 20d183876cf..6ef58113b66 100644
--- a/pkg/query-service/app/server.go
+++ b/pkg/query-service/app/server.go
@@ -24,6 +24,7 @@ import (
 	"github.com/soheilhy/cmux"
 	"go.signoz.io/signoz/pkg/query-service/agentConf"
 	"go.signoz.io/signoz/pkg/query-service/app/clickhouseReader"
+	"go.signoz.io/signoz/pkg/query-service/app/cloudintegrations"
 	"go.signoz.io/signoz/pkg/query-service/app/dashboards"
 	"go.signoz.io/signoz/pkg/query-service/app/integrations"
 	"go.signoz.io/signoz/pkg/query-service/app/logparsingpipeline"
@@ -181,6 +182,11 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		return nil, fmt.Errorf("couldn't create integrations controller: %w", err)
 	}
 
+	cloudIntegrationsController, err := cloudintegrations.NewController(localDB)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't create cloud provider integrations controller: %w", err)
+	}
+
 	logParsingPipelineController, err := logparsingpipeline.NewLogParsingPipelinesController(
 		localDB, "sqlite", integrationsController.GetPipelinesForInstalledIntegrations,
 	)
@@ -200,6 +206,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		RuleManager:                   rm,
 		FeatureFlags:                  fm,
 		IntegrationsController:        integrationsController,
+		CloudIntegrationsController:   cloudIntegrationsController,
 		LogsParsingPipelineController: logParsingPipelineController,
 		Cache:                         c,
 		FluxInterval:                  fluxInterval,
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 957f306a4cf..7f012cf71ef 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -10,7 +10,7 @@ import (
 	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
 	"github.com/stretchr/testify/require"
 	"go.signoz.io/signoz/pkg/query-service/app"
-	cloudintegrations "go.signoz.io/signoz/pkg/query-service/app/cloud_integrations"
+	"go.signoz.io/signoz/pkg/query-service/app/cloudintegrations"
 	"go.signoz.io/signoz/pkg/query-service/auth"
 	"go.signoz.io/signoz/pkg/query-service/dao"
 	"go.signoz.io/signoz/pkg/query-service/featureManager"

From d36848e1db813d7232c68007b82859aae09d12b2 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 3 Jan 2025 18:21:18 +0530
Subject: [PATCH 05/38] feat: cloud integrations: get routes started

---
 ee/query-service/app/server.go                |  1 +
 .../app/cloudintegrations/controller.go       | 34 ++++++++++++++++++-
 .../app/cloudintegrations/manager.go          |  1 -
 pkg/query-service/app/http_handler.go         | 26 ++++++++++++++
 pkg/query-service/app/server.go               |  1 +
 .../signoz_cloud_integrations_test.go         |  2 +-
 6 files changed, 62 insertions(+), 3 deletions(-)
 delete mode 100644 pkg/query-service/app/cloudintegrations/manager.go

diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go
index 51c320111ae..a604860b701 100644
--- a/ee/query-service/app/server.go
+++ b/ee/query-service/app/server.go
@@ -379,6 +379,7 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web *web.Web) (*
 	apiHandler.RegisterRoutes(r, am)
 	apiHandler.RegisterLogsRoutes(r, am)
 	apiHandler.RegisterIntegrationRoutes(r, am)
+	apiHandler.RegisterCloudIntegrationsRoutes(r, am)
 	apiHandler.RegisterQueryRangeV3Routes(r, am)
 	apiHandler.RegisterInfraMetricsRoutes(r, am)
 	apiHandler.RegisterQueryRangeV4Routes(r, am)
diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 8be06bafaee..e85d25deae9 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -1,6 +1,24 @@
 package cloudintegrations
 
-import "github.com/jmoiron/sqlx"
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	"github.com/jmoiron/sqlx"
+	"go.signoz.io/signoz/pkg/query-service/model"
+)
+
+var SupportedCloudProviders = []string{
+	"aws",
+}
+
+func validateCloudProviderName(name string) *model.ApiError {
+	if !slices.Contains(SupportedCloudProviders, name) {
+		return model.BadRequest(fmt.Errorf("invalid cloud provider: %s", name))
+	}
+	return nil
+}
 
 type Controller struct {
 	db *sqlx.DB
@@ -17,3 +35,17 @@ func NewController(db *sqlx.DB) (
 type AccountsListResponse struct {
 	Accounts []Account `json:"accounts"`
 }
+
+func (c *Controller) ListConnectedAccounts(
+	ctx context.Context, cloudProvider string,
+) (
+	*AccountsListResponse, *model.ApiError,
+) {
+	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
+		return nil, apiErr
+	}
+
+	return &AccountsListResponse{
+		Accounts: []Account{},
+	}, nil
+}
diff --git a/pkg/query-service/app/cloudintegrations/manager.go b/pkg/query-service/app/cloudintegrations/manager.go
deleted file mode 100644
index b1b0ea798cb..00000000000
--- a/pkg/query-service/app/cloudintegrations/manager.go
+++ /dev/null
@@ -1 +0,0 @@
-package cloudintegrations
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index e2ce9d83c4c..e4e6e05518b 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3864,6 +3864,32 @@ func (aH *APIHandler) UninstallIntegration(
 	aH.Respond(w, map[string]interface{}{})
 }
 
+// cloud provider integrations
+func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *AuthMiddleware) {
+	subRouter := router.PathPrefix("/api/v1/cloud-integrations").Subrouter()
+
+	subRouter.HandleFunc(
+		"/{cloudProvider}/accounts", am.ViewAccess(aH.CloudIntegrationsListConnectedAccounts),
+	).Methods(http.MethodGet)
+
+}
+
+func (aH *APIHandler) CloudIntegrationsListConnectedAccounts(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+
+	resp, apiErr := aH.CloudIntegrationsController.ListConnectedAccounts(
+		r.Context(), cloudProvider,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, fmt.Sprintf("Failed to fetch connected %s accounts", cloudProvider))
+		return
+	}
+	aH.Respond(w, resp)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()
diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go
index 6ef58113b66..25ea1e8c938 100644
--- a/pkg/query-service/app/server.go
+++ b/pkg/query-service/app/server.go
@@ -317,6 +317,7 @@ func (s *Server) createPublicServer(api *APIHandler) (*http.Server, error) {
 	api.RegisterRoutes(r, am)
 	api.RegisterLogsRoutes(r, am)
 	api.RegisterIntegrationRoutes(r, am)
+	api.RegisterCloudIntegrationsRoutes(r, am)
 	api.RegisterQueryRangeV3Routes(r, am)
 	api.RegisterInfraMetricsRoutes(r, am)
 	api.RegisterWebSocketPaths(r, am)
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 7f012cf71ef..1e7e20f17a7 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -88,7 +88,7 @@ func NewCloudIntegrationsTestBed(t *testing.T, testDB *sqlx.DB) *CloudIntegratio
 	router := app.NewRouter()
 	am := app.NewAuthMiddleware(auth.GetUserFromRequest)
 	apiHandler.RegisterRoutes(router, am)
-	apiHandler.RegisterIntegrationRoutes(router, am)
+	apiHandler.RegisterCloudIntegrationsRoutes(router, am)
 
 	user, apiErr := createTestUser()
 	if apiErr != nil {

From f864c0f51cb089a5abf6e5199c0b6a955f350036 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 3 Jan 2025 20:04:09 +0530
Subject: [PATCH 06/38] feat: cloud integrations: get accounts table schema
 started

---
 .../app/cloudintegrations/repo.go             | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100644 pkg/query-service/app/cloudintegrations/repo.go

diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
new file mode 100644
index 00000000000..1c919331cf0
--- /dev/null
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -0,0 +1,38 @@
+package cloudintegrations
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jmoiron/sqlx"
+	"go.signoz.io/signoz/pkg/query-service/model"
+)
+
+type CloudProviderAccountsRepository interface {
+	listConnectedAccounts(context.Context) ([]Account, *model.ApiError)
+}
+
+func InitSqliteDBIfNeeded(db *sqlx.DB) error {
+	if db == nil {
+		return fmt.Errorf("db is required")
+	}
+
+	createTablesStatements := `
+		CREATE TABLE IF NOT EXISTS cloud_integrations_accounts(
+			id TEXT PRIMARY KEY NOT NULL,
+			config_json TEXT NOT NULL,
+			cloud_account_id TEXT,
+			last_agent_report_json TEXT,
+			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
+			removed_at TIMESTAMP,
+		)
+	`
+	_, err := db.Exec(createTablesStatements)
+	if err != nil {
+		return fmt.Errorf(
+			"could not ensure cloud provider integrations schema in sqlite DB: %w", err,
+		)
+	}
+
+	return nil
+}

From 3434819b2720b8ba1a29d270eb8ad7fa144b943f Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Sun, 5 Jan 2025 14:24:05 +0530
Subject: [PATCH 07/38] feat: cloud integrations: get
 cloudProviderAccountsSQLRepository started

---
 .../app/cloudintegrations/controller.go       | 16 +++++++++---
 .../app/cloudintegrations/repo.go             | 26 +++++++++++++++++--
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index e85d25deae9..2dd09e96262 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -21,14 +21,19 @@ func validateCloudProviderName(name string) *model.ApiError {
 }
 
 type Controller struct {
-	db *sqlx.DB
+	repo cloudProviderAccountsRepository
 }
 
 func NewController(db *sqlx.DB) (
 	*Controller, error,
 ) {
+	repo, err := newCloudProviderAccountsRepository(db)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't create cloud provider accounts repo", err)
+	}
+
 	return &Controller{
-		db: db,
+		repo: repo,
 	}, nil
 }
 
@@ -45,7 +50,12 @@ func (c *Controller) ListConnectedAccounts(
 		return nil, apiErr
 	}
 
+	accounts, apiErr := c.repo.listConnectedAccounts(ctx)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't list cloud accounts")
+	}
+
 	return &AccountsListResponse{
-		Accounts: []Account{},
+		Accounts: accounts,
 	}, nil
 }
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 1c919331cf0..60fd4243d40 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -8,10 +8,22 @@ import (
 	"go.signoz.io/signoz/pkg/query-service/model"
 )
 
-type CloudProviderAccountsRepository interface {
+type cloudProviderAccountsRepository interface {
 	listConnectedAccounts(context.Context) ([]Account, *model.ApiError)
 }
 
+func newCloudProviderAccountsRepository(db *sqlx.DB) (
+	*cloudProviderAccountsSQLRepository, error,
+) {
+	if err := InitSqliteDBIfNeeded(db); err != nil {
+		return nil, fmt.Errorf("could not init sqlite DB for cloudintegrations: %w", err)
+	}
+
+	return &cloudProviderAccountsSQLRepository{
+		db: db,
+	}, nil
+}
+
 func InitSqliteDBIfNeeded(db *sqlx.DB) error {
 	if db == nil {
 		return fmt.Errorf("db is required")
@@ -24,7 +36,7 @@ func InitSqliteDBIfNeeded(db *sqlx.DB) error {
 			cloud_account_id TEXT,
 			last_agent_report_json TEXT,
 			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
-			removed_at TIMESTAMP,
+			removed_at TIMESTAMP
 		)
 	`
 	_, err := db.Exec(createTablesStatements)
@@ -36,3 +48,13 @@ func InitSqliteDBIfNeeded(db *sqlx.DB) error {
 
 	return nil
 }
+
+type cloudProviderAccountsSQLRepository struct {
+	db *sqlx.DB
+}
+
+func (repo *cloudProviderAccountsSQLRepository) listConnectedAccounts(context.Context) (
+	[]Account, *model.ApiError,
+) {
+	return []Account{}, nil
+}

From cc074367e2dad6193a8d6d991e758e5579f44934 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Sun, 5 Jan 2025 14:35:54 +0530
Subject: [PATCH 08/38] feat: cloud integrations:
 cloudProviderAccountsSQLRepository.listAccounts

---
 .../app/cloudintegrations/model.go            | 55 +++++++++++++++++++
 .../app/cloudintegrations/repo.go             | 25 ++++++++-
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index f600fec9624..4a1c753c72c 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -1,5 +1,60 @@
 package cloudintegrations
 
+import (
+	"database/sql/driver"
+	"encoding/json"
+	"fmt"
+	"time"
+)
+
 // Represents a cloud provider account for cloud integrations
 type Account struct {
+	Id              string        `json:"id" db:"id"`
+	Config          AccountConfig `json:"config_json" db:"config_json"`
+	CloudAccountId  string        `json:"cloud_account_id" db:"cloud_account_id"`
+	LastAgentReport AgentReport   `json:"last_agent_report_json" db:"last_agent_report_json"`
+	CreatedAt       time.Time     `json:"created_at" db:"created_at"`
+	RemovedAt       time.Time     `json:"removed_at" db:"removed_at"`
+}
+
+type AccountConfig map[string]any
+
+// For serializing from db
+func (c *AccountConfig) Scan(src interface{}) error {
+	if data, ok := src.([]byte); ok {
+		return json.Unmarshal(data, &c)
+	}
+	return nil
+}
+
+// For serializing to db
+func (c *AccountConfig) Value() (driver.Value, error) {
+	serialized, err := json.Marshal(c)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"couldn't serialize cloud account config to JSON: %w", err,
+		)
+	}
+	return serialized, nil
+}
+
+type AgentReport map[string]any
+
+// For serializing from db
+func (c *AgentReport) Scan(src interface{}) error {
+	if data, ok := src.([]byte); ok {
+		return json.Unmarshal(data, &c)
+	}
+	return nil
+}
+
+// For serializing to db
+func (c *AgentReport) Value() (driver.Value, error) {
+	serialized, err := json.Marshal(c)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"couldn't serialize cloud account config to JSON: %w", err,
+		)
+	}
+	return serialized, nil
 }
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 60fd4243d40..20e18ea25ab 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -53,8 +53,29 @@ type cloudProviderAccountsSQLRepository struct {
 	db *sqlx.DB
 }
 
-func (repo *cloudProviderAccountsSQLRepository) listConnectedAccounts(context.Context) (
+func (r *cloudProviderAccountsSQLRepository) listConnectedAccounts(ctx context.Context) (
 	[]Account, *model.ApiError,
 ) {
-	return []Account{}, nil
+	accounts := []Account{}
+
+	err := r.db.SelectContext(
+		ctx, &accounts, `
+			select
+				id,
+				config_json,
+				cloud_account_id,
+				last_agent_report_json,
+				created_at,
+				removed_at
+			from  cloud_integrations_accounts
+			order by created_at
+		`,
+	)
+	if err != nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"could not query connected cloud accounts: %w", err,
+		))
+	}
+
+	return accounts, nil
 }

From 218b4f8dd985ee3086b7e9d738726402a24061da Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Sun, 5 Jan 2025 16:12:12 +0530
Subject: [PATCH 09/38] feat: cloud integrations: http handler and controller
 plumbing for /generate-connection-url

---
 .../app/cloudintegrations/controller.go       | 55 ++++++++++++++++++-
 pkg/query-service/app/http_handler.go         | 29 +++++++++-
 .../signoz_cloud_integrations_test.go         | 44 ++++++++++++++-
 3 files changed, 123 insertions(+), 5 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 2dd09e96262..80dd987f750 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -29,7 +29,7 @@ func NewController(db *sqlx.DB) (
 ) {
 	repo, err := newCloudProviderAccountsRepository(db)
 	if err != nil {
-		return nil, fmt.Errorf("couldn't create cloud provider accounts repo", err)
+		return nil, fmt.Errorf("couldn't create cloud provider accounts repo: %w", err)
 	}
 
 	return &Controller{
@@ -59,3 +59,56 @@ func (c *Controller) ListConnectedAccounts(
 		Accounts: accounts,
 	}, nil
 }
+
+type UpsertAccountRequest struct {
+	AccountConfig CloudAccountConfig `json:"account_config"`
+	// Optional. To be specified for updates.
+	AccountId string `json:"account_id,omitempty"`
+}
+
+type CloudAccountConfig struct {
+	EnabledRegions []string `json:"regions"`
+}
+
+func (c *Controller) UpsertAccount(
+	ctx context.Context, cloudProvider string, req UpsertAccountRequest,
+) (
+	Account, *model.ApiError,
+) {
+	var account Account
+	return account, nil
+}
+
+type GenerateConnectionUrlRequest struct {
+	UpsertAccountRequest
+
+	AgentConfig SigNozAgentConfig `json:"agent_config"`
+}
+
+type SigNozAgentConfig struct {
+	Region string `json:"region"`
+}
+
+type GenerateConnectionUrlResponse struct {
+	AccountId     string `json:"account_id"`
+	ConnectionUrl string `json:"connection_url"`
+}
+
+func (c *Controller) GenerateConnectionUrl(
+	ctx context.Context, cloudProvider string, req GenerateConnectionUrlRequest,
+) (*GenerateConnectionUrlResponse, *model.ApiError) {
+	// Account connection with a simple connection URL may not be available for all providers.
+	if cloudProvider != "aws" {
+		return nil, model.BadRequest(fmt.Errorf("unsupported cloud provider: %s", cloudProvider))
+	}
+
+	account, apiErr := c.UpsertAccount(ctx, cloudProvider, req.UpsertAccountRequest)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
+	}
+
+	return &GenerateConnectionUrlResponse{
+		AccountId:     account.Id,
+		ConnectionUrl: "",
+	}, nil
+}
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index e4e6e05518b..80793fb9600 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3868,6 +3868,10 @@ func (aH *APIHandler) UninstallIntegration(
 func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/cloud-integrations").Subrouter()
 
+	subRouter.HandleFunc(
+		"/{cloudProvider}/accounts/generate-connection-url", am.EditAccess(aH.CloudIntegrationsGenerateConnectionUrl),
+	).Methods(http.MethodPost)
+
 	subRouter.HandleFunc(
 		"/{cloudProvider}/accounts", am.ViewAccess(aH.CloudIntegrationsListConnectedAccounts),
 	).Methods(http.MethodGet)
@@ -3884,12 +3888,35 @@ func (aH *APIHandler) CloudIntegrationsListConnectedAccounts(
 	)
 
 	if apiErr != nil {
-		RespondError(w, apiErr, fmt.Sprintf("Failed to fetch connected %s accounts", cloudProvider))
+		RespondError(w, apiErr, nil)
 		return
 	}
 	aH.Respond(w, resp)
 }
 
+func (aH *APIHandler) CloudIntegrationsGenerateConnectionUrl(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+
+	req := cloudintegrations.GenerateConnectionUrlRequest{}
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
+	}
+
+	result, apiErr := aH.CloudIntegrationsController.GenerateConnectionUrl(
+		r.Context(), cloudProvider, req,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, nil)
+		return
+	}
+
+	aH.Respond(w, result)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 1e7e20f17a7..bb7ac6bc7b7 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"testing"
 
+	"github.com/davecgh/go-spew/spew"
 	"github.com/jmoiron/sqlx"
 	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
 	"github.com/stretchr/testify/require"
@@ -30,6 +31,21 @@ func TestCloudIntegrationLifecycle(t *testing.T) {
 	)
 
 	// Should be able to generate connection url - initializing an account
+	connectionUrlResp := testbed.GenerateConnectionUrlFromQS("aws", cloudintegrations.GenerateConnectionUrlRequest{
+		AgentConfig: cloudintegrations.SigNozAgentConfig{
+			Region: "us-east-1",
+		},
+		UpsertAccountRequest: cloudintegrations.UpsertAccountRequest{
+			AccountConfig: cloudintegrations.CloudAccountConfig{
+				EnabledRegions: []string{"us-east-1", "us-east-2"},
+			},
+		},
+	})
+	spew.Dump(connectionUrlResp)
+	connectionUrl := connectionUrlResp.ConnectionUrl
+	require.NotEmpty(connectionUrl)
+	accountId := connectionUrlResp.AccountId
+	require.NotEmpty(accountId)
 
 	// Should be able to poll for account connection status
 
@@ -113,13 +129,35 @@ func (tb *CloudIntegrationsTestBed) GetAccountsListFromQS(
 		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
 	}
 
-	var accountsListResp cloudintegrations.AccountsListResponse
-	err = json.Unmarshal(dataJson, &accountsListResp)
+	var resp cloudintegrations.AccountsListResponse
+	err = json.Unmarshal(dataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into AccountsListResponse")
 	}
 
-	return &accountsListResp
+	return &resp
+}
+
+func (tb *CloudIntegrationsTestBed) GenerateConnectionUrlFromQS(
+	cloudProvider string, req cloudintegrations.GenerateConnectionUrlRequest,
+) *cloudintegrations.GenerateConnectionUrlResponse {
+	result := tb.RequestQS(
+		fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts/generate-connection-url", cloudProvider),
+		req,
+	)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp cloudintegrations.GenerateConnectionUrlResponse
+	err = json.Unmarshal(dataJson, &result)
+	if err != nil {
+		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into map[string]any")
+	}
+
+	return &resp
 }
 
 func (tb *CloudIntegrationsTestBed) RequestQS(

From c3e4957ad598e67605571dccb38e97c820b2e806 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Mon, 6 Jan 2025 21:05:23 +0530
Subject: [PATCH 10/38] feat: cloud integrations:
 cloudProviderAccountsSQLRepository.upsert

---
 .../app/cloudintegrations/model.go            |  13 +-
 .../app/cloudintegrations/repo.go             | 150 +++++++++++++++++-
 2 files changed, 156 insertions(+), 7 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 4a1c753c72c..293e7770dfe 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -12,12 +12,14 @@ type Account struct {
 	Id              string        `json:"id" db:"id"`
 	Config          AccountConfig `json:"config_json" db:"config_json"`
 	CloudAccountId  string        `json:"cloud_account_id" db:"cloud_account_id"`
-	LastAgentReport AgentReport   `json:"last_agent_report_json" db:"last_agent_report_json"`
+	LastAgentReport *AgentReport  `json:"last_agent_report_json" db:"last_agent_report_json"`
 	CreatedAt       time.Time     `json:"created_at" db:"created_at"`
-	RemovedAt       time.Time     `json:"removed_at" db:"removed_at"`
+	RemovedAt       *time.Time    `json:"removed_at" db:"removed_at"`
 }
 
-type AccountConfig map[string]any
+type AccountConfig struct {
+	EnabledRegions []string `json:"regions"`
+}
 
 // For serializing from db
 func (c *AccountConfig) Scan(src interface{}) error {
@@ -38,7 +40,10 @@ func (c *AccountConfig) Value() (driver.Value, error) {
 	return serialized, nil
 }
 
-type AgentReport map[string]any
+type AgentReport struct {
+	TimestampSeconds int64          `json:"timestamp"`
+	Data             map[string]any `json:"data"`
+}
 
 // For serializing from db
 func (c *AgentReport) Scan(src interface{}) error {
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 20e18ea25ab..21b7c35e16c 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -3,13 +3,28 @@ package cloudintegrations
 import (
 	"context"
 	"fmt"
+	"strings"
+	"time"
 
+	"github.com/google/uuid"
 	"github.com/jmoiron/sqlx"
 	"go.signoz.io/signoz/pkg/query-service/model"
 )
 
 type cloudProviderAccountsRepository interface {
-	listConnectedAccounts(context.Context) ([]Account, *model.ApiError)
+	listConnected(context.Context) ([]Account, *model.ApiError)
+
+	get(ctx context.Context, ids []string) (map[string]*Account, *model.ApiError)
+
+	// Insert an account or update it by ID for specified non-empty fields
+	upsert(
+		ctx context.Context,
+		id string,
+		config *AccountConfig,
+		cloudAccountId string,
+		agentReport *AgentReport,
+		removedAt *time.Time,
+	) (*Account, *model.ApiError)
 }
 
 func newCloudProviderAccountsRepository(db *sqlx.DB) (
@@ -53,7 +68,7 @@ type cloudProviderAccountsSQLRepository struct {
 	db *sqlx.DB
 }
 
-func (r *cloudProviderAccountsSQLRepository) listConnectedAccounts(ctx context.Context) (
+func (r *cloudProviderAccountsSQLRepository) listConnected(ctx context.Context) (
 	[]Account, *model.ApiError,
 ) {
 	accounts := []Account{}
@@ -67,7 +82,10 @@ func (r *cloudProviderAccountsSQLRepository) listConnectedAccounts(ctx context.C
 				last_agent_report_json,
 				created_at,
 				removed_at
-			from  cloud_integrations_accounts
+			from cloud_integrations_accounts
+      where
+        removed_at is NULL
+        and cloud_account_id is not NULL
 			order by created_at
 		`,
 	)
@@ -79,3 +97,129 @@ func (r *cloudProviderAccountsSQLRepository) listConnectedAccounts(ctx context.C
 
 	return accounts, nil
 }
+
+func (r *cloudProviderAccountsSQLRepository) get(
+	ctx context.Context, ids []string,
+) (map[string]*Account, *model.ApiError) {
+	accounts := []Account{}
+
+	idPlaceholders := []string{}
+	idValues := []interface{}{}
+	for _, id := range ids {
+		idPlaceholders = append(idPlaceholders, "?")
+		idValues = append(idValues, id)
+	}
+
+	err := r.db.SelectContext(
+		ctx, &accounts, fmt.Sprintf(`
+			select
+				id,
+				config_json,
+				cloud_account_id,
+				last_agent_report_json,
+				created_at,
+				removed_at
+			from cloud_integrations_accounts
+			where id in (%s)`,
+			strings.Join(idPlaceholders, ", "),
+		),
+		idValues...,
+	)
+	if err != nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"could not query cloud provider account: %w", err,
+		))
+	}
+
+	result := map[string]*Account{}
+	for _, a := range accounts {
+		result[a.Id] = &a
+	}
+
+	return result, nil
+}
+
+func (r *cloudProviderAccountsSQLRepository) upsert(
+	ctx context.Context,
+	id string,
+	config *AccountConfig,
+	cloudAccountId string,
+	agentReport *AgentReport,
+	removedAt *time.Time,
+) (*Account, *model.ApiError) {
+	// Insert
+	if len(id) < 1 {
+		// config must be specified when inserting
+		if config == nil {
+			return nil, model.BadRequest(fmt.Errorf("account config is required"))
+		}
+
+		id = uuid.NewString()
+	}
+
+	// Prepare clause for setting values in `on conflict do update`
+	onConflictUpdates := []string{}
+	updateColStmt := func(col string) string {
+		return fmt.Sprintf("set %s=excluded.%s", col, col)
+	}
+
+	if config != nil {
+		onConflictUpdates = append(
+			onConflictUpdates, updateColStmt("config_json"),
+		)
+	}
+
+	if len(cloudAccountId) > 0 {
+		onConflictUpdates = append(
+			onConflictUpdates, updateColStmt("cloud_account_id"),
+		)
+	}
+
+	if agentReport != nil {
+		onConflictUpdates = append(
+			onConflictUpdates, updateColStmt("last_agent_report_json"),
+		)
+	}
+
+	if removedAt != nil {
+		onConflictUpdates = append(
+			onConflictUpdates, updateColStmt("removed_at"),
+		)
+	}
+
+	_, dbErr := r.db.ExecContext(
+		ctx, fmt.Sprintf(`
+			INSERT INTO cloud_integrations_accounts (
+				id,
+				config_json,
+				cloud_account_id,
+				last_agent_report_json,
+				removed_at
+			) values ($1, $2, $3, $4, $5)
+			on conflict(id) do update
+      %s
+		`, strings.Join(onConflictUpdates, "\n")),
+		id, config, cloudAccountId, agentReport, removedAt,
+	)
+	if dbErr != nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"could not upsert cloud account record: %w", dbErr,
+		))
+	}
+
+	res, apiErr := r.get(ctx, []string{id})
+	if apiErr != nil {
+		return nil, model.WrapApiError(
+			apiErr, "couldn't fetch upserted account by id",
+		)
+	}
+
+	upsertedAccount := res[id]
+	if upsertedAccount == nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"couldn't fetch upserted account by id",
+		))
+	}
+
+	return upsertedAccount, nil
+}

From 21d9253a236b0fc7d2b01cb8bf23f56b786840ef Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Mon, 6 Jan 2025 21:06:14 +0530
Subject: [PATCH 11/38] feat: cloud integrations: finish up with
 /generate-connection-url

---
 .../app/cloudintegrations/controller.go       | 35 ++++++++-----------
 .../signoz_cloud_integrations_test.go         | 14 +++-----
 2 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 80dd987f750..e98f1370bc3 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -50,7 +50,7 @@ func (c *Controller) ListConnectedAccounts(
 		return nil, apiErr
 	}
 
-	accounts, apiErr := c.repo.listConnectedAccounts(ctx)
+	accounts, apiErr := c.repo.listConnected(ctx)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't list cloud accounts")
 	}
@@ -60,27 +60,11 @@ func (c *Controller) ListConnectedAccounts(
 	}, nil
 }
 
-type UpsertAccountRequest struct {
-	AccountConfig CloudAccountConfig `json:"account_config"`
+type GenerateConnectionUrlRequest struct {
 	// Optional. To be specified for updates.
 	AccountId string `json:"account_id,omitempty"`
-}
-
-type CloudAccountConfig struct {
-	EnabledRegions []string `json:"regions"`
-}
-
-func (c *Controller) UpsertAccount(
-	ctx context.Context, cloudProvider string, req UpsertAccountRequest,
-) (
-	Account, *model.ApiError,
-) {
-	var account Account
-	return account, nil
-}
 
-type GenerateConnectionUrlRequest struct {
-	UpsertAccountRequest
+	AccountConfig AccountConfig `json:"account_config"`
 
 	AgentConfig SigNozAgentConfig `json:"agent_config"`
 }
@@ -102,13 +86,22 @@ func (c *Controller) GenerateConnectionUrl(
 		return nil, model.BadRequest(fmt.Errorf("unsupported cloud provider: %s", cloudProvider))
 	}
 
-	account, apiErr := c.UpsertAccount(ctx, cloudProvider, req.UpsertAccountRequest)
+	account, apiErr := c.repo.upsert(
+		ctx, req.AccountId, &req.AccountConfig, "", nil, nil,
+	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
 	}
 
+	// TODO(Raj): Add actual connection url after shipping
+	// cloudformation template for AWS integration
+	connectionUrl := fmt.Sprintf(
+		"https://%s.console.aws.amazon.com/cloudformation/home?region=%s#/stacks/quickcreate?stackName=SigNozIntegration/",
+		req.AgentConfig.Region, req.AgentConfig.Region,
+	)
+
 	return &GenerateConnectionUrlResponse{
 		AccountId:     account.Id,
-		ConnectionUrl: "",
+		ConnectionUrl: connectionUrl,
 	}, nil
 }
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index bb7ac6bc7b7..635816dfaec 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"testing"
 
-	"github.com/davecgh/go-spew/spew"
 	"github.com/jmoiron/sqlx"
 	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
 	"github.com/stretchr/testify/require"
@@ -35,17 +34,14 @@ func TestCloudIntegrationLifecycle(t *testing.T) {
 		AgentConfig: cloudintegrations.SigNozAgentConfig{
 			Region: "us-east-1",
 		},
-		UpsertAccountRequest: cloudintegrations.UpsertAccountRequest{
-			AccountConfig: cloudintegrations.CloudAccountConfig{
-				EnabledRegions: []string{"us-east-1", "us-east-2"},
-			},
+		AccountConfig: cloudintegrations.AccountConfig{
+			EnabledRegions: []string{"us-east-1", "us-east-2"},
 		},
 	})
-	spew.Dump(connectionUrlResp)
-	connectionUrl := connectionUrlResp.ConnectionUrl
-	require.NotEmpty(connectionUrl)
 	accountId := connectionUrlResp.AccountId
 	require.NotEmpty(accountId)
+	connectionUrl := connectionUrlResp.ConnectionUrl
+	require.NotEmpty(connectionUrl)
 
 	// Should be able to poll for account connection status
 
@@ -152,7 +148,7 @@ func (tb *CloudIntegrationsTestBed) GenerateConnectionUrlFromQS(
 	}
 
 	var resp cloudintegrations.GenerateConnectionUrlResponse
-	err = json.Unmarshal(dataJson, &result)
+	err = json.Unmarshal(dataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into map[string]any")
 	}

From 19fca6ffddcf461f03daf22cd7d7a2d5c42b5815 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 10:55:40 +0530
Subject: [PATCH 12/38] feat: cloud integrations: add
 cloudProviderAccountsRepository.get

---
 .../app/cloudintegrations/repo.go             | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 21b7c35e16c..221e445cd0d 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -14,7 +14,9 @@ import (
 type cloudProviderAccountsRepository interface {
 	listConnected(context.Context) ([]Account, *model.ApiError)
 
-	get(ctx context.Context, ids []string) (map[string]*Account, *model.ApiError)
+	getByIds(ctx context.Context, ids []string) (map[string]*Account, *model.ApiError)
+
+	get(ctx context.Context, id string) (*Account, *model.ApiError)
 
 	// Insert an account or update it by ID for specified non-empty fields
 	upsert(
@@ -98,7 +100,7 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(ctx context.Context)
 	return accounts, nil
 }
 
-func (r *cloudProviderAccountsSQLRepository) get(
+func (r *cloudProviderAccountsSQLRepository) getByIds(
 	ctx context.Context, ids []string,
 ) (map[string]*Account, *model.ApiError) {
 	accounts := []Account{}
@@ -139,6 +141,24 @@ func (r *cloudProviderAccountsSQLRepository) get(
 	return result, nil
 }
 
+func (r *cloudProviderAccountsSQLRepository) get(
+	ctx context.Context, id string,
+) (*Account, *model.ApiError) {
+	res, apiErr := r.getByIds(ctx, []string{id})
+	if apiErr != nil {
+		return nil, apiErr
+	}
+
+	account := res[id]
+	if account == nil {
+		return nil, model.NotFoundError(fmt.Errorf(
+			"couldn't find account with Id %s", id,
+		))
+	}
+
+	return account, nil
+}
+
 func (r *cloudProviderAccountsSQLRepository) upsert(
 	ctx context.Context,
 	id string,
@@ -207,17 +227,10 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 		))
 	}
 
-	res, apiErr := r.get(ctx, []string{id})
+	upsertedAccount, apiErr := r.get(ctx, id)
 	if apiErr != nil {
-		return nil, model.WrapApiError(
-			apiErr, "couldn't fetch upserted account by id",
-		)
-	}
-
-	upsertedAccount := res[id]
-	if upsertedAccount == nil {
 		return nil, model.InternalError(fmt.Errorf(
-			"couldn't fetch upserted account by id",
+			"couldn't fetch upserted account by id: %w", apiErr.ToError(),
 		))
 	}
 

From 93a7087ebc39ac2b849b9100460a96de464ff81e Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 10:59:29 +0530
Subject: [PATCH 13/38] feat: cloud integrations: add API test expectation for
 being able to get account status

---
 .../signoz_cloud_integrations_test.go         | 33 ++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 635816dfaec..e7fce9887bf 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -18,8 +18,8 @@ import (
 	"go.signoz.io/signoz/pkg/query-service/utils"
 )
 
-func TestCloudIntegrationLifecycle(t *testing.T) {
-	// Test for the happy path of connecting and managing cloud integration accounts
+func TestAWSIntegrationLifecycle(t *testing.T) {
+	// Test for the happy path of connecting and managing AWS integration accounts
 
 	require := require.New(t)
 	testbed := NewCloudIntegrationsTestBed(t, nil)
@@ -44,6 +44,9 @@ func TestCloudIntegrationLifecycle(t *testing.T) {
 	require.NotEmpty(connectionUrl)
 
 	// Should be able to poll for account connection status
+	accountStatusResp := testbed.GetAccountStatusFromQS("aws", accountId)
+	require.Equal(accountId, accountStatusResp.Id)
+	require.Nil(accountStatusResp.Status.Integration.LastHeartbeatTsMillis)
 
 	// The unconnected account should not show up in accounts list yet
 
@@ -128,7 +131,7 @@ func (tb *CloudIntegrationsTestBed) GetAccountsListFromQS(
 	var resp cloudintegrations.AccountsListResponse
 	err = json.Unmarshal(dataJson, &resp)
 	if err != nil {
-		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into AccountsListResponse")
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AccountsListResponse")
 	}
 
 	return &resp
@@ -150,7 +153,29 @@ func (tb *CloudIntegrationsTestBed) GenerateConnectionUrlFromQS(
 	var resp cloudintegrations.GenerateConnectionUrlResponse
 	err = json.Unmarshal(dataJson, &resp)
 	if err != nil {
-		tb.t.Fatalf(" could not unmarshal apiResponse.Data json into map[string]any")
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into map[string]any")
+	}
+
+	return &resp
+}
+
+func (tb *CloudIntegrationsTestBed) GetAccountStatusFromQS(
+	cloudProvider string, accountId string,
+) *cloudintegrations.AccountStatusResponse {
+	result := tb.RequestQS(fmt.Sprintf(
+		"/api/v1/cloud-integrations/%s/accounts/%s/status",
+		cloudProvider, accountId,
+	), nil)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp cloudintegrations.AccountStatusResponse
+	err = json.Unmarshal(dataJson, &resp)
+	if err != nil {
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AccountStatusResponse")
 	}
 
 	return &resp

From 97156ddaf9b976c460e15801df3b17508bd0a72f Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 11:00:04 +0530
Subject: [PATCH 14/38] feat: cloud integrations: add http handler and
 controller method for getting account status

---
 .../app/cloudintegrations/controller.go       | 40 ++++++++++++++++++-
 .../app/cloudintegrations/repo.go             |  1 +
 pkg/query-service/app/http_handler.go         | 21 ++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index e98f1370bc3..1601e927e0b 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -93,8 +93,7 @@ func (c *Controller) GenerateConnectionUrl(
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
 	}
 
-	// TODO(Raj): Add actual connection url after shipping
-	// cloudformation template for AWS integration
+	// TODO(Raj): Add actual cloudformation template for AWS integration after it has been shipped.
 	connectionUrl := fmt.Sprintf(
 		"https://%s.console.aws.amazon.com/cloudformation/home?region=%s#/stacks/quickcreate?stackName=SigNozIntegration/",
 		req.AgentConfig.Region, req.AgentConfig.Region,
@@ -105,3 +104,40 @@ func (c *Controller) GenerateConnectionUrl(
 		ConnectionUrl: connectionUrl,
 	}, nil
 }
+
+type AccountStatusResponse struct {
+	Id     string        `json:"id"`
+	Status AccountStatus `json:"status"`
+}
+
+type AccountStatus struct {
+	Integration AccountIntegrationStatus `json:"integration"`
+}
+
+type AccountIntegrationStatus struct {
+	LastHeartbeatTsMillis *int64 `json:"last_heartbeat_ts_ms"`
+}
+
+func (c *Controller) GetAccountStatus(
+	ctx context.Context, cloudProvider string, accountId string,
+) (
+	*AccountStatusResponse, *model.ApiError,
+) {
+	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
+		return nil, apiErr
+	}
+
+	account, apiErr := c.repo.get(ctx, accountId)
+	if apiErr != nil {
+		return nil, apiErr
+	}
+
+	return &AccountStatusResponse{
+		Id: account.Id,
+		Status: AccountStatus{
+			Integration: AccountIntegrationStatus{
+				LastHeartbeatTsMillis: nil,
+			},
+		},
+	}, nil
+}
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 221e445cd0d..35f95435380 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -12,6 +12,7 @@ import (
 )
 
 type cloudProviderAccountsRepository interface {
+	// TODO(Raj): All methods should be scoped by cloud provider.
 	listConnected(context.Context) ([]Account, *model.ApiError)
 
 	getByIds(ctx context.Context, ids []string) (map[string]*Account, *model.ApiError)
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index 80793fb9600..a3d696b50d1 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3876,6 +3876,10 @@ func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *Au
 		"/{cloudProvider}/accounts", am.ViewAccess(aH.CloudIntegrationsListConnectedAccounts),
 	).Methods(http.MethodGet)
 
+	subRouter.HandleFunc(
+		"/{cloudProvider}/accounts/{accountId}/status", am.ViewAccess(aH.CloudIntegrationsGetAccountStatus),
+	).Methods(http.MethodGet)
+
 }
 
 func (aH *APIHandler) CloudIntegrationsListConnectedAccounts(
@@ -3917,6 +3921,23 @@ func (aH *APIHandler) CloudIntegrationsGenerateConnectionUrl(
 	aH.Respond(w, result)
 }
 
+func (aH *APIHandler) CloudIntegrationsGetAccountStatus(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+	accountId := mux.Vars(r)["accountId"]
+
+	resp, apiErr := aH.CloudIntegrationsController.GetAccountStatus(
+		r.Context(), cloudProvider, accountId,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, nil)
+		return
+	}
+	aH.Respond(w, resp)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()

From 946610b67805b4d46727e02b6739b238fc19041a Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 12:50:18 +0530
Subject: [PATCH 15/38] feat: cloud integrations: ensure unconnected accounts
 aren't included in list of connected accounts

---
 .../app/cloudintegrations/controller.go        |  4 ++--
 .../app/cloudintegrations/model.go             |  2 +-
 .../app/cloudintegrations/repo.go              | 18 ++++++++++--------
 .../signoz_cloud_integrations_test.go          |  4 ++++
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 1601e927e0b..5aae3dfc2d4 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -62,7 +62,7 @@ func (c *Controller) ListConnectedAccounts(
 
 type GenerateConnectionUrlRequest struct {
 	// Optional. To be specified for updates.
-	AccountId string `json:"account_id,omitempty"`
+	AccountId *string `json:"account_id,omitempty"`
 
 	AccountConfig AccountConfig `json:"account_config"`
 
@@ -87,7 +87,7 @@ func (c *Controller) GenerateConnectionUrl(
 	}
 
 	account, apiErr := c.repo.upsert(
-		ctx, req.AccountId, &req.AccountConfig, "", nil, nil,
+		ctx, req.AccountId, &req.AccountConfig, nil, nil, nil,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 293e7770dfe..93f250aa873 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -11,7 +11,7 @@ import (
 type Account struct {
 	Id              string        `json:"id" db:"id"`
 	Config          AccountConfig `json:"config_json" db:"config_json"`
-	CloudAccountId  string        `json:"cloud_account_id" db:"cloud_account_id"`
+	CloudAccountId  *string       `json:"cloud_account_id" db:"cloud_account_id"`
 	LastAgentReport *AgentReport  `json:"last_agent_report_json" db:"last_agent_report_json"`
 	CreatedAt       time.Time     `json:"created_at" db:"created_at"`
 	RemovedAt       *time.Time    `json:"removed_at" db:"removed_at"`
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 35f95435380..bdc2d400bd8 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -22,9 +22,9 @@ type cloudProviderAccountsRepository interface {
 	// Insert an account or update it by ID for specified non-empty fields
 	upsert(
 		ctx context.Context,
-		id string,
+		id *string,
 		config *AccountConfig,
-		cloudAccountId string,
+		cloudAccountId *string,
 		agentReport *AgentReport,
 		removedAt *time.Time,
 	) (*Account, *model.ApiError)
@@ -162,20 +162,22 @@ func (r *cloudProviderAccountsSQLRepository) get(
 
 func (r *cloudProviderAccountsSQLRepository) upsert(
 	ctx context.Context,
-	id string,
+	id *string,
 	config *AccountConfig,
-	cloudAccountId string,
+	cloudAccountId *string,
 	agentReport *AgentReport,
 	removedAt *time.Time,
 ) (*Account, *model.ApiError) {
 	// Insert
-	if len(id) < 1 {
+	if id == nil {
 		// config must be specified when inserting
 		if config == nil {
 			return nil, model.BadRequest(fmt.Errorf("account config is required"))
 		}
 
-		id = uuid.NewString()
+		newId := uuid.NewString()
+
+		id = &newId
 	}
 
 	// Prepare clause for setting values in `on conflict do update`
@@ -190,7 +192,7 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 		)
 	}
 
-	if len(cloudAccountId) > 0 {
+	if cloudAccountId != nil {
 		onConflictUpdates = append(
 			onConflictUpdates, updateColStmt("cloud_account_id"),
 		)
@@ -228,7 +230,7 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 		))
 	}
 
-	upsertedAccount, apiErr := r.get(ctx, id)
+	upsertedAccount, apiErr := r.get(ctx, *id)
 	if apiErr != nil {
 		return nil, model.InternalError(fmt.Errorf(
 			"couldn't fetch upserted account by id: %w", apiErr.ToError(),
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index e7fce9887bf..5f30dd1ee7f 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -49,6 +49,10 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Nil(accountStatusResp.Status.Integration.LastHeartbeatTsMillis)
 
 	// The unconnected account should not show up in accounts list yet
+	accountsListResp1 := testbed.GetAccountsListFromQS("aws")
+	require.Equal(0, len(accountsListResp1.Accounts),
+		"No accounts should be connected at the beginning",
+	)
 
 	// An agent should be able to check in to the new account
 	// Should have the settings that were specified while generating connection url

From 2cae20297be07ffb5e1ded77e0f0b15a263d1f01 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 13:36:56 +0530
Subject: [PATCH 16/38] feat: cloud integrations: add test expectation for
 agent check in request

---
 .../signoz_cloud_integrations_test.go         | 63 +++++++++++++++----
 1 file changed, 50 insertions(+), 13 deletions(-)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 5f30dd1ee7f..d7bf73e7604 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/jmoiron/sqlx"
 	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
@@ -21,6 +22,7 @@ import (
 func TestAWSIntegrationLifecycle(t *testing.T) {
 	// Test for the happy path of connecting and managing AWS integration accounts
 
+	t0 := time.Now()
 	require := require.New(t)
 	testbed := NewCloudIntegrationsTestBed(t, nil)
 
@@ -30,22 +32,24 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	)
 
 	// Should be able to generate connection url - initializing an account
-	connectionUrlResp := testbed.GenerateConnectionUrlFromQS("aws", cloudintegrations.GenerateConnectionUrlRequest{
-		AgentConfig: cloudintegrations.SigNozAgentConfig{
-			Region: "us-east-1",
-		},
-		AccountConfig: cloudintegrations.AccountConfig{
-			EnabledRegions: []string{"us-east-1", "us-east-2"},
-		},
-	})
-	accountId := connectionUrlResp.AccountId
-	require.NotEmpty(accountId)
+	testAccountConfig := cloudintegrations.AccountConfig{
+		EnabledRegions: []string{"us-east-1", "us-east-2"},
+	}
+	connectionUrlResp := testbed.GenerateConnectionUrlFromQS(
+		"aws", cloudintegrations.GenerateConnectionUrlRequest{
+			AgentConfig: cloudintegrations.SigNozAgentConfig{
+				Region: "us-east-1",
+			},
+			AccountConfig: testAccountConfig,
+		})
+	testAccountId := connectionUrlResp.AccountId
+	require.NotEmpty(testAccountId)
 	connectionUrl := connectionUrlResp.ConnectionUrl
 	require.NotEmpty(connectionUrl)
 
 	// Should be able to poll for account connection status
-	accountStatusResp := testbed.GetAccountStatusFromQS("aws", accountId)
-	require.Equal(accountId, accountStatusResp.Id)
+	accountStatusResp := testbed.GetAccountStatusFromQS("aws", testAccountId)
+	require.Equal(testAccountId, accountStatusResp.Id)
 	require.Nil(accountStatusResp.Status.Integration.LastHeartbeatTsMillis)
 
 	// The unconnected account should not show up in accounts list yet
@@ -55,7 +59,19 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	)
 
 	// An agent should be able to check in to the new account
-	// Should have the settings that were specified while generating connection url
+	// Should get the settings that were specified while generating connection url
+	testAWSAccountId := "4563215233"
+	agentCheckInResp := testbed.CheckInAsAgentWithQS(
+		"aws", cloudintegrations.AgentCheckInRequest{
+			AccountId:      testAccountId,
+			CloudAccountId: testAWSAccountId,
+		},
+	)
+	require.Equal(testAccountId, agentCheckInResp.Account.Id)
+	require.Equal(testAccountConfig, agentCheckInResp.Account.Config)
+	require.Equal(testAWSAccountId, agentCheckInResp.Account.CloudAccountId)
+	require.LessOrEqual(t0, agentCheckInResp.Account.CreatedAt)
+	require.Nil(agentCheckInResp.Account.RemovedAt)
 
 	// Polling for connection status should return latest status now.
 
@@ -185,6 +201,27 @@ func (tb *CloudIntegrationsTestBed) GetAccountStatusFromQS(
 	return &resp
 }
 
+func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
+	cloudProvider string, req cloudintegrations.AgentCheckInRequest,
+) *cloudintegrations.AgentCheckInResponse {
+	result := tb.RequestQS(
+		fmt.Sprintf("/api/v1/cloud-integrations/%s/agent-check-in", cloudProvider), req,
+	)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp cloudintegrations.AgentCheckInResponse
+	err = json.Unmarshal(dataJson, &resp)
+	if err != nil {
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AgentCheckInResponse")
+	}
+
+	return &resp
+}
+
 func (tb *CloudIntegrationsTestBed) RequestQS(
 	path string,
 	postData interface{},

From e3ff59f3aa238622a108e1122a99ef5c1639df67 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 13:54:23 +0530
Subject: [PATCH 17/38] feat: cloud integrations: agent check in API

---
 .../app/cloudintegrations/controller.go       | 37 +++++++++++
 .../app/cloudintegrations/model.go            | 18 +++--
 .../app/cloudintegrations/repo.go             | 65 +++++++++++--------
 pkg/query-service/app/http_handler.go         | 27 ++++++++
 .../signoz_cloud_integrations_test.go         |  6 +-
 5 files changed, 116 insertions(+), 37 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 5aae3dfc2d4..96ea4b4fe6b 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"slices"
+	"time"
 
 	"github.com/jmoiron/sqlx"
 	"go.signoz.io/signoz/pkg/query-service/model"
@@ -141,3 +142,39 @@ func (c *Controller) GetAccountStatus(
 		},
 	}, nil
 }
+
+type AgentCheckInRequest struct {
+	AccountId      string `json:"account_id"`
+	CloudAccountId string `json:"cloud_account_id"`
+	// Arbitrary cloud specific Agent data
+	Data map[string]any `json:"data,omitempty"`
+}
+
+type AgentCheckInResponse struct {
+	Account Account `json:"account"`
+}
+
+func (c *Controller) CheckInAsAgent(
+	ctx context.Context, cloudProvider string, req AgentCheckInRequest,
+) (*AgentCheckInResponse, *model.ApiError) {
+	// Account connection with a simple connection URL may not be available for all providers.
+	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
+		return nil, apiErr
+	}
+
+	agentReport := AgentReport{
+		TimestampSeconds: time.Now().Unix(),
+		Data:             req.Data,
+	}
+
+	account, apiErr := c.repo.upsert(
+		ctx, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
+	)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
+	}
+
+	return &AgentCheckInResponse{
+		Account: *account,
+	}, nil
+}
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 93f250aa873..aa5c9aa7d04 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -9,18 +9,24 @@ import (
 
 // Represents a cloud provider account for cloud integrations
 type Account struct {
-	Id              string        `json:"id" db:"id"`
-	Config          AccountConfig `json:"config_json" db:"config_json"`
-	CloudAccountId  *string       `json:"cloud_account_id" db:"cloud_account_id"`
-	LastAgentReport *AgentReport  `json:"last_agent_report_json" db:"last_agent_report_json"`
-	CreatedAt       time.Time     `json:"created_at" db:"created_at"`
-	RemovedAt       *time.Time    `json:"removed_at" db:"removed_at"`
+	Id              string         `json:"id" db:"id"`
+	Config          *AccountConfig `json:"config_json" db:"config_json"`
+	CloudAccountId  *string        `json:"cloud_account_id" db:"cloud_account_id"`
+	LastAgentReport *AgentReport   `json:"last_agent_report_json" db:"last_agent_report_json"`
+	CreatedAt       time.Time      `json:"created_at" db:"created_at"`
+	RemovedAt       *time.Time     `json:"removed_at" db:"removed_at"`
 }
 
 type AccountConfig struct {
 	EnabledRegions []string `json:"regions"`
 }
 
+func DefaultAccountConfig() AccountConfig {
+	return AccountConfig{
+		EnabledRegions: []string{},
+	}
+}
+
 // For serializing from db
 func (c *AccountConfig) Scan(src interface{}) error {
 	if data, ok := src.([]byte); ok {
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index bdc2d400bd8..46980ce642c 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -50,7 +50,7 @@ func InitSqliteDBIfNeeded(db *sqlx.DB) error {
 	createTablesStatements := `
 		CREATE TABLE IF NOT EXISTS cloud_integrations_accounts(
 			id TEXT PRIMARY KEY NOT NULL,
-			config_json TEXT NOT NULL,
+			config_json TEXT,
 			cloud_account_id TEXT,
 			last_agent_report_json TEXT,
 			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
@@ -136,6 +136,12 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 
 	result := map[string]*Account{}
 	for _, a := range accounts {
+
+		if a.Config == nil {
+			config := DefaultAccountConfig()
+			a.Config = &config
+		}
+
 		result[a.Id] = &a
 	}
 
@@ -170,58 +176,61 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 ) (*Account, *model.ApiError) {
 	// Insert
 	if id == nil {
-		// config must be specified when inserting
-		if config == nil {
-			return nil, model.BadRequest(fmt.Errorf("account config is required"))
-		}
-
 		newId := uuid.NewString()
-
 		id = &newId
 	}
 
 	// Prepare clause for setting values in `on conflict do update`
-	onConflictUpdates := []string{}
+	onConflictSetStmts := []string{}
 	updateColStmt := func(col string) string {
-		return fmt.Sprintf("set %s=excluded.%s", col, col)
+		return fmt.Sprintf("%s=excluded.%s", col, col)
 	}
 
 	if config != nil {
-		onConflictUpdates = append(
-			onConflictUpdates, updateColStmt("config_json"),
+		onConflictSetStmts = append(
+			onConflictSetStmts, updateColStmt("config_json"),
 		)
 	}
 
 	if cloudAccountId != nil {
-		onConflictUpdates = append(
-			onConflictUpdates, updateColStmt("cloud_account_id"),
+		onConflictSetStmts = append(
+			onConflictSetStmts, updateColStmt("cloud_account_id"),
 		)
 	}
 
 	if agentReport != nil {
-		onConflictUpdates = append(
-			onConflictUpdates, updateColStmt("last_agent_report_json"),
+		onConflictSetStmts = append(
+			onConflictSetStmts, updateColStmt("last_agent_report_json"),
 		)
 	}
 
 	if removedAt != nil {
-		onConflictUpdates = append(
-			onConflictUpdates, updateColStmt("removed_at"),
+		onConflictSetStmts = append(
+			onConflictSetStmts, updateColStmt("removed_at"),
+		)
+	}
+
+	onConflictClause := ""
+	if len(onConflictSetStmts) > 0 {
+		onConflictClause = fmt.Sprintf(
+			"on conflict(id) do update SET\n%s",
+			strings.Join(onConflictSetStmts, ",\n"),
 		)
 	}
 
+	insertQuery := fmt.Sprintf(`
+    INSERT INTO cloud_integrations_accounts (
+      id,
+      config_json,
+      cloud_account_id,
+      last_agent_report_json,
+      removed_at
+    ) values ($1, $2, $3, $4, $5)
+    %s`, onConflictClause,
+	)
+
 	_, dbErr := r.db.ExecContext(
-		ctx, fmt.Sprintf(`
-			INSERT INTO cloud_integrations_accounts (
-				id,
-				config_json,
-				cloud_account_id,
-				last_agent_report_json,
-				removed_at
-			) values ($1, $2, $3, $4, $5)
-			on conflict(id) do update
-      %s
-		`, strings.Join(onConflictUpdates, "\n")),
+		ctx, insertQuery,
 		id, config, cloudAccountId, agentReport, removedAt,
 	)
 	if dbErr != nil {
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index a3d696b50d1..72c3392ef33 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3880,6 +3880,10 @@ func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *Au
 		"/{cloudProvider}/accounts/{accountId}/status", am.ViewAccess(aH.CloudIntegrationsGetAccountStatus),
 	).Methods(http.MethodGet)
 
+	subRouter.HandleFunc(
+		"/{cloudProvider}/agent-check-in", am.EditAccess(aH.CloudIntegrationsAgentCheckIn),
+	).Methods(http.MethodPost)
+
 }
 
 func (aH *APIHandler) CloudIntegrationsListConnectedAccounts(
@@ -3938,6 +3942,29 @@ func (aH *APIHandler) CloudIntegrationsGetAccountStatus(
 	aH.Respond(w, resp)
 }
 
+func (aH *APIHandler) CloudIntegrationsAgentCheckIn(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+
+	req := cloudintegrations.AgentCheckInRequest{}
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
+	}
+
+	result, apiErr := aH.CloudIntegrationsController.CheckInAsAgent(
+		r.Context(), cloudProvider, req,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, nil)
+		return
+	}
+
+	aH.Respond(w, result)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index d7bf73e7604..49f2ab5e88d 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -68,9 +68,9 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 		},
 	)
 	require.Equal(testAccountId, agentCheckInResp.Account.Id)
-	require.Equal(testAccountConfig, agentCheckInResp.Account.Config)
-	require.Equal(testAWSAccountId, agentCheckInResp.Account.CloudAccountId)
-	require.LessOrEqual(t0, agentCheckInResp.Account.CreatedAt)
+	require.Equal(testAccountConfig, *agentCheckInResp.Account.Config)
+	require.Equal(testAWSAccountId, *agentCheckInResp.Account.CloudAccountId)
+	require.LessOrEqual(t0.Unix(), agentCheckInResp.Account.CreatedAt.Unix())
 	require.Nil(agentCheckInResp.Account.RemovedAt)
 
 	// Polling for connection status should return latest status now.

From e4ee24b166f9806d22c331ff79b93e40249a3856 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 19:41:39 +0530
Subject: [PATCH 18/38] feat: cloud integrations: ensure polling for status
 after agent check in works

---
 .../app/cloudintegrations/controller.go       | 24 +++++++++++++++----
 .../app/cloudintegrations/model.go            |  4 ++--
 .../signoz_cloud_integrations_test.go         | 10 +++++++-
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 96ea4b4fe6b..12b8c0923e4 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -133,14 +133,26 @@ func (c *Controller) GetAccountStatus(
 		return nil, apiErr
 	}
 
-	return &AccountStatusResponse{
+	var lastAgentReportTsMillis *int64
+	if account.LastAgentReport != nil {
+		// TODO(Raj): Investigate why agent reports are scanned to empty values
+		// from NULL columns even when report is a pointer
+		lastReportTsMillis := account.LastAgentReport.TimestampMillis
+		if lastReportTsMillis > 0 {
+			lastAgentReportTsMillis = &lastReportTsMillis
+		}
+	}
+
+	resp := AccountStatusResponse{
 		Id: account.Id,
 		Status: AccountStatus{
 			Integration: AccountIntegrationStatus{
-				LastHeartbeatTsMillis: nil,
+				LastHeartbeatTsMillis: lastAgentReportTsMillis,
 			},
 		},
-	}, nil
+	}
+
+	return &resp, nil
 }
 
 type AgentCheckInRequest struct {
@@ -163,10 +175,12 @@ func (c *Controller) CheckInAsAgent(
 	}
 
 	agentReport := AgentReport{
-		TimestampSeconds: time.Now().Unix(),
-		Data:             req.Data,
+		TimestampMillis: time.Now().UnixMilli(),
+		Data:            req.Data,
 	}
 
+	// TODO(Raj): Consider ensuring cloudAccountId is not updated post insertion
+	// Consider getting the account by Id first to validate
 	account, apiErr := c.repo.upsert(
 		ctx, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
 	)
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index aa5c9aa7d04..bd22b99922e 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -47,8 +47,8 @@ func (c *AccountConfig) Value() (driver.Value, error) {
 }
 
 type AgentReport struct {
-	TimestampSeconds int64          `json:"timestamp"`
-	Data             map[string]any `json:"data"`
+	TimestampMillis int64          `json:"timestamp_millis"`
+	Data            map[string]any `json:"data"`
 }
 
 // For serializing from db
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 49f2ab5e88d..a81b5ef0011 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -60,6 +60,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 
 	// An agent should be able to check in to the new account
 	// Should get the settings that were specified while generating connection url
+	tsMillisBeforeAgentCheckIn := time.Now().UnixMilli()
 	testAWSAccountId := "4563215233"
 	agentCheckInResp := testbed.CheckInAsAgentWithQS(
 		"aws", cloudintegrations.AgentCheckInRequest{
@@ -73,7 +74,14 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.LessOrEqual(t0.Unix(), agentCheckInResp.Account.CreatedAt.Unix())
 	require.Nil(agentCheckInResp.Account.RemovedAt)
 
-	// Polling for connection status should return latest status now.
+	// Polling for connection status should now return latest status
+	accountStatusResp1 := testbed.GetAccountStatusFromQS("aws", testAccountId)
+	require.Equal(testAccountId, accountStatusResp1.Id)
+	require.NotNil(accountStatusResp1.Status.Integration.LastHeartbeatTsMillis)
+	require.LessOrEqual(
+		tsMillisBeforeAgentCheckIn,
+		*accountStatusResp1.Status.Integration.LastHeartbeatTsMillis,
+	)
 
 	// The account should now show up in list of connected accounts.
 

From 13d663f3005a31a98bf3cb01480bd4d5aba20e80 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 19:44:40 +0530
Subject: [PATCH 19/38] feat: cloud integrations: ensure account included in
 connected account list after agent check in

---
 .../tests/integration/signoz_cloud_integrations_test.go     | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index a81b5ef0011..1bbcaf58b6a 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -84,6 +84,12 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	)
 
 	// The account should now show up in list of connected accounts.
+	accountsListResp2 := testbed.GetAccountsListFromQS("aws")
+	require.Equal(len(accountsListResp2.Accounts), 1,
+		"No accounts should be connected at the beginning",
+	)
+	require.Equal(testAccountId, accountsListResp2.Accounts[0].Id)
+	require.Equal(testAWSAccountId, *accountsListResp2.Accounts[0].CloudAccountId)
 
 	// Should be able to update account settings
 

From c81c05d81f2cba8504d1379ae813741824b3848c Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 20:45:28 +0530
Subject: [PATCH 20/38] feat: cloud integrations: add API expectation for
 updating account config

---
 .../signoz_cloud_integrations_test.go         | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 1bbcaf58b6a..2552b36d10b 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -92,6 +92,14 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Equal(testAWSAccountId, *accountsListResp2.Accounts[0].CloudAccountId)
 
 	// Should be able to update account settings
+	testAccountConfig2 := cloudintegrations.AccountConfig{
+		EnabledRegions: []string{"us-east-2", "us-west-1"},
+	}
+	latestAccount := testbed.UpdateAccountSettingsWithQS(
+		"aws", testAccountId, testAccountConfig2,
+	)
+	require.Equal(testAccountId, latestAccount.Id)
+	require.Equal(testAccountConfig2, latestAccount.Config)
 
 	// The agent should now receive latest settings.
 
@@ -236,6 +244,32 @@ func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
 	return &resp
 }
 
+func (tb *CloudIntegrationsTestBed) UpdateAccountSettingsWithQS(
+	cloudProvider string, accountId string, newConfig cloudintegrations.AccountConfig,
+) *cloudintegrations.Account {
+	result := tb.RequestQS(
+		fmt.Sprintf(
+			"/api/v1/cloud-integrations/%s/accounts/%s/config",
+			cloudProvider, accountId,
+		), cloudintegrations.UpdateAccountSettingsRequest{
+			Config: newConfig,
+		},
+	)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp cloudintegrations.Account
+	err = json.Unmarshal(dataJson, &resp)
+	if err != nil {
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AgentCheckInResponse")
+	}
+
+	return &resp
+}
+
 func (tb *CloudIntegrationsTestBed) RequestQS(
 	path string,
 	postData interface{},

From 9029cb2c25d91ca5457551b477d9a78946cf6b06 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 20:53:23 +0530
Subject: [PATCH 21/38] feat: cloud integrations: API for updating cloud
 account config

---
 .../app/cloudintegrations/controller.go       | 25 +++++++++++++++++
 pkg/query-service/app/http_handler.go         | 28 +++++++++++++++++++
 .../signoz_cloud_integrations_test.go         |  8 +++---
 3 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 12b8c0923e4..eb23f66a186 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -39,6 +39,7 @@ func NewController(db *sqlx.DB) (
 }
 
 type AccountsListResponse struct {
+	// TODO(Raj): Update shape of account in API response
 	Accounts []Account `json:"accounts"`
 }
 
@@ -192,3 +193,27 @@ func (c *Controller) CheckInAsAgent(
 		Account: *account,
 	}, nil
 }
+
+type UpdateAccountConfigRequest struct {
+	Config AccountConfig `json:"config"`
+}
+
+func (c *Controller) UpdateAccountConfig(
+	ctx context.Context,
+	cloudProvider string,
+	accountId string,
+	req UpdateAccountConfigRequest,
+) (*Account, *model.ApiError) {
+	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
+		return nil, apiErr
+	}
+
+	account, apiErr := c.repo.upsert(
+		ctx, &accountId, &req.Config, nil, nil, nil,
+	)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
+	}
+
+	return account, nil
+}
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index 72c3392ef33..5a7d3f11ac2 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3880,6 +3880,10 @@ func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *Au
 		"/{cloudProvider}/accounts/{accountId}/status", am.ViewAccess(aH.CloudIntegrationsGetAccountStatus),
 	).Methods(http.MethodGet)
 
+	subRouter.HandleFunc(
+		"/{cloudProvider}/accounts/{accountId}/config", am.EditAccess(aH.CloudIntegrationsUpdateAccountConfig),
+	).Methods(http.MethodPost)
+
 	subRouter.HandleFunc(
 		"/{cloudProvider}/agent-check-in", am.EditAccess(aH.CloudIntegrationsAgentCheckIn),
 	).Methods(http.MethodPost)
@@ -3965,6 +3969,30 @@ func (aH *APIHandler) CloudIntegrationsAgentCheckIn(
 	aH.Respond(w, result)
 }
 
+func (aH *APIHandler) CloudIntegrationsUpdateAccountConfig(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+	accountId := mux.Vars(r)["accountId"]
+
+	req := cloudintegrations.UpdateAccountConfigRequest{}
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
+	}
+
+	result, apiErr := aH.CloudIntegrationsController.UpdateAccountConfig(
+		r.Context(), cloudProvider, accountId, req,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, nil)
+		return
+	}
+
+	aH.Respond(w, result)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 2552b36d10b..57999de0f95 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -95,11 +95,11 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	testAccountConfig2 := cloudintegrations.AccountConfig{
 		EnabledRegions: []string{"us-east-2", "us-west-1"},
 	}
-	latestAccount := testbed.UpdateAccountSettingsWithQS(
+	latestAccount := testbed.UpdateAccountConfigWithQS(
 		"aws", testAccountId, testAccountConfig2,
 	)
 	require.Equal(testAccountId, latestAccount.Id)
-	require.Equal(testAccountConfig2, latestAccount.Config)
+	require.Equal(testAccountConfig2, *latestAccount.Config)
 
 	// The agent should now receive latest settings.
 
@@ -244,14 +244,14 @@ func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
 	return &resp
 }
 
-func (tb *CloudIntegrationsTestBed) UpdateAccountSettingsWithQS(
+func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 	cloudProvider string, accountId string, newConfig cloudintegrations.AccountConfig,
 ) *cloudintegrations.Account {
 	result := tb.RequestQS(
 		fmt.Sprintf(
 			"/api/v1/cloud-integrations/%s/accounts/%s/config",
 			cloudProvider, accountId,
-		), cloudintegrations.UpdateAccountSettingsRequest{
+		), cloudintegrations.UpdateAccountConfigRequest{
 			Config: newConfig,
 		},
 	)

From e7e09bed56eaa9e37e203d29da2e38f7145e92ca Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 20:56:10 +0530
Subject: [PATCH 22/38] feat: cloud integrations: expectation for agent
 receiving latest config after account config update

---
 .../tests/integration/signoz_cloud_integrations_test.go  | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 57999de0f95..753fdc23628 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -102,6 +102,15 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Equal(testAccountConfig2, *latestAccount.Config)
 
 	// The agent should now receive latest settings.
+	agentCheckInResp1 := testbed.CheckInAsAgentWithQS(
+		"aws", cloudintegrations.AgentCheckInRequest{
+			AccountId:      testAccountId,
+			CloudAccountId: testAWSAccountId,
+		},
+	)
+	require.Equal(testAccountId, agentCheckInResp1.Account.Id)
+	require.Equal(testAccountConfig2, *agentCheckInResp1.Account.Config)
+	require.Equal(testAWSAccountId, *agentCheckInResp1.Account.CloudAccountId)
 
 	// Should be able to disconnect account.
 

From cece0f13a1bc169dd2620c68d10fca50f98e6fca Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 21:03:29 +0530
Subject: [PATCH 23/38] feat: cloud integrations: expectation for disconnecting
 cloud accounts from UI

---
 .../signoz_cloud_integrations_test.go         | 44 +++++++++++++++++--
 1 file changed, 41 insertions(+), 3 deletions(-)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 753fdc23628..da16b95c7e1 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -111,12 +111,26 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Equal(testAccountId, agentCheckInResp1.Account.Id)
 	require.Equal(testAccountConfig2, *agentCheckInResp1.Account.Config)
 	require.Equal(testAWSAccountId, *agentCheckInResp1.Account.CloudAccountId)
+	require.Nil(agentCheckInResp1.Account.RemovedAt)
 
 	// Should be able to disconnect account.
+	tsBeforeDisconnect := time.Now().Unix()
+	latestAccount = testbed.DisconnectAccountWithQS(
+		"aws", testAccountId,
+	)
+	require.Equal(testAccountId, latestAccount.Id)
+	require.LessOrEqual(tsBeforeDisconnect, *latestAccount.RemovedAt)
 
 	// The agent should receive the disconnected status post disconnection
-
-	require.Equal(1, 2)
+	agentCheckInResp2 := testbed.CheckInAsAgentWithQS(
+		"aws", cloudintegrations.AgentCheckInRequest{
+			AccountId:      testAccountId,
+			CloudAccountId: testAWSAccountId,
+		},
+	)
+	require.Equal(testAccountId, agentCheckInResp2.Account.Id)
+	require.Equal(testAWSAccountId, *agentCheckInResp2.Account.CloudAccountId)
+	require.LessOrEqual(tsBeforeDisconnect, *agentCheckInResp1.Account.RemovedAt)
 }
 
 type CloudIntegrationsTestBed struct {
@@ -273,7 +287,31 @@ func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 	var resp cloudintegrations.Account
 	err = json.Unmarshal(dataJson, &resp)
 	if err != nil {
-		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AgentCheckInResponse")
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")
+	}
+
+	return &resp
+}
+
+func (tb *CloudIntegrationsTestBed) DisconnectAccountWithQS(
+	cloudProvider string, accountId string,
+) *cloudintegrations.Account {
+	result := tb.RequestQS(
+		fmt.Sprintf(
+			"/api/v1/cloud-integrations/%s/accounts/%s/disconnect",
+			cloudProvider, accountId,
+		), map[string]any{},
+	)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp cloudintegrations.Account
+	err = json.Unmarshal(dataJson, &resp)
+	if err != nil {
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")
 	}
 
 	return &resp

From 045e5cbec451a2dedf4d61ca276f490460234e6c Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Tue, 7 Jan 2025 21:08:19 +0530
Subject: [PATCH 24/38] feat: cloud integrations: API for disconnecting cloud
 accounts

---
 .../app/cloudintegrations/controller.go       | 18 +++++++++++++++
 pkg/query-service/app/http_handler.go         | 22 +++++++++++++++++++
 .../signoz_cloud_integrations_test.go         |  4 ++--
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index eb23f66a186..0dcc350bdca 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -217,3 +217,21 @@ func (c *Controller) UpdateAccountConfig(
 
 	return account, nil
 }
+
+func (c *Controller) DisconnectAccount(
+	ctx context.Context, cloudProvider string, accountId string,
+) (*Account, *model.ApiError) {
+	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
+		return nil, apiErr
+	}
+
+	tsNow := time.Now()
+	account, apiErr := c.repo.upsert(
+		ctx, &accountId, nil, nil, nil, &tsNow,
+	)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
+	}
+
+	return account, nil
+}
diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go
index 5a7d3f11ac2..2fcc35a9fb6 100644
--- a/pkg/query-service/app/http_handler.go
+++ b/pkg/query-service/app/http_handler.go
@@ -3884,6 +3884,10 @@ func (aH *APIHandler) RegisterCloudIntegrationsRoutes(router *mux.Router, am *Au
 		"/{cloudProvider}/accounts/{accountId}/config", am.EditAccess(aH.CloudIntegrationsUpdateAccountConfig),
 	).Methods(http.MethodPost)
 
+	subRouter.HandleFunc(
+		"/{cloudProvider}/accounts/{accountId}/disconnect", am.EditAccess(aH.CloudIntegrationsDisconnectAccount),
+	).Methods(http.MethodPost)
+
 	subRouter.HandleFunc(
 		"/{cloudProvider}/agent-check-in", am.EditAccess(aH.CloudIntegrationsAgentCheckIn),
 	).Methods(http.MethodPost)
@@ -3993,6 +3997,24 @@ func (aH *APIHandler) CloudIntegrationsUpdateAccountConfig(
 	aH.Respond(w, result)
 }
 
+func (aH *APIHandler) CloudIntegrationsDisconnectAccount(
+	w http.ResponseWriter, r *http.Request,
+) {
+	cloudProvider := mux.Vars(r)["cloudProvider"]
+	accountId := mux.Vars(r)["accountId"]
+
+	result, apiErr := aH.CloudIntegrationsController.DisconnectAccount(
+		r.Context(), cloudProvider, accountId,
+	)
+
+	if apiErr != nil {
+		RespondError(w, apiErr, nil)
+		return
+	}
+
+	aH.Respond(w, result)
+}
+
 // logs
 func (aH *APIHandler) RegisterLogsRoutes(router *mux.Router, am *AuthMiddleware) {
 	subRouter := router.PathPrefix("/api/v1/logs").Subrouter()
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index da16b95c7e1..af3d65d93b1 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -114,7 +114,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Nil(agentCheckInResp1.Account.RemovedAt)
 
 	// Should be able to disconnect account.
-	tsBeforeDisconnect := time.Now().Unix()
+	tsBeforeDisconnect := time.Now()
 	latestAccount = testbed.DisconnectAccountWithQS(
 		"aws", testAccountId,
 	)
@@ -130,7 +130,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	)
 	require.Equal(testAccountId, agentCheckInResp2.Account.Id)
 	require.Equal(testAWSAccountId, *agentCheckInResp2.Account.CloudAccountId)
-	require.LessOrEqual(tsBeforeDisconnect, *agentCheckInResp1.Account.RemovedAt)
+	require.LessOrEqual(tsBeforeDisconnect, *agentCheckInResp2.Account.RemovedAt)
 }
 
 type CloudIntegrationsTestBed struct {

From 8028a18d44773335dafd6e7a057b4e1b373218d0 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 11:36:02 +0530
Subject: [PATCH 25/38] feat: cloud integrations: some cleanup

---
 .../app/cloudintegrations/controller.go       | 19 ++---------
 .../app/cloudintegrations/model.go            | 34 +++++++++++++------
 2 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 0dcc350bdca..f270575dd36 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -134,23 +134,10 @@ func (c *Controller) GetAccountStatus(
 		return nil, apiErr
 	}
 
-	var lastAgentReportTsMillis *int64
+	resp := AccountStatusResponse{Id: account.Id}
+
 	if account.LastAgentReport != nil {
-		// TODO(Raj): Investigate why agent reports are scanned to empty values
-		// from NULL columns even when report is a pointer
-		lastReportTsMillis := account.LastAgentReport.TimestampMillis
-		if lastReportTsMillis > 0 {
-			lastAgentReportTsMillis = &lastReportTsMillis
-		}
-	}
-
-	resp := AccountStatusResponse{
-		Id: account.Id,
-		Status: AccountStatus{
-			Integration: AccountIntegrationStatus{
-				LastHeartbeatTsMillis: lastAgentReportTsMillis,
-			},
-		},
+		resp.Status.Integration.LastHeartbeatTsMillis = &account.LastAgentReport.TimestampMillis
 	}
 
 	return &resp, nil
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index bd22b99922e..8bef801b961 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -28,15 +28,21 @@ func DefaultAccountConfig() AccountConfig {
 }
 
 // For serializing from db
-func (c *AccountConfig) Scan(src interface{}) error {
-	if data, ok := src.([]byte); ok {
-		return json.Unmarshal(data, &c)
+func (c *AccountConfig) Scan(src any) error {
+	data, ok := src.([]byte)
+	if !ok {
+		return fmt.Errorf("tried to scan from %T instead of bytes", src)
 	}
-	return nil
+
+	return json.Unmarshal(data, &c)
 }
 
 // For serializing to db
 func (c *AccountConfig) Value() (driver.Value, error) {
+	if c == nil {
+		return nil, nil
+	}
+
 	serialized, err := json.Marshal(c)
 	if err != nil {
 		return nil, fmt.Errorf(
@@ -52,19 +58,25 @@ type AgentReport struct {
 }
 
 // For serializing from db
-func (c *AgentReport) Scan(src interface{}) error {
-	if data, ok := src.([]byte); ok {
-		return json.Unmarshal(data, &c)
+func (r *AgentReport) Scan(src any) error {
+	data, ok := src.([]byte)
+	if !ok {
+		return fmt.Errorf("tried to scan from %T instead of bytes", src)
 	}
-	return nil
+
+	return json.Unmarshal(data, &r)
 }
 
 // For serializing to db
-func (c *AgentReport) Value() (driver.Value, error) {
-	serialized, err := json.Marshal(c)
+func (r *AgentReport) Value() (driver.Value, error) {
+	if r == nil {
+		return nil, nil
+	}
+
+	serialized, err := json.Marshal(r)
 	if err != nil {
 		return nil, fmt.Errorf(
-			"couldn't serialize cloud account config to JSON: %w", err,
+			"couldn't serialize agent report to JSON: %w", err,
 		)
 	}
 	return serialized, nil

From c8ff9b1b1ddf3f3deb0f9c8bac1dbd5d78e78e0a Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 14:26:37 +0530
Subject: [PATCH 26/38] feat: cloud integrations: some more cleanup

---
 .../signoz_cloud_integrations_test.go         | 115 ++++++------------
 1 file changed, 39 insertions(+), 76 deletions(-)

diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index af3d65d93b1..5ee105195ed 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -20,18 +20,18 @@ import (
 )
 
 func TestAWSIntegrationLifecycle(t *testing.T) {
-	// Test for the happy path of connecting and managing AWS integration accounts
+	// Test for happy path of connecting and managing AWS integration accounts
 
 	t0 := time.Now()
 	require := require.New(t)
 	testbed := NewCloudIntegrationsTestBed(t, nil)
 
-	accountsListResp := testbed.GetAccountsListFromQS("aws")
+	accountsListResp := testbed.GetConnectedAccountsListFromQS("aws")
 	require.Equal(len(accountsListResp.Accounts), 0,
 		"No accounts should be connected at the beginning",
 	)
 
-	// Should be able to generate connection url - initializing an account
+	// Should be able to generate a connection url from UI - initializing an integration account
 	testAccountConfig := cloudintegrations.AccountConfig{
 		EnabledRegions: []string{"us-east-1", "us-east-2"},
 	}
@@ -47,19 +47,16 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	connectionUrl := connectionUrlResp.ConnectionUrl
 	require.NotEmpty(connectionUrl)
 
-	// Should be able to poll for account connection status
+	// Should be able to poll for account connection status from the UI
 	accountStatusResp := testbed.GetAccountStatusFromQS("aws", testAccountId)
 	require.Equal(testAccountId, accountStatusResp.Id)
 	require.Nil(accountStatusResp.Status.Integration.LastHeartbeatTsMillis)
 
-	// The unconnected account should not show up in accounts list yet
-	accountsListResp1 := testbed.GetAccountsListFromQS("aws")
-	require.Equal(0, len(accountsListResp1.Accounts),
-		"No accounts should be connected at the beginning",
-	)
+	// The unconnected account should not show up in connected accounts list yet
+	accountsListResp1 := testbed.GetConnectedAccountsListFromQS("aws")
+	require.Equal(0, len(accountsListResp1.Accounts))
 
-	// An agent should be able to check in to the new account
-	// Should get the settings that were specified while generating connection url
+	// An agent installed in user's AWS account should be able to check in for the new integration account
 	tsMillisBeforeAgentCheckIn := time.Now().UnixMilli()
 	testAWSAccountId := "4563215233"
 	agentCheckInResp := testbed.CheckInAsAgentWithQS(
@@ -74,7 +71,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.LessOrEqual(t0.Unix(), agentCheckInResp.Account.CreatedAt.Unix())
 	require.Nil(agentCheckInResp.Account.RemovedAt)
 
-	// Polling for connection status should now return latest status
+	// Polling for connection status from UI should now return latest status
 	accountStatusResp1 := testbed.GetAccountStatusFromQS("aws", testAccountId)
 	require.Equal(testAccountId, accountStatusResp1.Id)
 	require.NotNil(accountStatusResp1.Status.Integration.LastHeartbeatTsMillis)
@@ -84,14 +81,12 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	)
 
 	// The account should now show up in list of connected accounts.
-	accountsListResp2 := testbed.GetAccountsListFromQS("aws")
-	require.Equal(len(accountsListResp2.Accounts), 1,
-		"No accounts should be connected at the beginning",
-	)
+	accountsListResp2 := testbed.GetConnectedAccountsListFromQS("aws")
+	require.Equal(len(accountsListResp2.Accounts), 1)
 	require.Equal(testAccountId, accountsListResp2.Accounts[0].Id)
 	require.Equal(testAWSAccountId, *accountsListResp2.Accounts[0].CloudAccountId)
 
-	// Should be able to update account settings
+	// Should be able to update account config from UI
 	testAccountConfig2 := cloudintegrations.AccountConfig{
 		EnabledRegions: []string{"us-east-2", "us-west-1"},
 	}
@@ -101,7 +96,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Equal(testAccountId, latestAccount.Id)
 	require.Equal(testAccountConfig2, *latestAccount.Config)
 
-	// The agent should now receive latest settings.
+	// The agent should now receive latest account config.
 	agentCheckInResp1 := testbed.CheckInAsAgentWithQS(
 		"aws", cloudintegrations.AgentCheckInRequest{
 			AccountId:      testAccountId,
@@ -113,15 +108,13 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	require.Equal(testAWSAccountId, *agentCheckInResp1.Account.CloudAccountId)
 	require.Nil(agentCheckInResp1.Account.RemovedAt)
 
-	// Should be able to disconnect account.
+	// Should be able to disconnect/remove account from UI.
 	tsBeforeDisconnect := time.Now()
-	latestAccount = testbed.DisconnectAccountWithQS(
-		"aws", testAccountId,
-	)
+	latestAccount = testbed.DisconnectAccountWithQS("aws", testAccountId)
 	require.Equal(testAccountId, latestAccount.Id)
 	require.LessOrEqual(tsBeforeDisconnect, *latestAccount.RemovedAt)
 
-	// The agent should receive the disconnected status post disconnection
+	// The agent should receive the disconnected status in account config post disconnection
 	agentCheckInResp2 := testbed.CheckInAsAgentWithQS(
 		"aws", cloudintegrations.AgentCheckInRequest{
 			AccountId:      testAccountId,
@@ -152,11 +145,7 @@ func NewCloudIntegrationsTestBed(t *testing.T, testDB *sqlx.DB) *CloudIntegratio
 	}
 
 	fm := featureManager.StartManager()
-	reader, mockClickhouse := NewMockClickhouseReader(t, testDB, fm)
-	mockClickhouse.MatchExpectationsInOrder(false)
-
 	apiHandler, err := app.NewAPIHandler(app.APIHandlerOpts{
-		Reader:                      reader,
 		AppDao:                      dao.DB(),
 		CloudIntegrationsController: controller,
 		FeatureFlags:                fm,
@@ -176,25 +165,19 @@ func NewCloudIntegrationsTestBed(t *testing.T, testDB *sqlx.DB) *CloudIntegratio
 	}
 
 	return &CloudIntegrationsTestBed{
-		t:              t,
-		testUser:       user,
-		qsHttpHandler:  router,
-		mockClickhouse: mockClickhouse,
+		t:             t,
+		testUser:      user,
+		qsHttpHandler: router,
 	}
 }
 
-func (tb *CloudIntegrationsTestBed) GetAccountsListFromQS(
+func (tb *CloudIntegrationsTestBed) GetConnectedAccountsListFromQS(
 	cloudProvider string,
 ) *cloudintegrations.AccountsListResponse {
-	result := tb.RequestQS(fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts", cloudProvider), nil)
-
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
+	respDataJson := tb.RequestQS(fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts", cloudProvider), nil)
 
 	var resp cloudintegrations.AccountsListResponse
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AccountsListResponse")
 	}
@@ -205,18 +188,13 @@ func (tb *CloudIntegrationsTestBed) GetAccountsListFromQS(
 func (tb *CloudIntegrationsTestBed) GenerateConnectionUrlFromQS(
 	cloudProvider string, req cloudintegrations.GenerateConnectionUrlRequest,
 ) *cloudintegrations.GenerateConnectionUrlResponse {
-	result := tb.RequestQS(
+	respDataJson := tb.RequestQS(
 		fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts/generate-connection-url", cloudProvider),
 		req,
 	)
 
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
-
 	var resp cloudintegrations.GenerateConnectionUrlResponse
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into map[string]any")
 	}
@@ -227,18 +205,13 @@ func (tb *CloudIntegrationsTestBed) GenerateConnectionUrlFromQS(
 func (tb *CloudIntegrationsTestBed) GetAccountStatusFromQS(
 	cloudProvider string, accountId string,
 ) *cloudintegrations.AccountStatusResponse {
-	result := tb.RequestQS(fmt.Sprintf(
+	respDataJson := tb.RequestQS(fmt.Sprintf(
 		"/api/v1/cloud-integrations/%s/accounts/%s/status",
 		cloudProvider, accountId,
 	), nil)
 
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
-
 	var resp cloudintegrations.AccountStatusResponse
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AccountStatusResponse")
 	}
@@ -249,17 +222,12 @@ func (tb *CloudIntegrationsTestBed) GetAccountStatusFromQS(
 func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
 	cloudProvider string, req cloudintegrations.AgentCheckInRequest,
 ) *cloudintegrations.AgentCheckInResponse {
-	result := tb.RequestQS(
+	respDataJson := tb.RequestQS(
 		fmt.Sprintf("/api/v1/cloud-integrations/%s/agent-check-in", cloudProvider), req,
 	)
 
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
-
 	var resp cloudintegrations.AgentCheckInResponse
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AgentCheckInResponse")
 	}
@@ -270,7 +238,7 @@ func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
 func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 	cloudProvider string, accountId string, newConfig cloudintegrations.AccountConfig,
 ) *cloudintegrations.Account {
-	result := tb.RequestQS(
+	respDataJson := tb.RequestQS(
 		fmt.Sprintf(
 			"/api/v1/cloud-integrations/%s/accounts/%s/config",
 			cloudProvider, accountId,
@@ -279,13 +247,8 @@ func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 		},
 	)
 
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
-
 	var resp cloudintegrations.Account
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")
 	}
@@ -296,20 +259,15 @@ func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 func (tb *CloudIntegrationsTestBed) DisconnectAccountWithQS(
 	cloudProvider string, accountId string,
 ) *cloudintegrations.Account {
-	result := tb.RequestQS(
+	respDataJson := tb.RequestQS(
 		fmt.Sprintf(
 			"/api/v1/cloud-integrations/%s/accounts/%s/disconnect",
 			cloudProvider, accountId,
 		), map[string]any{},
 	)
 
-	dataJson, err := json.Marshal(result.Data)
-	if err != nil {
-		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
-	}
-
 	var resp cloudintegrations.Account
-	err = json.Unmarshal(dataJson, &resp)
+	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")
 	}
@@ -320,7 +278,7 @@ func (tb *CloudIntegrationsTestBed) DisconnectAccountWithQS(
 func (tb *CloudIntegrationsTestBed) RequestQS(
 	path string,
 	postData interface{},
-) *app.ApiResponse {
+) (responseDataJson []byte) {
 	req, err := AuthenticatedRequestForTest(
 		tb.testUser, path, postData,
 	)
@@ -332,5 +290,10 @@ func (tb *CloudIntegrationsTestBed) RequestQS(
 	if err != nil {
 		tb.t.Fatalf("test request failed: %v", err)
 	}
-	return result
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+	return dataJson
 }

From b883e98c82d50e2721fd6a7f7d586d622c83af07 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 15:49:11 +0530
Subject: [PATCH 27/38] feat: cloud integrations: repo: scope rows by cloud
 provider

---
 .../app/cloudintegrations/controller.go       | 12 ++---
 .../app/cloudintegrations/model.go            |  1 +
 .../app/cloudintegrations/repo.go             | 54 +++++++++++--------
 3 files changed, 39 insertions(+), 28 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index f270575dd36..644d2de1a2f 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -52,7 +52,7 @@ func (c *Controller) ListConnectedAccounts(
 		return nil, apiErr
 	}
 
-	accounts, apiErr := c.repo.listConnected(ctx)
+	accounts, apiErr := c.repo.listConnected(ctx, cloudProvider)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't list cloud accounts")
 	}
@@ -89,7 +89,7 @@ func (c *Controller) GenerateConnectionUrl(
 	}
 
 	account, apiErr := c.repo.upsert(
-		ctx, req.AccountId, &req.AccountConfig, nil, nil, nil,
+		ctx, cloudProvider, req.AccountId, &req.AccountConfig, nil, nil, nil,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
@@ -129,7 +129,7 @@ func (c *Controller) GetAccountStatus(
 		return nil, apiErr
 	}
 
-	account, apiErr := c.repo.get(ctx, accountId)
+	account, apiErr := c.repo.get(ctx, cloudProvider, accountId)
 	if apiErr != nil {
 		return nil, apiErr
 	}
@@ -170,7 +170,7 @@ func (c *Controller) CheckInAsAgent(
 	// TODO(Raj): Consider ensuring cloudAccountId is not updated post insertion
 	// Consider getting the account by Id first to validate
 	account, apiErr := c.repo.upsert(
-		ctx, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
+		ctx, cloudProvider, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
@@ -196,7 +196,7 @@ func (c *Controller) UpdateAccountConfig(
 	}
 
 	account, apiErr := c.repo.upsert(
-		ctx, &accountId, &req.Config, nil, nil, nil,
+		ctx, cloudProvider, &accountId, &req.Config, nil, nil, nil,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
@@ -214,7 +214,7 @@ func (c *Controller) DisconnectAccount(
 
 	tsNow := time.Now()
 	account, apiErr := c.repo.upsert(
-		ctx, &accountId, nil, nil, nil, &tsNow,
+		ctx, cloudProvider, &accountId, nil, nil, nil, &tsNow,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 8bef801b961..bc69ec27c00 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -9,6 +9,7 @@ import (
 
 // Represents a cloud provider account for cloud integrations
 type Account struct {
+	CloudProvider   string         `json:"cloud_provider" db:"cloud_provider"`
 	Id              string         `json:"id" db:"id"`
 	Config          *AccountConfig `json:"config_json" db:"config_json"`
 	CloudAccountId  *string        `json:"cloud_account_id" db:"cloud_account_id"`
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 46980ce642c..f616802b042 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -12,16 +12,16 @@ import (
 )
 
 type cloudProviderAccountsRepository interface {
-	// TODO(Raj): All methods should be scoped by cloud provider.
-	listConnected(context.Context) ([]Account, *model.ApiError)
+	listConnected(ctx context.Context, cloudProvider string) ([]Account, *model.ApiError)
 
-	getByIds(ctx context.Context, ids []string) (map[string]*Account, *model.ApiError)
+	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*Account, *model.ApiError)
 
-	get(ctx context.Context, id string) (*Account, *model.ApiError)
+	get(ctx context.Context, cloudProvider string, id string) (*Account, *model.ApiError)
 
 	// Insert an account or update it by ID for specified non-empty fields
 	upsert(
 		ctx context.Context,
+		cloudProvider string,
 		id *string,
 		config *AccountConfig,
 		cloudAccountId *string,
@@ -49,12 +49,14 @@ func InitSqliteDBIfNeeded(db *sqlx.DB) error {
 
 	createTablesStatements := `
 		CREATE TABLE IF NOT EXISTS cloud_integrations_accounts(
-			id TEXT PRIMARY KEY NOT NULL,
+			cloud_provider TEXT NOT NULL,
+			id TEXT NOT NULL,
 			config_json TEXT,
 			cloud_account_id TEXT,
 			last_agent_report_json TEXT,
 			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
-			removed_at TIMESTAMP
+			removed_at TIMESTAMP,
+			UNIQUE(cloud_provider, id)
 		)
 	`
 	_, err := db.Exec(createTablesStatements)
@@ -71,14 +73,15 @@ type cloudProviderAccountsSQLRepository struct {
 	db *sqlx.DB
 }
 
-func (r *cloudProviderAccountsSQLRepository) listConnected(ctx context.Context) (
-	[]Account, *model.ApiError,
-) {
+func (r *cloudProviderAccountsSQLRepository) listConnected(
+	ctx context.Context, cloudProvider string,
+) ([]Account, *model.ApiError) {
 	accounts := []Account{}
 
 	err := r.db.SelectContext(
 		ctx, &accounts, `
 			select
+				cloud_provider,
 				id,
 				config_json,
 				cloud_account_id,
@@ -87,10 +90,11 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(ctx context.Context)
 				removed_at
 			from cloud_integrations_accounts
       where
-        removed_at is NULL
+        cloud_provider=$1
+        and removed_at is NULL
         and cloud_account_id is not NULL
 			order by created_at
-		`,
+		`, cloudProvider,
 	)
 	if err != nil {
 		return nil, model.InternalError(fmt.Errorf(
@@ -102,20 +106,22 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(ctx context.Context)
 }
 
 func (r *cloudProviderAccountsSQLRepository) getByIds(
-	ctx context.Context, ids []string,
+	ctx context.Context, cloudProvider string, ids []string,
 ) (map[string]*Account, *model.ApiError) {
 	accounts := []Account{}
 
 	idPlaceholders := []string{}
-	idValues := []interface{}{}
+	queryArgs := []any{cloudProvider}
+
 	for _, id := range ids {
 		idPlaceholders = append(idPlaceholders, "?")
-		idValues = append(idValues, id)
+		queryArgs = append(queryArgs, id)
 	}
 
 	err := r.db.SelectContext(
 		ctx, &accounts, fmt.Sprintf(`
 			select
+				cloud_provider,
 				id,
 				config_json,
 				cloud_account_id,
@@ -123,10 +129,12 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 				created_at,
 				removed_at
 			from cloud_integrations_accounts
-			where id in (%s)`,
+			where
+				cloud_provider=?
+				and id in (%s)`,
 			strings.Join(idPlaceholders, ", "),
 		),
-		idValues...,
+		queryArgs...,
 	)
 	if err != nil {
 		return nil, model.InternalError(fmt.Errorf(
@@ -149,9 +157,9 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 }
 
 func (r *cloudProviderAccountsSQLRepository) get(
-	ctx context.Context, id string,
+	ctx context.Context, cloudProvider string, id string,
 ) (*Account, *model.ApiError) {
-	res, apiErr := r.getByIds(ctx, []string{id})
+	res, apiErr := r.getByIds(ctx, cloudProvider, []string{id})
 	if apiErr != nil {
 		return nil, apiErr
 	}
@@ -168,6 +176,7 @@ func (r *cloudProviderAccountsSQLRepository) get(
 
 func (r *cloudProviderAccountsSQLRepository) upsert(
 	ctx context.Context,
+	cloudProvider string,
 	id *string,
 	config *AccountConfig,
 	cloudAccountId *string,
@@ -213,25 +222,26 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 	onConflictClause := ""
 	if len(onConflictSetStmts) > 0 {
 		onConflictClause = fmt.Sprintf(
-			"on conflict(id) do update SET\n%s",
+			"on conflict(cloud_provider, id) do update SET\n%s",
 			strings.Join(onConflictSetStmts, ",\n"),
 		)
 	}
 
 	insertQuery := fmt.Sprintf(`
     INSERT INTO cloud_integrations_accounts (
+      cloud_provider,
       id,
       config_json,
       cloud_account_id,
       last_agent_report_json,
       removed_at
-    ) values ($1, $2, $3, $4, $5)
+    ) values ($1, $2, $3, $4, $5, $6)
     %s`, onConflictClause,
 	)
 
 	_, dbErr := r.db.ExecContext(
 		ctx, insertQuery,
-		id, config, cloudAccountId, agentReport, removedAt,
+		cloudProvider, id, config, cloudAccountId, agentReport, removedAt,
 	)
 	if dbErr != nil {
 		return nil, model.InternalError(fmt.Errorf(
@@ -239,7 +249,7 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 		))
 	}
 
-	upsertedAccount, apiErr := r.get(ctx, *id)
+	upsertedAccount, apiErr := r.get(ctx, cloudProvider, *id)
 	if apiErr != nil {
 		return nil, model.InternalError(fmt.Errorf(
 			"couldn't fetch upserted account by id: %w", apiErr.ToError(),

From c9f4397c9964b22f9a2aef7ffeb92e505d0392e9 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 18:10:36 +0530
Subject: [PATCH 28/38] feat: testutils: refactor out helper for creating a
 test sqlite DB

---
 pkg/query-service/utils/testutils.go | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/pkg/query-service/utils/testutils.go b/pkg/query-service/utils/testutils.go
index d8989d93239..69c7a9bf24d 100644
--- a/pkg/query-service/utils/testutils.go
+++ b/pkg/query-service/utils/testutils.go
@@ -5,24 +5,31 @@ import (
 	"testing"
 
 	"github.com/jmoiron/sqlx"
+	_ "github.com/mattn/go-sqlite3"
 	"go.signoz.io/signoz/pkg/query-service/app/dashboards"
 	"go.signoz.io/signoz/pkg/query-service/dao"
 )
 
-func NewQueryServiceDBForTests(t *testing.T) *sqlx.DB {
+func NewTestSqliteDB(t *testing.T) (testDB *sqlx.DB, testDBFilePath string) {
 	testDBFile, err := os.CreateTemp("", "test-signoz-db-*")
 	if err != nil {
 		t.Fatalf("could not create temp file for test db: %v", err)
 	}
-	testDBFilePath := testDBFile.Name()
+	testDBFilePath = testDBFile.Name()
 	t.Cleanup(func() { os.Remove(testDBFilePath) })
 	testDBFile.Close()
 
-	testDB, err := sqlx.Open("sqlite3", testDBFilePath)
+	testDB, err = sqlx.Open("sqlite3", testDBFilePath)
 	if err != nil {
 		t.Fatalf("could not open test db sqlite file: %v", err)
 	}
 
+	return testDB, testDBFilePath
+}
+
+func NewQueryServiceDBForTests(t *testing.T) *sqlx.DB {
+	testDB, testDBFilePath := NewTestSqliteDB(t)
+
 	// TODO(Raj): This should not require passing in the DB file path
 	dao.InitDao("sqlite", testDBFilePath)
 	dashboards.InitDB(testDBFilePath)

From 34d480a95814a1d157614946e489e747e79e24c8 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 18:13:22 +0530
Subject: [PATCH 29/38] feat: cloud integrations: controller: add test
 validating regeneration of connection url

---
 .../app/cloudintegrations/controller_test.go  | 53 +++++++++++++++++++
 .../app/cloudintegrations/repo.go             |  3 ++
 .../app/cloudintegrations/repo_test.go        | 17 ++++++
 3 files changed, 73 insertions(+)
 create mode 100644 pkg/query-service/app/cloudintegrations/controller_test.go
 create mode 100644 pkg/query-service/app/cloudintegrations/repo_test.go

diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
new file mode 100644
index 00000000000..5dd6cd626c6
--- /dev/null
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -0,0 +1,53 @@
+package cloudintegrations
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/pkg/query-service/utils"
+)
+
+func TestRegenerateConnectionUrlWithUpdatedConfig(t *testing.T) {
+	require := require.New(t)
+	testDB, _ := utils.NewTestSqliteDB(t)
+	controller, err := NewController(testDB)
+	require.NoError(err)
+
+	// should be able to generate connection url for
+	// same account id again with updated config
+	testAccountConfig1 := AccountConfig{EnabledRegions: []string{"us-east-1", "us-west-1"}}
+	resp1, apiErr := controller.GenerateConnectionUrl(
+		context.TODO(), "aws", GenerateConnectionUrlRequest{
+			AccountConfig: testAccountConfig1,
+			AgentConfig:   SigNozAgentConfig{Region: "us-east-2"},
+		},
+	)
+	require.Nil(apiErr)
+	require.NotEmpty(resp1.ConnectionUrl)
+	require.NotEmpty(resp1.AccountId)
+
+	testAccountId := resp1.AccountId
+	account, apiErr := controller.repo.get(
+		context.TODO(), "aws", testAccountId,
+	)
+	require.Nil(apiErr)
+	require.Equal(testAccountConfig1, *account.Config)
+
+	testAccountConfig2 := AccountConfig{EnabledRegions: []string{"us-east-2", "us-west-2"}}
+	resp2, apiErr := controller.GenerateConnectionUrl(
+		context.TODO(), "aws", GenerateConnectionUrlRequest{
+			AccountId:     &testAccountId,
+			AccountConfig: testAccountConfig2,
+			AgentConfig:   SigNozAgentConfig{Region: "us-east-2"},
+		},
+	)
+	require.Nil(apiErr)
+	require.Equal(testAccountId, resp2.AccountId)
+
+	account, apiErr = controller.repo.get(
+		context.TODO(), "aws", testAccountId,
+	)
+	require.Nil(apiErr)
+	require.Equal(testAccountConfig2, *account.Config)
+}
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index f616802b042..e62490503cb 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -12,10 +12,12 @@ import (
 )
 
 type cloudProviderAccountsRepository interface {
+	// list connected cloud provider accounts.
 	listConnected(ctx context.Context, cloudProvider string) ([]Account, *model.ApiError)
 
 	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*Account, *model.ApiError)
 
+	// TODO(Raj): Move this to being a helper on the controller?
 	get(ctx context.Context, cloudProvider string, id string) (*Account, *model.ApiError)
 
 	// Insert an account or update it by ID for specified non-empty fields
@@ -93,6 +95,7 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(
         cloud_provider=$1
         and removed_at is NULL
         and cloud_account_id is not NULL
+        and last_agent_report_json is not NULL
 			order by created_at
 		`, cloudProvider,
 	)
diff --git a/pkg/query-service/app/cloudintegrations/repo_test.go b/pkg/query-service/app/cloudintegrations/repo_test.go
new file mode 100644
index 00000000000..fbf96be4c0d
--- /dev/null
+++ b/pkg/query-service/app/cloudintegrations/repo_test.go
@@ -0,0 +1,17 @@
+package cloudintegrations
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRepo(t *testing.T) {
+	require := require.New(t)
+
+	// record Ids should be scoped by cloud provider.
+
+	// querying records should be scoped by cloud provider.
+
+	require.Equal(1, 2)
+}

From 85ce9801b7ba50628adb337b68f902052aed0a7b Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 18:27:35 +0530
Subject: [PATCH 30/38] feat: cloud integrations: controller: validations for
 agent check ins

---
 .../app/cloudintegrations/controller.go       | 13 +++++---
 .../app/cloudintegrations/controller_test.go  | 33 +++++++++++++++++++
 .../app/cloudintegrations/repo.go             |  1 -
 .../app/cloudintegrations/repo_test.go        | 17 ----------
 4 files changed, 42 insertions(+), 22 deletions(-)
 delete mode 100644 pkg/query-service/app/cloudintegrations/repo_test.go

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 644d2de1a2f..bd54ee30ad5 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -157,19 +157,24 @@ type AgentCheckInResponse struct {
 func (c *Controller) CheckInAsAgent(
 	ctx context.Context, cloudProvider string, req AgentCheckInRequest,
 ) (*AgentCheckInResponse, *model.ApiError) {
-	// Account connection with a simple connection URL may not be available for all providers.
 	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
 		return nil, apiErr
 	}
 
+	account, apiErr := c.repo.get(ctx, cloudProvider, req.AccountId)
+	if account != nil && account.CloudAccountId != nil && *account.CloudAccountId != req.CloudAccountId {
+		return nil, model.BadRequest(fmt.Errorf(
+			"can't check in with new %s account id %s for account %s with existing %s id %s",
+			cloudProvider, req.CloudAccountId, account.Id, cloudProvider, *account.CloudAccountId,
+		))
+	}
+
 	agentReport := AgentReport{
 		TimestampMillis: time.Now().UnixMilli(),
 		Data:            req.Data,
 	}
 
-	// TODO(Raj): Consider ensuring cloudAccountId is not updated post insertion
-	// Consider getting the account by Id first to validate
-	account, apiErr := c.repo.upsert(
+	account, apiErr = c.repo.upsert(
 		ctx, cloudProvider, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
 	)
 	if apiErr != nil {
diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
index 5dd6cd626c6..a34d57a0c13 100644
--- a/pkg/query-service/app/cloudintegrations/controller_test.go
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"go.signoz.io/signoz/pkg/query-service/utils"
 )
@@ -51,3 +52,35 @@ func TestRegenerateConnectionUrlWithUpdatedConfig(t *testing.T) {
 	require.Nil(apiErr)
 	require.Equal(testAccountConfig2, *account.Config)
 }
+
+func TestAgentCheckIns(t *testing.T) {
+	require := require.New(t)
+	testDB, _ := utils.NewTestSqliteDB(t)
+	controller, err := NewController(testDB)
+	require.NoError(err)
+
+	// An agent should be able to check in from a cloud account even
+	// if no connection url was requested (no account with agent's account id exists)
+	testAccountId1 := uuid.NewString()
+	testCloudAccountId1 := "546311234"
+	resp1, apiErr := controller.CheckInAsAgent(
+		context.TODO(), "aws", AgentCheckInRequest{
+			AccountId:      testAccountId1,
+			CloudAccountId: testCloudAccountId1,
+		},
+	)
+	require.Nil(apiErr)
+	require.Equal(testAccountId1, resp1.Account.Id)
+	require.Equal(testCloudAccountId1, *resp1.Account.CloudAccountId)
+
+	// The agent should not be able to check in with a different
+	// cloud account id for the same account.
+	testCloudAccountId2 := "99999999"
+	_, apiErr = controller.CheckInAsAgent(
+		context.TODO(), "aws", AgentCheckInRequest{
+			AccountId:      testAccountId1,
+			CloudAccountId: testCloudAccountId2,
+		},
+	)
+	require.NotNil(apiErr)
+}
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index e62490503cb..9e44246f2e4 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -17,7 +17,6 @@ type cloudProviderAccountsRepository interface {
 
 	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*Account, *model.ApiError)
 
-	// TODO(Raj): Move this to being a helper on the controller?
 	get(ctx context.Context, cloudProvider string, id string) (*Account, *model.ApiError)
 
 	// Insert an account or update it by ID for specified non-empty fields
diff --git a/pkg/query-service/app/cloudintegrations/repo_test.go b/pkg/query-service/app/cloudintegrations/repo_test.go
deleted file mode 100644
index fbf96be4c0d..00000000000
--- a/pkg/query-service/app/cloudintegrations/repo_test.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package cloudintegrations
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestRepo(t *testing.T) {
-	require := require.New(t)
-
-	// record Ids should be scoped by cloud provider.
-
-	// querying records should be scoped by cloud provider.
-
-	require.Equal(1, 2)
-}

From d26e388a54eb55533f2d3e59cbe5472337e37e8e Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 20:35:09 +0530
Subject: [PATCH 31/38] feat: cloud integrations: connected account response
 structure

---
 .../app/cloudintegrations/controller.go       | 44 +++++++++++--------
 .../app/cloudintegrations/model.go            | 39 +++++++++++++++-
 .../signoz_cloud_integrations_test.go         |  6 +--
 3 files changed, 66 insertions(+), 23 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index bd54ee30ad5..5b506548e87 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/jmoiron/sqlx"
 	"go.signoz.io/signoz/pkg/query-service/model"
+	"go.uber.org/zap"
 )
 
 var SupportedCloudProviders = []string{
@@ -38,15 +39,21 @@ func NewController(db *sqlx.DB) (
 	}, nil
 }
 
-type AccountsListResponse struct {
-	// TODO(Raj): Update shape of account in API response
-	Accounts []Account `json:"accounts"`
+type ConnectedAccount struct {
+	Id             string        `json:"id"`
+	CloudAccountId string        `json:"cloud_account_id"`
+	Config         AccountConfig `json:"config"`
+	Status         AccountStatus `json:"status"`
+}
+
+type ConnectedAccountsListResponse struct {
+	Accounts []ConnectedAccount `json:"accounts"`
 }
 
 func (c *Controller) ListConnectedAccounts(
 	ctx context.Context, cloudProvider string,
 ) (
-	*AccountsListResponse, *model.ApiError,
+	*ConnectedAccountsListResponse, *model.ApiError,
 ) {
 	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
 		return nil, apiErr
@@ -57,8 +64,18 @@ func (c *Controller) ListConnectedAccounts(
 		return nil, model.WrapApiError(apiErr, "couldn't list cloud accounts")
 	}
 
-	return &AccountsListResponse{
-		Accounts: accounts,
+	connectedAccounts := []ConnectedAccount{}
+	for _, a := range accounts {
+		ca, err := a.connectedAccount()
+		if err != nil {
+			zap.L().Error("invalid connected account record", zap.String("accountId", a.Id), zap.Error(err))
+		} else {
+			connectedAccounts = append(connectedAccounts, *ca)
+		}
+	}
+
+	return &ConnectedAccountsListResponse{
+		Accounts: connectedAccounts,
 	}, nil
 }
 
@@ -112,14 +129,6 @@ type AccountStatusResponse struct {
 	Status AccountStatus `json:"status"`
 }
 
-type AccountStatus struct {
-	Integration AccountIntegrationStatus `json:"integration"`
-}
-
-type AccountIntegrationStatus struct {
-	LastHeartbeatTsMillis *int64 `json:"last_heartbeat_ts_ms"`
-}
-
 func (c *Controller) GetAccountStatus(
 	ctx context.Context, cloudProvider string, accountId string,
 ) (
@@ -134,10 +143,9 @@ func (c *Controller) GetAccountStatus(
 		return nil, apiErr
 	}
 
-	resp := AccountStatusResponse{Id: account.Id}
-
-	if account.LastAgentReport != nil {
-		resp.Status.Integration.LastHeartbeatTsMillis = &account.LastAgentReport.TimestampMillis
+	resp := AccountStatusResponse{
+		Id:     account.Id,
+		Status: account.status(),
 	}
 
 	return &resp, nil
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index bc69ec27c00..122be73a1d0 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -11,9 +11,9 @@ import (
 type Account struct {
 	CloudProvider   string         `json:"cloud_provider" db:"cloud_provider"`
 	Id              string         `json:"id" db:"id"`
-	Config          *AccountConfig `json:"config_json" db:"config_json"`
+	Config          *AccountConfig `json:"config" db:"config_json"`
 	CloudAccountId  *string        `json:"cloud_account_id" db:"cloud_account_id"`
-	LastAgentReport *AgentReport   `json:"last_agent_report_json" db:"last_agent_report_json"`
+	LastAgentReport *AgentReport   `json:"last_agent_report" db:"last_agent_report_json"`
 	CreatedAt       time.Time      `json:"created_at" db:"created_at"`
 	RemovedAt       *time.Time     `json:"removed_at" db:"removed_at"`
 }
@@ -82,3 +82,38 @@ func (r *AgentReport) Value() (driver.Value, error) {
 	}
 	return serialized, nil
 }
+
+type AccountStatus struct {
+	Integration AccountIntegrationStatus `json:"integration"`
+}
+
+type AccountIntegrationStatus struct {
+	LastHeartbeatTsMillis *int64 `json:"last_heartbeat_ts_ms"`
+}
+
+func (a *Account) status() AccountStatus {
+	status := AccountStatus{}
+	if a.LastAgentReport != nil {
+		lastHeartbeat := a.LastAgentReport.TimestampMillis
+		status.Integration.LastHeartbeatTsMillis = &lastHeartbeat
+	}
+	return status
+}
+
+func (a *Account) connectedAccount() (*ConnectedAccount, error) {
+	ca := ConnectedAccount{Id: a.Id, Status: a.status()}
+
+	if a.CloudAccountId == nil {
+		return nil, fmt.Errorf("account %s has nil cloud account id", a.Id)
+	} else {
+		ca.CloudAccountId = *a.CloudAccountId
+	}
+
+	if a.Config == nil {
+		return nil, fmt.Errorf("account %s has nil config", a.Id)
+	} else {
+		ca.Config = *a.Config
+	}
+
+	return &ca, nil
+}
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index 5ee105195ed..ee2f9292484 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -84,7 +84,7 @@ func TestAWSIntegrationLifecycle(t *testing.T) {
 	accountsListResp2 := testbed.GetConnectedAccountsListFromQS("aws")
 	require.Equal(len(accountsListResp2.Accounts), 1)
 	require.Equal(testAccountId, accountsListResp2.Accounts[0].Id)
-	require.Equal(testAWSAccountId, *accountsListResp2.Accounts[0].CloudAccountId)
+	require.Equal(testAWSAccountId, accountsListResp2.Accounts[0].CloudAccountId)
 
 	// Should be able to update account config from UI
 	testAccountConfig2 := cloudintegrations.AccountConfig{
@@ -173,10 +173,10 @@ func NewCloudIntegrationsTestBed(t *testing.T, testDB *sqlx.DB) *CloudIntegratio
 
 func (tb *CloudIntegrationsTestBed) GetConnectedAccountsListFromQS(
 	cloudProvider string,
-) *cloudintegrations.AccountsListResponse {
+) *cloudintegrations.ConnectedAccountsListResponse {
 	respDataJson := tb.RequestQS(fmt.Sprintf("/api/v1/cloud-integrations/%s/accounts", cloudProvider), nil)
 
-	var resp cloudintegrations.AccountsListResponse
+	var resp cloudintegrations.ConnectedAccountsListResponse
 	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into AccountsListResponse")

From 12cc83f223f51f7d28df9ed23942893613ff0a9e Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Wed, 8 Jan 2025 20:40:11 +0530
Subject: [PATCH 32/38] feat: cloud integrations: API response account
 structure

---
 .../app/cloudintegrations/controller.go       | 28 ++++++++-----------
 .../app/cloudintegrations/model.go            | 18 +++++-------
 .../app/cloudintegrations/repo.go             | 22 +++++++--------
 .../signoz_cloud_integrations_test.go         |  8 +++---
 4 files changed, 34 insertions(+), 42 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index 5b506548e87..f502282817f 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/jmoiron/sqlx"
 	"go.signoz.io/signoz/pkg/query-service/model"
-	"go.uber.org/zap"
 )
 
 var SupportedCloudProviders = []string{
@@ -39,7 +38,7 @@ func NewController(db *sqlx.DB) (
 	}, nil
 }
 
-type ConnectedAccount struct {
+type Account struct {
 	Id             string        `json:"id"`
 	CloudAccountId string        `json:"cloud_account_id"`
 	Config         AccountConfig `json:"config"`
@@ -47,7 +46,7 @@ type ConnectedAccount struct {
 }
 
 type ConnectedAccountsListResponse struct {
-	Accounts []ConnectedAccount `json:"accounts"`
+	Accounts []Account `json:"accounts"`
 }
 
 func (c *Controller) ListConnectedAccounts(
@@ -59,19 +58,14 @@ func (c *Controller) ListConnectedAccounts(
 		return nil, apiErr
 	}
 
-	accounts, apiErr := c.repo.listConnected(ctx, cloudProvider)
+	accountRecords, apiErr := c.repo.listConnected(ctx, cloudProvider)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't list cloud accounts")
 	}
 
-	connectedAccounts := []ConnectedAccount{}
-	for _, a := range accounts {
-		ca, err := a.connectedAccount()
-		if err != nil {
-			zap.L().Error("invalid connected account record", zap.String("accountId", a.Id), zap.Error(err))
-		} else {
-			connectedAccounts = append(connectedAccounts, *ca)
-		}
+	connectedAccounts := []Account{}
+	for _, a := range accountRecords {
+		connectedAccounts = append(connectedAccounts, a.account())
 	}
 
 	return &ConnectedAccountsListResponse{
@@ -159,7 +153,7 @@ type AgentCheckInRequest struct {
 }
 
 type AgentCheckInResponse struct {
-	Account Account `json:"account"`
+	Account AccountRecord `json:"account"`
 }
 
 func (c *Controller) CheckInAsAgent(
@@ -208,19 +202,21 @@ func (c *Controller) UpdateAccountConfig(
 		return nil, apiErr
 	}
 
-	account, apiErr := c.repo.upsert(
+	accountRecord, apiErr := c.repo.upsert(
 		ctx, cloudProvider, &accountId, &req.Config, nil, nil, nil,
 	)
 	if apiErr != nil {
 		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
 	}
 
-	return account, nil
+	account := accountRecord.account()
+
+	return &account, nil
 }
 
 func (c *Controller) DisconnectAccount(
 	ctx context.Context, cloudProvider string, accountId string,
-) (*Account, *model.ApiError) {
+) (*AccountRecord, *model.ApiError) {
 	if apiErr := validateCloudProviderName(cloudProvider); apiErr != nil {
 		return nil, apiErr
 	}
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 122be73a1d0..11714469035 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -8,7 +8,7 @@ import (
 )
 
 // Represents a cloud provider account for cloud integrations
-type Account struct {
+type AccountRecord struct {
 	CloudProvider   string         `json:"cloud_provider" db:"cloud_provider"`
 	Id              string         `json:"id" db:"id"`
 	Config          *AccountConfig `json:"config" db:"config_json"`
@@ -91,7 +91,7 @@ type AccountIntegrationStatus struct {
 	LastHeartbeatTsMillis *int64 `json:"last_heartbeat_ts_ms"`
 }
 
-func (a *Account) status() AccountStatus {
+func (a *AccountRecord) status() AccountStatus {
 	status := AccountStatus{}
 	if a.LastAgentReport != nil {
 		lastHeartbeat := a.LastAgentReport.TimestampMillis
@@ -100,20 +100,16 @@ func (a *Account) status() AccountStatus {
 	return status
 }
 
-func (a *Account) connectedAccount() (*ConnectedAccount, error) {
-	ca := ConnectedAccount{Id: a.Id, Status: a.status()}
+func (a *AccountRecord) account() Account {
+	ca := Account{Id: a.Id, Status: a.status()}
 
-	if a.CloudAccountId == nil {
-		return nil, fmt.Errorf("account %s has nil cloud account id", a.Id)
-	} else {
+	if a.CloudAccountId != nil {
 		ca.CloudAccountId = *a.CloudAccountId
 	}
 
-	if a.Config == nil {
-		return nil, fmt.Errorf("account %s has nil config", a.Id)
-	} else {
+	if a.Config != nil {
 		ca.Config = *a.Config
 	}
 
-	return &ca, nil
+	return ca
 }
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 9e44246f2e4..2e326e1b6a9 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -13,11 +13,11 @@ import (
 
 type cloudProviderAccountsRepository interface {
 	// list connected cloud provider accounts.
-	listConnected(ctx context.Context, cloudProvider string) ([]Account, *model.ApiError)
+	listConnected(ctx context.Context, cloudProvider string) ([]AccountRecord, *model.ApiError)
 
-	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*Account, *model.ApiError)
+	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*AccountRecord, *model.ApiError)
 
-	get(ctx context.Context, cloudProvider string, id string) (*Account, *model.ApiError)
+	get(ctx context.Context, cloudProvider string, id string) (*AccountRecord, *model.ApiError)
 
 	// Insert an account or update it by ID for specified non-empty fields
 	upsert(
@@ -28,7 +28,7 @@ type cloudProviderAccountsRepository interface {
 		cloudAccountId *string,
 		agentReport *AgentReport,
 		removedAt *time.Time,
-	) (*Account, *model.ApiError)
+	) (*AccountRecord, *model.ApiError)
 }
 
 func newCloudProviderAccountsRepository(db *sqlx.DB) (
@@ -76,8 +76,8 @@ type cloudProviderAccountsSQLRepository struct {
 
 func (r *cloudProviderAccountsSQLRepository) listConnected(
 	ctx context.Context, cloudProvider string,
-) ([]Account, *model.ApiError) {
-	accounts := []Account{}
+) ([]AccountRecord, *model.ApiError) {
+	accounts := []AccountRecord{}
 
 	err := r.db.SelectContext(
 		ctx, &accounts, `
@@ -109,8 +109,8 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(
 
 func (r *cloudProviderAccountsSQLRepository) getByIds(
 	ctx context.Context, cloudProvider string, ids []string,
-) (map[string]*Account, *model.ApiError) {
-	accounts := []Account{}
+) (map[string]*AccountRecord, *model.ApiError) {
+	accounts := []AccountRecord{}
 
 	idPlaceholders := []string{}
 	queryArgs := []any{cloudProvider}
@@ -144,7 +144,7 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 		))
 	}
 
-	result := map[string]*Account{}
+	result := map[string]*AccountRecord{}
 	for _, a := range accounts {
 
 		if a.Config == nil {
@@ -160,7 +160,7 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 
 func (r *cloudProviderAccountsSQLRepository) get(
 	ctx context.Context, cloudProvider string, id string,
-) (*Account, *model.ApiError) {
+) (*AccountRecord, *model.ApiError) {
 	res, apiErr := r.getByIds(ctx, cloudProvider, []string{id})
 	if apiErr != nil {
 		return nil, apiErr
@@ -184,7 +184,7 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 	cloudAccountId *string,
 	agentReport *AgentReport,
 	removedAt *time.Time,
-) (*Account, *model.ApiError) {
+) (*AccountRecord, *model.ApiError) {
 	// Insert
 	if id == nil {
 		newId := uuid.NewString()
diff --git a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
index ee2f9292484..9f199b5b06c 100644
--- a/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
+++ b/pkg/query-service/tests/integration/signoz_cloud_integrations_test.go
@@ -237,7 +237,7 @@ func (tb *CloudIntegrationsTestBed) CheckInAsAgentWithQS(
 
 func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 	cloudProvider string, accountId string, newConfig cloudintegrations.AccountConfig,
-) *cloudintegrations.Account {
+) *cloudintegrations.AccountRecord {
 	respDataJson := tb.RequestQS(
 		fmt.Sprintf(
 			"/api/v1/cloud-integrations/%s/accounts/%s/config",
@@ -247,7 +247,7 @@ func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 		},
 	)
 
-	var resp cloudintegrations.Account
+	var resp cloudintegrations.AccountRecord
 	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")
@@ -258,7 +258,7 @@ func (tb *CloudIntegrationsTestBed) UpdateAccountConfigWithQS(
 
 func (tb *CloudIntegrationsTestBed) DisconnectAccountWithQS(
 	cloudProvider string, accountId string,
-) *cloudintegrations.Account {
+) *cloudintegrations.AccountRecord {
 	respDataJson := tb.RequestQS(
 		fmt.Sprintf(
 			"/api/v1/cloud-integrations/%s/accounts/%s/disconnect",
@@ -266,7 +266,7 @@ func (tb *CloudIntegrationsTestBed) DisconnectAccountWithQS(
 		), map[string]any{},
 	)
 
-	var resp cloudintegrations.Account
+	var resp cloudintegrations.AccountRecord
 	err := json.Unmarshal(respDataJson, &resp)
 	if err != nil {
 		tb.t.Fatalf("could not unmarshal apiResponse.Data json into Account")

From 5e1246b45a496a3d12ffc0ada587b4e8ba2fcacf Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 9 Jan 2025 09:47:29 +0530
Subject: [PATCH 33/38] feat: cloud integrations: some more cleanup

---
 .../app/cloudintegrations/controller.go       |  1 +
 .../app/cloudintegrations/model.go            |  2 +
 .../app/cloudintegrations/repo.go             | 53 +++++++++----------
 3 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index f502282817f..e3055f2a5ca 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -83,6 +83,7 @@ type GenerateConnectionUrlRequest struct {
 }
 
 type SigNozAgentConfig struct {
+	// The region in which SigNoz agent should be installed.
 	Region string `json:"region"`
 }
 
diff --git a/pkg/query-service/app/cloudintegrations/model.go b/pkg/query-service/app/cloudintegrations/model.go
index 11714469035..53ad4c853e2 100644
--- a/pkg/query-service/app/cloudintegrations/model.go
+++ b/pkg/query-service/app/cloudintegrations/model.go
@@ -109,6 +109,8 @@ func (a *AccountRecord) account() Account {
 
 	if a.Config != nil {
 		ca.Config = *a.Config
+	} else {
+		ca.Config = DefaultAccountConfig()
 	}
 
 	return ca
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 2e326e1b6a9..98caed6c471 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -12,14 +12,14 @@ import (
 )
 
 type cloudProviderAccountsRepository interface {
-	// list connected cloud provider accounts.
 	listConnected(ctx context.Context, cloudProvider string) ([]AccountRecord, *model.ApiError)
 
 	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*AccountRecord, *model.ApiError)
 
 	get(ctx context.Context, cloudProvider string, id string) (*AccountRecord, *model.ApiError)
 
-	// Insert an account or update it by ID for specified non-empty fields
+	// Insert an account or update it by (cloudProvider, id)
+	// for specified non-empty fields
 	upsert(
 		ctx context.Context,
 		cloudProvider string,
@@ -90,11 +90,11 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(
 				created_at,
 				removed_at
 			from cloud_integrations_accounts
-      where
-        cloud_provider=$1
-        and removed_at is NULL
-        and cloud_account_id is not NULL
-        and last_agent_report_json is not NULL
+			where
+				cloud_provider=$1
+				and removed_at is NULL
+				and cloud_account_id is not NULL
+				and last_agent_report_json is not NULL
 			order by created_at
 		`, cloudProvider,
 	)
@@ -110,7 +110,8 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(
 func (r *cloudProviderAccountsSQLRepository) getByIds(
 	ctx context.Context, cloudProvider string, ids []string,
 ) (map[string]*AccountRecord, *model.ApiError) {
-	accounts := []AccountRecord{}
+
+	records := []AccountRecord{}
 
 	idPlaceholders := []string{}
 	queryArgs := []any{cloudProvider}
@@ -121,7 +122,7 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 	}
 
 	err := r.db.SelectContext(
-		ctx, &accounts, fmt.Sprintf(`
+		ctx, &records, fmt.Sprintf(`
 			select
 				cloud_provider,
 				id,
@@ -145,13 +146,7 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 	}
 
 	result := map[string]*AccountRecord{}
-	for _, a := range accounts {
-
-		if a.Config == nil {
-			config := DefaultAccountConfig()
-			a.Config = &config
-		}
-
+	for _, a := range records {
 		result[a.Id] = &a
 	}
 
@@ -193,31 +188,31 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 
 	// Prepare clause for setting values in `on conflict do update`
 	onConflictSetStmts := []string{}
-	updateColStmt := func(col string) string {
+	setColStatement := func(col string) string {
 		return fmt.Sprintf("%s=excluded.%s", col, col)
 	}
 
 	if config != nil {
 		onConflictSetStmts = append(
-			onConflictSetStmts, updateColStmt("config_json"),
+			onConflictSetStmts, setColStatement("config_json"),
 		)
 	}
 
 	if cloudAccountId != nil {
 		onConflictSetStmts = append(
-			onConflictSetStmts, updateColStmt("cloud_account_id"),
+			onConflictSetStmts, setColStatement("cloud_account_id"),
 		)
 	}
 
 	if agentReport != nil {
 		onConflictSetStmts = append(
-			onConflictSetStmts, updateColStmt("last_agent_report_json"),
+			onConflictSetStmts, setColStatement("last_agent_report_json"),
 		)
 	}
 
 	if removedAt != nil {
 		onConflictSetStmts = append(
-			onConflictSetStmts, updateColStmt("removed_at"),
+			onConflictSetStmts, setColStatement("removed_at"),
 		)
 	}
 
@@ -230,14 +225,14 @@ func (r *cloudProviderAccountsSQLRepository) upsert(
 	}
 
 	insertQuery := fmt.Sprintf(`
-    INSERT INTO cloud_integrations_accounts (
-      cloud_provider,
-      id,
-      config_json,
-      cloud_account_id,
-      last_agent_report_json,
-      removed_at
-    ) values ($1, $2, $3, $4, $5, $6)
+		INSERT INTO cloud_integrations_accounts (
+			cloud_provider,
+			id,
+			config_json,
+			cloud_account_id,
+			last_agent_report_json,
+			removed_at
+		) values ($1, $2, $3, $4, $5, $6)
     %s`, onConflictClause,
 	)
 

From 4125815383a57facb77ac35ebaa90db54dd400f0 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 9 Jan 2025 17:18:21 +0530
Subject: [PATCH 34/38] feat: cloud integrations: remove
 cloudProviderAccountsRepository.GetById

---
 .../app/cloudintegrations/repo.go             | 63 +++++--------------
 1 file changed, 17 insertions(+), 46 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 98caed6c471..1d965312de8 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -2,6 +2,7 @@ package cloudintegrations
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"strings"
 	"time"
@@ -14,8 +15,6 @@ import (
 type cloudProviderAccountsRepository interface {
 	listConnected(ctx context.Context, cloudProvider string) ([]AccountRecord, *model.ApiError)
 
-	getByIds(ctx context.Context, cloudProvider string, ids []string) (map[string]*AccountRecord, *model.ApiError)
-
 	get(ctx context.Context, cloudProvider string, id string) (*AccountRecord, *model.ApiError)
 
 	// Insert an account or update it by (cloudProvider, id)
@@ -107,22 +106,13 @@ func (r *cloudProviderAccountsSQLRepository) listConnected(
 	return accounts, nil
 }
 
-func (r *cloudProviderAccountsSQLRepository) getByIds(
-	ctx context.Context, cloudProvider string, ids []string,
-) (map[string]*AccountRecord, *model.ApiError) {
-
-	records := []AccountRecord{}
-
-	idPlaceholders := []string{}
-	queryArgs := []any{cloudProvider}
-
-	for _, id := range ids {
-		idPlaceholders = append(idPlaceholders, "?")
-		queryArgs = append(queryArgs, id)
-	}
+func (r *cloudProviderAccountsSQLRepository) get(
+	ctx context.Context, cloudProvider string, id string,
+) (*AccountRecord, *model.ApiError) {
+	var result AccountRecord
 
-	err := r.db.SelectContext(
-		ctx, &records, fmt.Sprintf(`
+	err := r.db.GetContext(
+		ctx, &result, `
 			select
 				cloud_provider,
 				id,
@@ -133,42 +123,23 @@ func (r *cloudProviderAccountsSQLRepository) getByIds(
 				removed_at
 			from cloud_integrations_accounts
 			where
-				cloud_provider=?
-				and id in (%s)`,
-			strings.Join(idPlaceholders, ", "),
-		),
-		queryArgs...,
+				cloud_provider=$1
+				and id=$2
+		`,
+		cloudProvider, id,
 	)
-	if err != nil {
-		return nil, model.InternalError(fmt.Errorf(
-			"could not query cloud provider account: %w", err,
-		))
-	}
-
-	result := map[string]*AccountRecord{}
-	for _, a := range records {
-		result[a.Id] = &a
-	}
 
-	return result, nil
-}
-
-func (r *cloudProviderAccountsSQLRepository) get(
-	ctx context.Context, cloudProvider string, id string,
-) (*AccountRecord, *model.ApiError) {
-	res, apiErr := r.getByIds(ctx, cloudProvider, []string{id})
-	if apiErr != nil {
-		return nil, apiErr
-	}
-
-	account := res[id]
-	if account == nil {
+	if err == sql.ErrNoRows {
 		return nil, model.NotFoundError(fmt.Errorf(
 			"couldn't find account with Id %s", id,
 		))
+	} else if err != nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"couldn't query cloud provider accounts: %w", err,
+		))
 	}
 
-	return account, nil
+	return &result, nil
 }
 
 func (r *cloudProviderAccountsSQLRepository) upsert(

From 26bd4e32731302513307d28d6079e594b4f4e84b Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 9 Jan 2025 17:25:34 +0530
Subject: [PATCH 35/38] feat: cloud integrations: shouldn't be able to
 disconnect non-existent account

---
 .../app/cloudintegrations/controller.go          |  9 +++++++--
 .../app/cloudintegrations/controller_test.go     | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index e3055f2a5ca..c9f64ed56c4 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -222,12 +222,17 @@ func (c *Controller) DisconnectAccount(
 		return nil, apiErr
 	}
 
+	account, apiErr := c.repo.get(ctx, cloudProvider, accountId)
+	if apiErr != nil {
+		return nil, model.WrapApiError(apiErr, "couldn't disconnect account")
+	}
+
 	tsNow := time.Now()
-	account, apiErr := c.repo.upsert(
+	account, apiErr = c.repo.upsert(
 		ctx, cloudProvider, &accountId, nil, nil, nil, &tsNow,
 	)
 	if apiErr != nil {
-		return nil, model.WrapApiError(apiErr, "couldn't upsert cloud account")
+		return nil, model.WrapApiError(apiErr, "couldn't disconnect account")
 	}
 
 	return account, nil
diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
index a34d57a0c13..10024119e45 100644
--- a/pkg/query-service/app/cloudintegrations/controller_test.go
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/ee/query-service/model"
 	"go.signoz.io/signoz/pkg/query-service/utils"
 )
 
@@ -84,3 +85,18 @@ func TestAgentCheckIns(t *testing.T) {
 	)
 	require.NotNil(apiErr)
 }
+
+func TestCantDisconnectNonExistentAccount(t *testing.T) {
+	require := require.New(t)
+	testDB, _ := utils.NewTestSqliteDB(t)
+	controller, err := NewController(testDB)
+	require.NoError(err)
+
+	// Attempting to disconnect a non-existent account should return error
+	account, apiErr := controller.DisconnectAccount(
+		context.TODO(), "aws", uuid.NewString(),
+	)
+	require.NotNil(apiErr)
+	require.Equal(model.ErrorNotFound, apiErr.Type())
+	require.Nil(account)
+}

From 2a33c292d177fadd80ec1984ed9db37b91d3ab39 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 9 Jan 2025 18:44:17 +0530
Subject: [PATCH 36/38] feat: cloud integrations: validate agents can't check
 in to cloud account with 2 signoz ids

---
 .../app/cloudintegrations/controller_test.go  | 51 +++++++++++++++++++
 .../app/cloudintegrations/repo.go             | 42 +++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
index 10024119e45..197f84bb388 100644
--- a/pkg/query-service/app/cloudintegrations/controller_test.go
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -84,6 +84,57 @@ func TestAgentCheckIns(t *testing.T) {
 		},
 	)
 	require.NotNil(apiErr)
+
+	// The agent should not be able to check-in with a particular cloud account id
+	// if another connected AccountRecord exists for same cloud account
+	// i.e. there can't be 2 connected account records for the same cloud account id
+	// at any point in time.
+	existingConnected, apiErr := controller.repo.getConnectedCloudAccount(
+		context.TODO(), "aws", testCloudAccountId1,
+	)
+	require.Nil(apiErr)
+	require.NotNil(existingConnected)
+	require.Equal(testCloudAccountId1, *existingConnected.CloudAccountId)
+	require.Nil(existingConnected.RemovedAt)
+
+	testAccountId2 := uuid.NewString()
+	_, apiErr = controller.CheckInAsAgent(
+		context.TODO(), "aws", AgentCheckInRequest{
+			AccountId:      testAccountId2,
+			CloudAccountId: testCloudAccountId1,
+		},
+	)
+	require.NotNil(apiErr)
+
+	// After disconnecting existing account record, the agent should be able to
+	// connected for a particular cloud account id
+	_, apiErr = controller.DisconnectAccount(
+		context.TODO(), "aws", testAccountId1,
+	)
+
+	existingConnected, apiErr = controller.repo.getConnectedCloudAccount(
+		context.TODO(), "aws", testCloudAccountId1,
+	)
+	require.Nil(existingConnected)
+	require.NotNil(apiErr)
+	require.Equal(model.ErrorNotFound, apiErr.Type())
+
+	_, apiErr = controller.CheckInAsAgent(
+		context.TODO(), "aws", AgentCheckInRequest{
+			AccountId:      testAccountId2,
+			CloudAccountId: testCloudAccountId1,
+		},
+	)
+	require.NotNil(apiErr)
+
+	// should be able to keep checking in
+	_, apiErr = controller.CheckInAsAgent(
+		context.TODO(), "aws", AgentCheckInRequest{
+			AccountId:      testAccountId2,
+			CloudAccountId: testCloudAccountId1,
+		},
+	)
+	require.NotNil(apiErr)
 }
 
 func TestCantDisconnectNonExistentAccount(t *testing.T) {
diff --git a/pkg/query-service/app/cloudintegrations/repo.go b/pkg/query-service/app/cloudintegrations/repo.go
index 1d965312de8..e5c4e4cc94f 100644
--- a/pkg/query-service/app/cloudintegrations/repo.go
+++ b/pkg/query-service/app/cloudintegrations/repo.go
@@ -17,6 +17,10 @@ type cloudProviderAccountsRepository interface {
 
 	get(ctx context.Context, cloudProvider string, id string) (*AccountRecord, *model.ApiError)
 
+	getConnectedCloudAccount(
+		ctx context.Context, cloudProvider string, cloudAccountId string,
+	) (*AccountRecord, *model.ApiError)
+
 	// Insert an account or update it by (cloudProvider, id)
 	// for specified non-empty fields
 	upsert(
@@ -142,6 +146,44 @@ func (r *cloudProviderAccountsSQLRepository) get(
 	return &result, nil
 }
 
+func (r *cloudProviderAccountsSQLRepository) getConnectedCloudAccount(
+	ctx context.Context, cloudProvider string, cloudAccountId string,
+) (*AccountRecord, *model.ApiError) {
+	var result AccountRecord
+
+	err := r.db.GetContext(
+		ctx, &result, `
+			select
+				cloud_provider,
+				id,
+				config_json,
+				cloud_account_id,
+				last_agent_report_json,
+				created_at,
+				removed_at
+			from cloud_integrations_accounts
+			where
+				cloud_provider=$1
+				and cloud_account_id=$2
+				and last_agent_report_json is not NULL
+				and removed_at is NULL
+		`,
+		cloudProvider, cloudAccountId,
+	)
+
+	if err == sql.ErrNoRows {
+		return nil, model.NotFoundError(fmt.Errorf(
+			"couldn't find connected cloud account %s", cloudAccountId,
+		))
+	} else if err != nil {
+		return nil, model.InternalError(fmt.Errorf(
+			"couldn't query cloud provider accounts: %w", err,
+		))
+	}
+
+	return &result, nil
+}
+
 func (r *cloudProviderAccountsSQLRepository) upsert(
 	ctx context.Context,
 	cloudProvider string,

From 2ff511ef0e5c52f11a876514a2b41dbcb60d4e3c Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Thu, 9 Jan 2025 18:57:03 +0530
Subject: [PATCH 37/38] feat: cloud integrations: ensure agents can't check in
 to cloud account with 2 signoz ids

---
 .../app/cloudintegrations/controller.go          | 16 ++++++++++++----
 .../app/cloudintegrations/controller_test.go     |  4 ++--
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller.go b/pkg/query-service/app/cloudintegrations/controller.go
index c9f64ed56c4..9cfe9fb3e18 100644
--- a/pkg/query-service/app/cloudintegrations/controller.go
+++ b/pkg/query-service/app/cloudintegrations/controller.go
@@ -164,11 +164,19 @@ func (c *Controller) CheckInAsAgent(
 		return nil, apiErr
 	}
 
-	account, apiErr := c.repo.get(ctx, cloudProvider, req.AccountId)
-	if account != nil && account.CloudAccountId != nil && *account.CloudAccountId != req.CloudAccountId {
+	existingAccount, apiErr := c.repo.get(ctx, cloudProvider, req.AccountId)
+	if existingAccount != nil && existingAccount.CloudAccountId != nil && *existingAccount.CloudAccountId != req.CloudAccountId {
 		return nil, model.BadRequest(fmt.Errorf(
 			"can't check in with new %s account id %s for account %s with existing %s id %s",
-			cloudProvider, req.CloudAccountId, account.Id, cloudProvider, *account.CloudAccountId,
+			cloudProvider, req.CloudAccountId, existingAccount.Id, cloudProvider, *existingAccount.CloudAccountId,
+		))
+	}
+
+	existingAccount, apiErr = c.repo.getConnectedCloudAccount(ctx, cloudProvider, req.CloudAccountId)
+	if existingAccount != nil && existingAccount.Id != req.AccountId {
+		return nil, model.BadRequest(fmt.Errorf(
+			"can't check in to %s account %s with id %s. already connected with id %s",
+			cloudProvider, req.CloudAccountId, req.AccountId, existingAccount.Id,
 		))
 	}
 
@@ -177,7 +185,7 @@ func (c *Controller) CheckInAsAgent(
 		Data:            req.Data,
 	}
 
-	account, apiErr = c.repo.upsert(
+	account, apiErr := c.repo.upsert(
 		ctx, cloudProvider, &req.AccountId, nil, &req.CloudAccountId, &agentReport, nil,
 	)
 	if apiErr != nil {
diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
index 197f84bb388..6dfeb479bf1 100644
--- a/pkg/query-service/app/cloudintegrations/controller_test.go
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -125,7 +125,7 @@ func TestAgentCheckIns(t *testing.T) {
 			CloudAccountId: testCloudAccountId1,
 		},
 	)
-	require.NotNil(apiErr)
+	require.Nil(apiErr)
 
 	// should be able to keep checking in
 	_, apiErr = controller.CheckInAsAgent(
@@ -134,7 +134,7 @@ func TestAgentCheckIns(t *testing.T) {
 			CloudAccountId: testCloudAccountId1,
 		},
 	)
-	require.NotNil(apiErr)
+	require.Nil(apiErr)
 }
 
 func TestCantDisconnectNonExistentAccount(t *testing.T) {

From 35e10272a3a017203fde499308d100f31e68dd36 Mon Sep 17 00:00:00 2001
From: Raj <rkssisodiya@gmail.com>
Date: Fri, 10 Jan 2025 10:34:08 +0530
Subject: [PATCH 38/38] feat: cloud integrations: remove stray import of
 ee/model in cloudintegrations controller

---
 pkg/query-service/app/cloudintegrations/controller_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/query-service/app/cloudintegrations/controller_test.go b/pkg/query-service/app/cloudintegrations/controller_test.go
index 6dfeb479bf1..2d85cc1b958 100644
--- a/pkg/query-service/app/cloudintegrations/controller_test.go
+++ b/pkg/query-service/app/cloudintegrations/controller_test.go
@@ -6,7 +6,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
-	"go.signoz.io/signoz/ee/query-service/model"
+	"go.signoz.io/signoz/pkg/query-service/model"
 	"go.signoz.io/signoz/pkg/query-service/utils"
 )