yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-09-02 13:43:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix #4089: Test Connection - Ingestion Resource; Fix #4129: Status -Ingestion Resource (#4215)

Browse Source 
* Fix #4089: Test Connection - Ingestion Resource; Fix #4129: Status - Ingestion Resource

* use TestServiceConnectionRequest

* Simplify test service connection

* Add exception info

* Update pipelineStatus

* Deprecate old class

* Use JSON Schema models

Co-authored-by: Pere Miquel Brull <peremiquelbrull@gmail.com>
This commit is contained in:
Sriharsha Chintalapani 2022-04-19 09:54:10 -07:00 committed by GitHub
parent 09f01fc7e3
commit 6b77ab29f2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 229 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
catalog-rest-service/src/main/java/org/openmetadata/catalog/airflow/AirflowRESTClient.java
View File

@ -21,12 +21,15 @@ import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.List;
import javax.ws.rs.core.Response;
import lombok.extern.slf4j.Slf4j;
import org.json.JSONObject;
import org.openmetadata.catalog.airflow.models.AirflowAuthRequest;
import org.openmetadata.catalog.airflow.models.AirflowAuthResponse;
import org.openmetadata.catalog.api.services.ingestionPipelines.TestServiceConnection;
import org.openmetadata.catalog.entity.services.ingestionPipelines.IngestionPipeline;
import org.openmetadata.catalog.entity.services.ingestionPipelines.PipelineStatus;
import org.openmetadata.catalog.exception.AirflowException;
import org.openmetadata.catalog.exception.IngestionPipelineDeploymentException;
import org.openmetadata.catalog.util.JsonUtils;
@ -150,4 +153,57 @@ public class AirflowRESTClient {
throw IngestionPipelineDeploymentException.byMessage(pipelineName, e.getMessage());
}
}
public IngestionPipeline getStatus(IngestionPipeline ingestionPipeline) {
try {
String token = authenticate();
String authToken = String.format(AUTH_TOKEN, token);
String statusEndPoint = "%s/rest_api/api?api=dag_status&dag_id=%s";
String statusUrl = String.format(statusEndPoint, airflowURL, ingestionPipeline.getName());
JSONObject requestPayload = new JSONObject();
HttpRequest request =
HttpRequest.newBuilder(URI.create(statusUrl))
.header(CONTENT_HEADER, CONTENT_TYPE)
.header(AUTH_HEADER, authToken)
.POST(HttpRequest.BodyPublishers.ofString(requestPayload.toString()))
.build();
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
List<PipelineStatus> statuses = JsonUtils.readObjects(response.body(), PipelineStatus.class);
ingestionPipeline.setPipelineStatuses(statuses);
return ingestionPipeline;
}
throw AirflowException.byMessage(
ingestionPipeline.getName(),
"Failed to fetch ingestion pipeline runs",
Response.Status.fromStatusCode(response.statusCode()));
} catch (Exception e) {
throw AirflowException.byMessage(ingestionPipeline.getName(), e.getMessage());
}
}
public HttpResponse<String> testConnection(TestServiceConnection testServiceConnection) {
try {
String token = authenticate();
String authToken = String.format(AUTH_TOKEN, token);
String statusEndPoint = "%s/rest_api/api?api=test_connection";
String statusUrl = String.format(statusEndPoint, airflowURL);
String connectionPayload = JsonUtils.pojoToJson(testServiceConnection);
HttpRequest request =
HttpRequest.newBuilder(URI.create(statusUrl))
.header(CONTENT_HEADER, CONTENT_TYPE)
.header(AUTH_HEADER, authToken)
.POST(HttpRequest.BodyPublishers.ofString(connectionPayload))
.build();
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
return response;
}
throw AirflowException.byMessage(
"Failed to test connection.", String.valueOf(Response.Status.fromStatusCode(response.statusCode())));
} catch (Exception e) {
throw AirflowException.byMessage("Failed to test connection.", e.getMessage());
}
}
}

58
catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/ingestionpipelines/IngestionPipelineResource.java
View File

