From 2c007541e2999827648ab45e8298afa9439f97e1 Mon Sep 17 00:00:00 2001
From: fritz-astronomer <80706212+fritz-astronomer@users.noreply.github.com>
Date: Thu, 23 May 2024 16:47:43 -0400
Subject: [PATCH] add Telescope page
---
astronomer_starship/src/App.jsx | 4 +-
astronomer_starship/src/State.jsx | 16 ++
astronomer_starship/src/constants.js | 1 +
astronomer_starship/src/pages/SetupPage.jsx | 10 ++
.../src/pages/TelescopePage.jsx | 151 ++++++++++++++++++
astronomer_starship/starship_api.py | 101 +++++++++++-
6 files changed, 278 insertions(+), 5 deletions(-)
create mode 100644 astronomer_starship/src/pages/TelescopePage.jsx
diff --git a/astronomer_starship/src/App.jsx b/astronomer_starship/src/App.jsx
index 6ae7a51..580e644 100644
--- a/astronomer_starship/src/App.jsx
+++ b/astronomer_starship/src/App.jsx
@@ -19,6 +19,7 @@ import {
} from './State';
import './index.css';
import AppLoading from './component/AppLoading';
+import TelescopePage from './pages/TelescopePage';
export default function App() {
const [state, dispatch] = useReducer(reducer, initialState, getInitialState);
@@ -118,7 +119,8 @@ export default function App() {
} />
} />
} />
- ,
+ } />
+ ,
),
);
return (
diff --git a/astronomer_starship/src/State.jsx b/astronomer_starship/src/State.jsx
index 127ffd8..5a5014e 100644
--- a/astronomer_starship/src/State.jsx
+++ b/astronomer_starship/src/State.jsx
@@ -32,6 +32,8 @@ export const initialState = {
isTokenTouched: false,
token: null,
deploymentId: null,
+ telescopeOrganizationId: '',
+ telescopePresignedUrl: '',
// Software Specific:
releaseName: null,
@@ -133,6 +135,20 @@ export const reducer = (state, action) => {
};
}
+ // ### Telescope ###
+ case 'set-telescope-org': {
+ return {
+ ...state,
+ telescopeOrganizationId: action.telescopeOrganizationId,
+ };
+ }
+ case 'set-telescope-presigned-url': {
+ return {
+ ...state,
+ telescopePresignedUrl: action.telescopePresignedUrl,
+ };
+ }
+
// ### VARIABLES PAGE ####
case 'set-variables-loading': {
return {
diff --git a/astronomer_starship/src/constants.js b/astronomer_starship/src/constants.js
index 68746f1..873cba8 100644
--- a/astronomer_starship/src/constants.js
+++ b/astronomer_starship/src/constants.js
@@ -1,4 +1,5 @@
const constants = {
+ TELESCOPE_ROUTE: '/api/starship/telescope',
ENV_VAR_ROUTE: '/api/starship/env_vars',
POOL_ROUTE: '/api/starship/pools',
CONNECTIONS_ROUTE: '/api/starship/connections',
diff --git a/astronomer_starship/src/pages/SetupPage.jsx b/astronomer_starship/src/pages/SetupPage.jsx
index 945b279..05834f4 100644
--- a/astronomer_starship/src/pages/SetupPage.jsx
+++ b/astronomer_starship/src/pages/SetupPage.jsx
@@ -24,6 +24,8 @@ import PropTypes from 'prop-types';
import {
CheckIcon, ExternalLinkIcon, RepeatIcon,
} from '@chakra-ui/icons';
+import { IoTelescopeOutline } from 'react-icons/io5';
+import { NavLink } from 'react-router-dom';
import { getHoustonRoute, getTargetUrlFromParts, proxyHeaders, proxyUrl, tokenUrlFromAirflowUrl } from '../util';
import ValidatedUrlCheckbox from '../component/ValidatedUrlCheckbox';
import axios from "axios";
@@ -88,6 +90,14 @@ export default function SetupPage({ state, dispatch }) {
Starship is a utility to migrate Airflow metadata between instances
+ }
+ as={NavLink}
+ to="/telescope"
+ >
+ Telescope
+
}
diff --git a/astronomer_starship/src/pages/TelescopePage.jsx b/astronomer_starship/src/pages/TelescopePage.jsx
new file mode 100644
index 0000000..401e22f
--- /dev/null
+++ b/astronomer_starship/src/pages/TelescopePage.jsx
@@ -0,0 +1,151 @@
+import {
+ Box, Button,
+ Divider, FormControl,
+ FormErrorMessage,
+ FormHelperText,
+ FormLabel, Link,
+ Input, InputGroup,
+ Text, Tooltip, VStack, CircularProgress, useToast,
+} from '@chakra-ui/react';
+import React, { useEffect, useMemo, useState } from 'react';
+import { GoDownload, GoUpload } from 'react-icons/go';
+import axios from "axios";
+import constants from "../constants.js";
+import { localRoute } from "../util.js";
+import { FaCheck } from "react-icons/fa";
+
+
+export default function TelescopePage({ state, dispatch }) {
+ const [isUploading, setIsUploading] = useState(false);
+ const [isUploadComplete, setIsUploadComplete] = useState(false);
+ const toast = useToast();
+ const [route, setRoute] = useState('');
+ const [filename, setFilename] = useState('');
+ const [error, setError] = useState(null);
+ useEffect(() => {
+ const _route = localRoute(
+ constants.TELESCOPE_ROUTE +
+ (
+ state.telescopeOrganizationId ? (
+ `?organization=${state.telescopeOrganizationId}` +
+ (state.telescopePresignedUrl ? `&presigned_url=${encodeURIComponent(state.telescopePresignedUrl)}` : '')
+ ) : ''
+ )
+ )
+ setRoute(_route);
+ const _filename = `${state.telescopeOrganizationId}.${(new Date()).toISOString().slice(0,10)}.data.json`
+ setFilename(_filename);
+ console.log(_route, _filename);
+ }, [state]);
+ return (
+
+
+ Telescope is a tool for gathering metadata from an Airflow instance which can be processed to collect insights.
+
+
+
+
+
+
+ Organization
+
+
+ dispatch({
+ type: 'set-telescope-org',
+ telescopeOrganizationId: e.target.value,
+ })}
+ />
+
+
+ Organization name
+
+
+
+
+ Pre-signed URL
+
+
+ dispatch({
+ type: 'set-telescope-presigned-url',
+ telescopePresignedUrl: e.target.value,
+ })}
+ />
+
+
+ (Optional) Enter a pre-signed URL to submit the report,
+ or contact an Astronomer Representative to receive one.
+
+ Please fill both parts.
+
+
+ :
+ isUploadComplete ? :
+
+ }
+ marginX="5px"
+ isDisabled={
+ isUploading || isUploadComplete || !state.telescopePresignedUrl || !state.telescopeOrganizationId
+ }
+ colorScheme={
+ isUploadComplete ? 'green' :
+ error ? 'red' :
+ 'gray'
+ }
+ onClick={() => {
+ setIsUploading(true);
+ const errFn = (err) => {
+ toast({
+ title: err.response?.data?.error || err.response?.data || err.message,
+ status: 'error',
+ isClosable: true,
+ });
+ setError(err);
+ };
+ axios.get(route)
+ .then((res) => {
+ console.log(res.data);
+ setIsUploadComplete(true);
+ })
+ .catch(errFn)
+ .finally(() => {
+ setIsUploading(false);
+ });
+ }}
+ >
+ {isUploading ? :
+ isUploadComplete ? '' :
+ error ? 'Error!' :
+ 'Upload'
+ }
+
+
+
+ }
+ marginX="5px"
+ isDisabled={!state.telescopeOrganizationId}
+ href={route}
+ as={Link}
+ download={filename}
+ target="_blank"
+ >
+ Download
+
+
+
+
+
+ );
+}
diff --git a/astronomer_starship/starship_api.py b/astronomer_starship/starship_api.py
index 4de6dcd..dffb7c1 100644
--- a/astronomer_starship/starship_api.py
+++ b/astronomer_starship/starship_api.py
@@ -2,20 +2,66 @@
from functools import partial
import flask
+import requests
from airflow.plugins_manager import AirflowPlugin
from airflow.www.app import csrf
from flask import Blueprint, request, jsonify
from flask_appbuilder import expose, BaseView
+import os
+from typing import Any, Dict, List, Union
+import base64
+import logging
+from json import JSONDecodeError
+
+from astronomer_starship.compat.starship_compatability import (
+ StarshipCompatabilityLayer,
+ get_kwargs_fn,
+)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Callable
-from astronomer_starship.compat.starship_compatability import (
- StarshipCompatabilityLayer,
- get_kwargs_fn,
-)
+
+def get_json_or_clean_str(o: str) -> Union[List[Any], Dict[Any, Any], Any]:
+ """For Aeroscope - Either load JSON (if we can) or strip and split the string, while logging the error"""
+ try:
+ return json.loads(o)
+ except (JSONDecodeError, TypeError) as e:
+ logging.debug(e)
+ logging.debug(o)
+ return o.strip()
+
+
+def clean_airflow_report_output(log_string: str) -> Union[dict, str]:
+ r"""For Aeroscope - Look for the magic string from the Airflow report and then decode the base64 and convert to json
+ Or return output as a list, trimmed and split on newlines
+ >>> clean_airflow_report_output('INFO 123 - xyz - abc\n\n\nERROR - 1234\n%%%%%%%\naGVsbG8gd29ybGQ=')
+ 'hello world'
+ >>> clean_airflow_report_output(
+ ... 'INFO 123 - xyz - abc\n\n\nERROR - 1234\n%%%%%%%\neyJvdXRwdXQiOiAiaGVsbG8gd29ybGQifQ=='
+ ... )
+ {'output': 'hello world'}
+ """
+
+ log_lines = log_string.split("\n")
+ enumerated_log_lines = list(enumerate(log_lines))
+ found_i = -1
+ for i, line in enumerated_log_lines:
+ if "%%%%%%%" in line:
+ found_i = i + 1
+ break
+ if found_i != -1:
+ output = base64.decodebytes(
+ "\n".join(log_lines[found_i:]).encode("utf-8")
+ ).decode("utf-8")
+ try:
+ return json.loads(output)
+ except JSONDecodeError:
+ return get_json_or_clean_str(output)
+ else:
+ return get_json_or_clean_str(log_string)
def starship_route(
@@ -126,6 +172,53 @@ def ok():
return starship_route(get=ok)
+ @expose("/telescope", methods=["GET"])
+ @csrf.exempt
+ def telescope(self):
+ from socket import gethostname
+ import io
+ import runpy
+ from urllib.request import urlretrieve
+ from contextlib import redirect_stdout, redirect_stderr
+ from urllib.error import HTTPError
+ from datetime import datetime, timezone
+
+ aero_version = os.getenv("TELESCOPE_REPORT_RELEASE_VERSION", "latest")
+ a = "airflow_report.pyz"
+ aero_url = (
+ "https://github.com/astronomer/telescope/releases/latest/download/airflow_report.pyz"
+ if aero_version == "latest"
+ else f"https://github.com/astronomer/telescope/releases/download/{aero_version}/airflow_report.pyz"
+ )
+ try:
+ urlretrieve(aero_url, a)
+ except HTTPError as e:
+ raise RuntimeError(
+ f"Error finding specified version:{aero_version} -- Reason:{e.reason}"
+ )
+
+ s = io.StringIO()
+ with redirect_stdout(s), redirect_stderr(s):
+ runpy.run_path(a)
+ report = {
+ "telescope_version": "aeroscope-latest",
+ "report_date": datetime.now(timezone.utc).isoformat()[:10],
+ "organization_name": request.args["organization"],
+ "local": {
+ gethostname(): {
+ "airflow_report": clean_airflow_report_output(s.getvalue())
+ }
+ },
+ }
+ presigned_url = request.args.get("presigned_url", False)
+ if presigned_url:
+ try:
+ upload = requests.put(presigned_url, data=json.dumps(report))
+ return upload.content, upload.status_code
+ except requests.exceptions.ConnectionError as e:
+ return str(e), 400
+ return report
+
@expose("/airflow_version", methods=["GET"])
@csrf.exempt
def airflow_version(self) -> str: