[CLODUGA-20271] Add wait for DB audit commands (#231)

Co-authored-by: aquatiko <[email protected]>
yugabyte · Apr 29, 2024 · 30ff1a1 · 30ff1a1
1 parent 1fb067f
commit 30ff1a1
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 12 deletions.
diff --git a/cmd/db_audit_logs_exporter/db_audit_logs_exporter.go b/cmd/db_audit_logs_exporter/db_audit_logs_exporter.go
@@ -80,19 +80,38 @@ var assignDbAuditLogsExporterCmd = &cobra.Command{
 			logrus.Debugf("Full HTTP response: %v", r)
 			logrus.Fatalf(ybmAuthClient.GetApiErrorDetails(err))
 		}
+		respData := resp.GetData()
 
-		dbAuditIntegrationId := resp.GetData().Info.Id
+		dbAuditLogConfigId := respData.Info.Id
 
-		msg := fmt.Sprintf("The db audit exporter config %s is being created", formatter.Colorize(dbAuditIntegrationId, formatter.GREEN_COLOR))
+		msg := fmt.Sprintf("The db audit exporter config %s is being created", formatter.Colorize(dbAuditLogConfigId, formatter.GREEN_COLOR))
 
-		fmt.Println(msg)
+		if viper.GetBool("wait") {
+			returnStatus, err := authApi.WaitForTaskCompletion(clusterId, ybmclient.ENTITYTYPEENUM_CLUSTER, ybmclient.TASKTYPEENUM_ENABLE_DATABASE_AUDIT_LOGGING, []string{"FAILED", "SUCCEEDED"}, msg)
+			if err != nil {
+				logrus.Fatalf("error when getting task status: %s", err)
+			}
+			if returnStatus != "SUCCEEDED" {
+				logrus.Fatalf("Operation failed with error: %s", returnStatus)
+			}
+			fmt.Printf("DB audit logging %v has been started on the cluster %v\n", formatter.Colorize(dbAuditLogConfigId, formatter.GREEN_COLOR), formatter.Colorize(clusterName, formatter.GREEN_COLOR))
+
+			respC, r, err := authApi.ListDbAuditExporterConfig(clusterId).Execute()
+			if err != nil {
+				logrus.Debugf("Full HTTP response: %v", r)
+				logrus.Fatalf(ybmAuthClient.GetApiErrorDetails(err))
+			}
+			respData = respC.GetData()[0]
+		} else {
+			fmt.Println(msg)
+		}
 
 		dbAuditLogsExporterCtx := formatter.Context{
 			Output: os.Stdout,
 			Format: formatter.NewDbAuditLogsExporterFormat(viper.GetString("output")),
 		}
 
-		formatter.DbAuditLogsExporterWrite(dbAuditLogsExporterCtx, []openapi.DbAuditExporterConfigurationData{resp.GetData()})
+		formatter.DbAuditLogsExporterWrite(dbAuditLogsExporterCtx, []openapi.DbAuditExporterConfigurationData{respData})
 	},
 }
 
@@ -137,18 +156,38 @@ var updateDbAuditLogsExporterCmd = &cobra.Command{
 			logrus.Fatalf(ybmAuthClient.GetApiErrorDetails(err))
 		}
 
-		dbAuditIntegrationId := resp.GetData().Info.Id
+		respData := resp.GetData()
+
+		dbAuditLogConfigId := respData.Info.Id
+
+		msg := fmt.Sprintf("The db audit exporter config %s is being updated", formatter.Colorize(dbAuditLogConfigId, formatter.GREEN_COLOR))
 
-		msg := fmt.Sprintf("The db audit exporter config %s is being updated", formatter.Colorize(dbAuditIntegrationId, formatter.GREEN_COLOR))
+		if viper.GetBool("wait") {
+			returnStatus, err := authApi.WaitForTaskCompletion(clusterId, ybmclient.ENTITYTYPEENUM_CLUSTER, ybmclient.TASKTYPEENUM_EDIT_DATABASE_AUDIT_LOGGING, []string{"FAILED", "SUCCEEDED"}, msg)
+			if err != nil {
+				logrus.Fatalf("error when getting task status: %s", err)
+			}
+			if returnStatus != "SUCCEEDED" {
+				logrus.Fatalf("Operation failed with error: %s", returnStatus)
+			}
+			fmt.Printf("DB audit logging configuration %v has been updated on the cluster %v\n", formatter.Colorize(dbAuditLogConfigId, formatter.GREEN_COLOR), formatter.Colorize(clusterName, formatter.GREEN_COLOR))
 
-		fmt.Println(msg)
+			respC, r, err := authApi.ListDbAuditExporterConfig(clusterId).Execute()
+			if err != nil {
+				logrus.Debugf("Full HTTP response: %v", r)
+				logrus.Fatalf(ybmAuthClient.GetApiErrorDetails(err))
+			}
+			respData = respC.GetData()[0]
+		} else {
+			fmt.Println(msg)
+		}
 
 		dbAuditLogsExporterCtx := formatter.Context{
 			Output: os.Stdout,
 			Format: formatter.NewDbAuditLogsExporterFormat(viper.GetString("output")),
 		}
 