@ -17,6 +17,7 @@ import static org.openmetadata.catalog.Entity.FIELD_OWNER;
import static org.openmetadata.catalog.security.SecurityUtil.ADMIN;
import static org.openmetadata.catalog.security.SecurityUtil.BOT;
import static org.openmetadata.catalog.security.SecurityUtil.OWNER;
import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
import io.swagger.annotations.Api;
import io.swagger.v3.oas.annotations.ExternalDocumentation;
@ -28,6 +29,7 @@ import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.parameters.RequestBody;
import io.swagger.v3.oas.annotations.responses.ApiResponse;
import java.io.IOException;
import java.net.http.HttpResponse;
import java.util.List;
import java.util.Locale;
import java.util.UUID;
@ -57,6 +59,7 @@ import org.openmetadata.catalog.Entity;
import org.openmetadata.catalog.airflow.AirflowConfiguration;
import org.openmetadata.catalog.airflow.AirflowRESTClient;
import org.openmetadata.catalog.api.services.ingestionPipelines.CreateIngestionPipeline;
import org.openmetadata.catalog.api.services.ingestionPipelines.TestServiceConnection;
import org.openmetadata.catalog.entity.services.DashboardService;
import org.openmetadata.catalog.entity.services.DatabaseService;
import org.openmetadata.catalog.entity.services.MessagingService;
@ -164,7 +167,12 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
Include include)
throws IOException {
ListFilter filter = new ListFilter(include).addQueryParam("service", serviceParam);
return super.listInternal(uriInfo, securityContext, fieldsParam, filter, limitParam, before, after);
ResultList<IngestionPipeline> ingestionPipelines =
super.listInternal(uriInfo, securityContext, fieldsParam, filter, limitParam, before, after);
if (fieldsParam != null && fieldsParam.contains("pipelineStatuses")) {
addStatus(ingestionPipelines.getData());
}
return ingestionPipelines;
}
@GET
@ -217,7 +225,11 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
@DefaultValue("non-deleted")
Include include)
throws IOException {
return getInternal(uriInfo, securityContext, id, fieldsParam, include);
IngestionPipeline ingestionPipeline = getInternal(uriInfo, securityContext, id, fieldsParam, include);
if (fieldsParam != null && fieldsParam.contains("pipelineStatuses")) {
ingestionPipeline = addStatus(ingestionPipeline);
}
return ingestionPipeline;
}
@GET
@ -279,7 +291,11 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
@DefaultValue("non-deleted")
Include include)
throws IOException {
return getByNameInternal(uriInfo, securityContext, fqn, fieldsParam, include);
IngestionPipeline ingestionPipeline = getByNameInternal(uriInfo, securityContext, fqn, fieldsParam, include);
if (fieldsParam != null && fieldsParam.contains("pipelineStatuses")) {
ingestionPipeline = addStatus(ingestionPipeline);
}
return ingestionPipeline;
}
@POST
@ -376,6 +392,29 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
return addHref(uriInfo, dao.get(uriInfo, id, fields));
}
@POST
@Path("/testConnection")
@Operation(
summary = "Test Connection of a Service",
tags = "IngestionPipelines",
description = "Test Connection of a Service.",
responses = {
@ApiResponse(
responseCode = "200",
description = "The ingestion",
content =
@Content(mediaType = "application/json", schema = @Schema(implementation = IngestionPipeline.class))),
@ApiResponse(responseCode = "404", description = "Ingestion for instance {name} is not found")
})
public Response testIngestion(
@Context UriInfo uriInfo,
@Context SecurityContext securityContext,
@Valid TestServiceConnection testServiceConnection)
throws IOException {
HttpResponse<String> response = airflowRESTClient.testConnection(testServiceConnection);
return Response.status(200, response.body()).build();
}
@DELETE
@Path("/{id}")
@Operation(
@ -463,4 +502,17 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
source.setSourceConfig(create.getSourceConfig());
return source;
}
public void addStatus(List<IngestionPipeline> ingestionPipelines) {
listOrEmpty(ingestionPipelines).forEach(this::addStatus);
}
private IngestionPipeline addStatus(IngestionPipeline ingestionPipeline) {
try {
ingestionPipeline = airflowRESTClient.getStatus(ingestionPipeline);
} catch (Exception e) {
LOG.error("Failed to fetch status for {} due to {}", ingestionPipeline.getName(), e);
}
return ingestionPipeline;
}
}

40
catalog-rest-service/src/main/resources/json/schema/api/services/ingestionPipelines/testServiceConnection.json Normal file
View File

@ -0,0 +1,40 @@
{
"$id": "https://open-metadata.org/schema/api/services/ingestionPipelines/testServiceConnection.json",
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "TestServiceConnectionRequest",
"description": "Test Service Connection to test user provided configuration is valid or not.",
"type": "object",
"properties": {
"connection": {
"description": "Database Connection.",
"oneOf": [
{
"$ref": "../../../entity/services/databaseService.json#/definitions/databaseConnection"
},
{
"$ref": "../../../entity/services/dashboardService.json#/definitions/dashboardConnection"
},
{
"$ref": "../../../entity/services/messagingService.json#/definitions/messagingConnection"
}
]
},
"connectionType": {
"description": "Type of database service such as MySQL, BigQuery, Snowflake, Redshift, Postgres...",
"type": "string",
"enum": ["Database", "Dashboard", "Messaging"],
"javaEnums": [
{
"name": "Database"
},
{
"name": "Dashboard"
},
{
"name": "Messaging"
}
]
}
},
"additionalProperties": false
}

36
catalog-rest-service/src/main/resources/json/schema/entity/services/ingestionPipelines/ingestionPipeline.json
View File

