Skip to content

Commit eb28d9b

Browse files
keyihuangkeyihuang
committed
feat: introduce diagnostic tool api (#13845)
1 parent f264427 commit eb28d9b

File tree

4 files changed

+262
-0
lines changed

4 files changed

+262
-0
lines changed

‎airbyte-api/server-api/src/main/openapi/config.yaml

+30
Original file line numberDiff line numberDiff line change
@@ -5198,6 +5198,28 @@ paths:
51985198
schema:
51995199
$ref: "#/components/schemas/UserInvitationRead"
52005200

5201+
/v1/diagnostic_tool/generate_report:
5202+
post:
5203+
summary: Generate a diagnostic report
5204+
tags:
5205+
- diagnostic_tool
5206+
operationId: generateDiagnosticReport
5207+
requestBody:
5208+
content:
5209+
application/json:
5210+
schema:
5211+
$ref: "#/components/schemas/DiagnosticReportRequestBody"
5212+
responses:
5213+
"200":
5214+
description: Successfully generated diagnostic report
5215+
content:
5216+
application/zip:
5217+
schema:
5218+
type: string
5219+
format: binary
5220+
"422":
5221+
$ref: "#/components/responses/InvalidInputResponse"
5222+
52015223
# Airbyte API
52025224
/public:
52035225
get:
@@ -14233,6 +14255,14 @@ components:
1423314255
format: date-time
1423414256
x-field-extra-annotation: '@com.fasterxml.jackson.annotation.JsonFormat(pattern="yyyy-MM-dd''T''HH:mm:ss.SSS''Z''")'
1423514257

14258+
DiagnosticReportRequestBody:
14259+
type: object
14260+
required:
14261+
- organizationId
14262+
properties:
14263+
organizationId:
14264+
$ref: "#/components/schemas/OrganizationId"
14265+
1423614266
responses:
1423714267
NotFoundResponse:
1423814268
description: Object with given id was not found.

‎airbyte-commons-server/src/main/java/io/airbyte/commons/server/handlers/DiagnosticToolHandler.java

+116
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved.
3+
*/
4+
5+
package io.airbyte.commons.server.handlers;
6+
7+
import jakarta.inject.Singleton;
8+
import java.io.ByteArrayOutputStream;
9+
import java.io.File;
10+
import java.io.FileOutputStream;
11+
import java.io.IOException;
12+
import java.util.zip.ZipEntry;
13+
import java.util.zip.ZipOutputStream;
14+
import org.slf4j.Logger;
15+
import org.slf4j.LoggerFactory;
16+
17+
/**
18+
* DiagnosticToolHandler.
19+
*/
20+
@Singleton
21+
public class DiagnosticToolHandler {
22+
23+
private static final Logger LOGGER = LoggerFactory.getLogger(JobInputHandler.class);
24+
public static final String APPLICATION_YAML = "application.yaml";
25+
public static final String DEPLOYMENT_YAML = "deployment.yaml";
26+
public static final String DIAGNOSTIC_REPORT_FILE_NAME = "diagnostic_report";
27+
public static final String DIAGNOSTIC_REPORT_FILE_FORMAT = ".zip";
28+
29+
private byte[] generateZipInMemory() throws IOException {
30+
ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
31+
ZipOutputStream zipOut = new ZipOutputStream(byteArrayOutputStream);
32+
try {
33+
addApplicationYaml(zipOut);
34+
} catch (IOException e) {
35+
LOGGER.error("Error in writing application data.", e);
36+
}
37+
try {
38+
addDeploymentYaml(zipOut);
39+
} catch (IOException e) {
40+
LOGGER.error("Error in writing deployment data.", e);
41+
}
42+
zipOut.finish();
43+
return byteArrayOutputStream.toByteArray();
44+
}
45+
46+
private void addApplicationYaml(ZipOutputStream zipOut) throws IOException {
47+
48+
// In-memory construct an entry (application.yaml file) in the final zip output.
49+
ZipEntry applicationYaml = new ZipEntry(APPLICATION_YAML);
50+
zipOut.putNextEntry(applicationYaml);
51+
52+
// Write application information to the zip entry: application.yaml.
53+
String applicationYamlContent = generateApplicationYaml();
54+
zipOut.write(applicationYamlContent.getBytes());
55+
zipOut.closeEntry();
56+
}
57+
58+
private String generateApplicationYaml() {
59+
// Collect workspace information and write it to `application.yaml`.
60+
return collectWorkspaceInfo();
61+
// Collect other information here, e.g. logs
62+
}
63+
64+
/**
65+
* Collect workspace information and write it to `application.yaml`.
66+
*/
67+
private String collectWorkspaceInfo() {
68+
LOGGER.info("Collecting workspace data...");
69+
// TODO: implement this
70+
return "Workspaces:\n...";
71+
}
72+
73+
private void addDeploymentYaml(ZipOutputStream zipOut) throws IOException {
74+
// In-memory construct an entry (deployment.yaml file) in the final zip output.
75+
ZipEntry zipEntry = new ZipEntry(DEPLOYMENT_YAML);
76+
zipOut.putNextEntry(zipEntry);
77+
String deploymentYamlContent = generateDeploymentYaml();
78+
zipOut.write(deploymentYamlContent.getBytes());
79+
zipOut.closeEntry();
80+
}
81+
82+
/**
83+
* Collect Kubernetes information and write it to deployment_info.yaml.
84+
*/
85+
private String generateDeploymentYaml() {
86+
LOGGER.info("Collecting k8s data...");
87+
// TODO: implement this
88+
return "k8s: \n";
89+
}
90+
91+
/**
92+
* Generate diagnostic report by collecting relevant data and zipping them into a single file.
93+
*
94+
* @return File - generated zip file as a diagnostic report
95+
*/
96+
public File generateDiagnosticReport() {
97+
try {
98+
// Generate the zip file content in memory as byte[]
99+
byte[] zipFileContent = generateZipInMemory();
100+
101+
// Write the byte[] to a temporary file
102+
File tempFile = File.createTempFile(DIAGNOSTIC_REPORT_FILE_NAME, DIAGNOSTIC_REPORT_FILE_FORMAT);
103+
try (FileOutputStream fos = new FileOutputStream(tempFile)) {
104+
fos.write(zipFileContent);
105+
}
106+
107+
// Return the temporary file
108+
return tempFile;
109+
110+
} catch (IOException e) {
111+
LOGGER.error("Error generating diagnostic report", e);
112+
return null;
113+
}
114+
}
115+
116+
}

‎airbyte-commons-server/src/test/java/io/airbyte/commons/server/handlers/DiagnosticToolHandlerTest.java

+74
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved.
3+
*/
4+
5+
package io.airbyte.commons.server.handlers;
6+
7+
import java.io.File;
8+
import java.io.FileInputStream;
9+
import java.io.IOException;
10+
import java.util.zip.ZipEntry;
11+
import java.util.zip.ZipInputStream;
12+
import org.junit.jupiter.api.Assertions;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
16+
/**
17+
* DiagnosticToolHandlerTest.
18+
*/
19+
@SuppressWarnings("PMD")
20+
class DiagnosticToolHandlerTest {
21+
22+
private DiagnosticToolHandler diagnosticToolHandler;
23+
24+
@BeforeEach
25+
void beforeEach() {
26+
diagnosticToolHandler = new DiagnosticToolHandler();
27+
}
28+
29+
@Test
30+
void testGenerateDiagnosticReport() throws IOException {
31+
final File zipFile = diagnosticToolHandler.generateDiagnosticReport();
32+
Assertions.assertTrue(zipFile.exists());
33+
// Check the content of the zip file
34+
try (FileInputStream fis = new FileInputStream(zipFile);
35+
ZipInputStream zis = new ZipInputStream(fis)) {
36+
37+
ZipEntry entry;
38+
boolean foundApplicationYaml = false;
39+
boolean foundDeploymentYaml = false;
40+
41+
// Iterate through the entries in the zip
42+
while ((entry = zis.getNextEntry()) != null) {
43+
if (entry.getName().equals(DiagnosticToolHandler.APPLICATION_YAML)) {
44+
foundApplicationYaml = true;
45+
46+
// Check the content of application.yaml
47+
byte[] buffer = new byte[1024];
48+
int bytesRead;
49+
StringBuilder content = new StringBuilder();
50+
while ((bytesRead = zis.read(buffer)) != -1) {
51+
content.append(new String(buffer, 0, bytesRead));
52+
}
53+
Assertions.assertTrue(content.toString().contains("Workspaces"));
54+
} else if (entry.getName().equals(DiagnosticToolHandler.DEPLOYMENT_YAML)) {
55+
foundDeploymentYaml = true;
56+
57+
// Check the content of deployment.yaml
58+
byte[] buffer = new byte[1024];
59+
int bytesRead;
60+
StringBuilder content = new StringBuilder();
61+
while ((bytesRead = zis.read(buffer)) != -1) {
62+
content.append(new String(buffer, 0, bytesRead));
63+
}
64+
Assertions.assertTrue(content.toString().contains("k8s"));
65+
}
66+
}
67+
68+
// Ensure both files are present in the zip
69+
Assertions.assertTrue(foundApplicationYaml);
70+
Assertions.assertTrue(foundDeploymentYaml);
71+
}
72+
}
73+
74+
}

‎airbyte-server/src/main/java/io/airbyte/server/apis/DiagnosticToolApiController.java

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved.
3+
*/
4+
5+
package io.airbyte.server.apis;
6+
7+
import static io.airbyte.commons.auth.AuthRoleConstants.ORGANIZATION_READER;
8+
import static io.airbyte.commons.auth.AuthRoleConstants.WORKSPACE_READER;
9+
10+
import io.airbyte.api.generated.DiagnosticToolApi;
11+
import io.airbyte.api.model.generated.DiagnosticReportRequestBody;
12+
import io.airbyte.commons.server.handlers.DiagnosticToolHandler;
13+
import io.airbyte.commons.server.scheduling.AirbyteTaskExecutors;
14+
import io.micronaut.http.MediaType;
15+
import io.micronaut.http.annotation.Body;
16+
import io.micronaut.http.annotation.Controller;
17+
import io.micronaut.http.annotation.Post;
18+
import io.micronaut.scheduling.annotation.ExecuteOn;
19+
import io.micronaut.security.annotation.Secured;
20+
import io.micronaut.security.rules.SecurityRule;
21+
import java.io.File;
22+
23+
@Controller("/api/v1/diagnostic_tool")
24+
@Secured(SecurityRule.IS_AUTHENTICATED)
25+
public class DiagnosticToolApiController implements DiagnosticToolApi {
26+
27+
private final DiagnosticToolHandler diagnosticToolHandler;
28+
29+
public DiagnosticToolApiController(final DiagnosticToolHandler diagnosticToolHandler) {
30+
this.diagnosticToolHandler = diagnosticToolHandler;
31+
}
32+
33+
@Override
34+
@Post(uri = "/generate_report",
35+
produces = MediaType.APPLICATION_ZIP)
36+
@ExecuteOn(AirbyteTaskExecutors.IO)
37+
@Secured({WORKSPACE_READER, ORGANIZATION_READER})
38+
public File generateDiagnosticReport(@Body final DiagnosticReportRequestBody diagnosticReportRequestBody) {
39+
return ApiHelper.execute(diagnosticToolHandler::generateDiagnosticReport);
40+
}
41+
42+
}

0 commit comments

Comments
 (0)