-		formatter.DbAuditLogsExporterWrite(dbAuditLogsExporterCtx, []openapi.DbAuditExporterConfigurationData{resp.GetData()})
+		formatter.DbAuditLogsExporterWrite(dbAuditLogsExporterCtx, []openapi.DbAuditExporterConfigurationData{respData})
 	},
 }
 
@@ -222,9 +261,24 @@ var removeDbAuditLogsExporterCmd = &cobra.Command{
 
 		if err != nil {
 			logrus.Debugf("Full HTTP response: %v", resp)
+			logrus.Fatalf(ybmAuthClient.GetApiErrorDetails(err))
 		}
 
-		fmt.Printf("Deleting Db Audit Logs Exporter Config %s\n", formatter.Colorize(exportConfigId, formatter.GREEN_COLOR))
+		msg := fmt.Sprintf("Request submitted to remove Db Audit Logging %s\n", formatter.Colorize(exportConfigId, formatter.GREEN_COLOR))
+
+		if viper.GetBool("wait") {
+			returnStatus, err := authApi.WaitForTaskCompletion(clusterId, ybmclient.ENTITYTYPEENUM_CLUSTER, ybmclient.TASKTYPEENUM_DISABLE_DATABASE_AUDIT_LOGGING, []string{"FAILED", "SUCCEEDED"}, msg)
+			if err != nil {
+				logrus.Fatalf("error when getting task status: %s", err)
+			}
+			if returnStatus != "SUCCEEDED" {
+				logrus.Fatalf("Operation failed with error: %s", returnStatus)
+			}
+			fmt.Printf("DB audit logging %v has been removed from the cluster %v\n", formatter.Colorize(exportConfigId, formatter.GREEN_COLOR), formatter.Colorize(clusterName, formatter.GREEN_COLOR))
+			return
+		} else {
+			fmt.Println(msg)
+		}
 	},
 }
 

diff --git a/cmd/db_audit_logs_exporter_test.go b/cmd/db_audit_logs_exporter_test.go
@@ -185,19 +185,19 @@ ID                                     Date Created               Cluster ID
 	Context("When removing db audit exporter config", func() {
 		It("should delete the config", func() {
 			statusCode = 200
-			err := loadJson("./test/fixtures/list-db-audit.json", &responseDbAuditList)
+			err := loadJson("./test/fixtures/db-audit-data.json", &responseDbAudit)
 			Expect(err).ToNot(HaveOccurred())
 			server.AppendHandlers(
 				ghttp.CombineHandlers(
 					ghttp.VerifyRequest(http.MethodDelete, "/api/public/v1/accounts/340af43a-8a7c-4659-9258-4876fd6a207b/projects/78d4459c-0f45-47a5-899a-45ddf43eba6e/clusters/5f80730f-ba3f-4f7e-8c01-f8fa4c90dad8/db-audit-log-exporter-configs/123e4567-e89b-12d3-a456-426614174000"),
-					ghttp.RespondWithJSONEncodedPtr(&statusCode, responseDbAuditList),
+					ghttp.RespondWithJSONEncodedPtr(&statusCode, responseDbAudit),
 				),
 			)
 			cmd := exec.Command(compiledCLIPath, "db-audit-logs-exporter", "unassign", "--cluster-name", "stunning-sole", "--export-config-id", "123e4567-e89b-12d3-a456-426614174000", "--force")
 			session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
 			Expect(err).NotTo(HaveOccurred())
 			session.Wait(2)
-			Expect(session.Out).Should(gbytes.Say(`Deleting Db Audit Logs Exporter Config 123e4567-e89b-12d3-a456-426614174000`))
+			Expect(session.Out).Should(gbytes.Say(`Request submitted to remove Db Audit Logging 123e4567-e89b-12d3-a456-426614174000`))
 			session.Kill()
 		})
 		It("should return required field name and type when not set", func() {