@ -11,6 +11,30 @@
"javaType": "org.openmetadata.catalog.entity.services.ingestionPipelines.PipelineType",
"enum": ["metadata", "usage"]
},
"pipelineStatus": {
"type": "object",
"javaType": "org.openmetadata.catalog.entity.services.ingestionPipelines.PipelineStatus",
"description": "This defines runtime status of Pipeline.",
"properties": {
"runId": {
"description": "Pipeline unique run ID.",
"type": "string"
},
"state": {
"description": "Pipeline status denotes if its failed or succeeded.",
"type": "string"
},
"startDate": {
"description": "startDate of the pipeline run for this particular execution.",
"type": "string"
},
"endDate": {
"description": "endDate of the pipeline run for this particular execution.",
"type": "string"
}
},
"additionalProperties": false
},
"airflowConfig": {
"description": "Properties to configure the Airflow pipeline that will run the workflow.",
"type": "object",
@ -151,6 +175,18 @@
"description": "Link to the database service where this database is hosted in.",
"$ref": "../../../type/entityReference.json"
},
"pipelineStatuses": {
"description": "List of executions and status for the Pipeline.",
"type": "array",
"items": {
"$ref": "#/definitions/pipelineStatus"
},
"default": null
},
"nextExecutionDate": {
"description": "Next execution date from the underlying pipeline platform once the pipeline scheduled.",
"$ref": "../../../type/basic.json#/definitions/date"
},
"href": {
"description": "Link to this ingestion pipeline resource.",
"$ref": "../../../type/basic.json#/definitions/href"

2
openmetadata-airflow-apis/src/openmetadata/api/apis_metadata.py
View File

@ -48,7 +48,7 @@ APIS_METADATA = [
"post_arguments": [
{
"name": "service_connection",
"description": "ServiceConnectionModel config to test",
"description": "TestServiceConnectionRequest config to test",
"required": True,
},
],

69
openmetadata-airflow-apis/src/openmetadata/api/response.py
View File

@ -10,11 +10,15 @@
# limitations under the License.
import json
from typing import Optional
from typing import Optional, Union
from airflow.models import DagRun
from flask import Response
from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
PipelineStatus,
)
class ApiResponse:
"""
@ -37,12 +41,8 @@ class ApiResponse:
return resp
@staticmethod
def success(response_obj: Optional[dict] = None):
if not response_obj:
response_obj = {}
response_obj["status"] = "success"
return ApiResponse.standard_response(ApiResponse.STATUS_OK, response_obj)
def success(response_obj: Union[Optional[dict], Optional[list]] = None):
return ApiResponse.standard_response(ApiResponse.STATUS_OK, response_obj or {})
@staticmethod
def error(status, error):
@ -70,47 +70,14 @@ class ResponseFormat:
pass
@staticmethod
def format_dag_run_state(dag_run: DagRun):
return {
"state": dag_run.get_state(),
"run_id": dag_run.run_id,
"startDate": (
None
if not dag_run.start_date
else dag_run.start_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
),
"endDate": (
None
if not dag_run.end_date
else dag_run.end_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
),
}
@staticmethod
def format_dag_task(task_instance):
return {
"taskId": task_instance.task_id,
"dagId": task_instance.dag_id,
"state": task_instance.state,
"tryNumber": (
None
if not task_instance._try_number
else str(task_instance._try_number)
),
"maxTries": (
None if not task_instance.max_tries else str(task_instance.max_tries)
),
"startDate": (
None
if not task_instance.start_date
else task_instance.start_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
),
"endDate": (
None
if not task_instance.end_date
else task_instance.end_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
),
"duration": (
None if not task_instance.duration else str(task_instance.duration)
),
}
def format_dag_run_state(dag_run: DagRun) -> PipelineStatus:
return PipelineStatus(
state=dag_run.get_state(),
runId=dag_run.run_id,
startDate=None
if not dag_run.start_date
else dag_run.start_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z"),
endDate=None
if not dag_run.end_date
else dag_run.end_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z"),
)

10
openmetadata-airflow-apis/src/openmetadata/api/rest_api.py
View File

@ -39,8 +39,8 @@ from openmetadata.operations.test_connection import test_source_connection
from openmetadata.operations.trigger import trigger
from pydantic.error_wrappers import ValidationError
from metadata.generated.schema.entity.services.connections.serviceConnection import (
ServiceConnectionModel,
from metadata.generated.schema.api.services.ingestionPipelines.testServiceConnection import (
TestServiceConnectionRequest,
)
from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
IngestionPipeline,
@ -172,10 +172,8 @@ class REST_API(AppBuilderBaseView):
json_request = request.get_json()
try:
service_connection_model = ServiceConnectionModel(**json_request)
response = test_source_connection(
service_connection_model.serviceConnection.__root__.config
)
test_service_connection = TestServiceConnectionRequest(**json_request)
response = test_source_connection(test_service_connection)
return response

16
openmetadata-airflow-apis/src/openmetadata/operations/status.py
View File

@ -11,6 +11,7 @@
"""
Module containing the logic to check a DAG status
"""
import json
from typing import Optional
from airflow import settings
@ -18,6 +19,10 @@ from airflow.models import DagRun
from flask import Response
from openmetadata.api.response import ApiResponse, ResponseFormat
from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
PipelineStatus,
)
def status(dag_id: str, run_id: Optional[str]) -> Response:
@ -34,9 +39,9 @@ def status(dag_id: str, run_id: Optional[str]) -> Response:
if dag_run is None:
return ApiResponse.not_found(f"DAG run {run_id} not found")
res_dag_run = ResponseFormat.format_dag_run_state(dag_run)
res_dag_run: PipelineStatus = ResponseFormat.format_dag_run_state(dag_run)
return ApiResponse.success({"message": f"{res_dag_run}"})
return ApiResponse.success(json.loads(res_dag_run.json()))
runs = (
query.filter(
@ -47,6 +52,9 @@ def status(dag_id: str, run_id: Optional[str]) -> Response:
.all()
)
formatted = [ResponseFormat.format_dag_run_state(dag_run) for dag_run in runs]
formatted = [
json.loads(ResponseFormat.format_dag_run_state(dag_run).json())
for dag_run in runs
]
return ApiResponse.success({"message": f"{formatted}"})
return ApiResponse.success(formatted)

9
openmetadata-airflow-apis/src/openmetadata/operations/test_connection.py
View File

@ -15,6 +15,9 @@ from a WorkflowSource
from flask import Response
from openmetadata.api.response import ApiResponse
from metadata.generated.schema.api.services.ingestionPipelines.testServiceConnection import (
TestServiceConnectionRequest,
)
from metadata.utils.engines import (
SourceConnectionException,
get_engine,
@ -22,13 +25,15 @@ from metadata.utils.engines import (
)
def test_source_connection(connection) -> Response:
def test_source_connection(
test_service_connection: TestServiceConnectionRequest,
) -> Response:
"""
Create the engine and test the connection
:param workflow_source: Source to test
:return: None or exception
"""
engine = get_engine(connection)
engine = get_engine(test_service_connection.connection.config)
try:
test_connection(engine)