Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Telescope page #103

Merged
merged 1 commit into from
May 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion astronomer_starship/src/App.jsx
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -118,7 +119,8 @@ export default function App() {
<Route path="pools" element={<PoolsPage key="pools" state={state} dispatch={dispatch} />} />
<Route path="env" element={<EnvVarsPage key="env-vars" state={state} dispatch={dispatch} />} />
<Route path="dags" element={<DAGHistoryPage key="dag-history" state={state} dispatch={dispatch} />} />
</Route>,
<Route path="telescope" element={<TelescopePage key="telescope" state={state} dispatch={dispatch} />} />
</Route>,
),
);
return (
Expand Down
16 changes: 16 additions & 0 deletions astronomer_starship/src/State.jsx
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ export const initialState = {
isTokenTouched: false,
token: null,
deploymentId: null,
telescopeOrganizationId: '',
telescopePresignedUrl: '',

// Software Specific:
releaseName: null,
Expand Down Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions astronomer_starship/src/constants.js
Original file line number Diff line number Diff line change
@@ -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',
Expand Down
10 changes: 10 additions & 0 deletions astronomer_starship/src/pages/SetupPage.jsx
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -88,6 +90,14 @@ export default function SetupPage({ state, dispatch }) {
<HStack>
<Text fontSize="xl">Starship is a utility to migrate Airflow metadata between instances</Text>
<Spacer />
<Button
size="sm"
leftIcon={<IoTelescopeOutline />}
as={NavLink}
to="/telescope"
>
Telescope
</Button>
<Button
size="sm"
leftIcon={<RepeatIcon />}
Expand Down
151 changes: 151 additions & 0 deletions astronomer_starship/src/pages/TelescopePage.jsx
Original file line number Diff line number Diff line change
@@ -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 (
<Box>
<Text fontSize="xl">
Telescope is a tool for gathering metadata from an Airflow instance which can be processed to collect insights.
</Text>
<Divider marginY="5px" />
<VStack width="60%" display="flex" alignItems="center">
<Box width="100%" margin="0 30px" alignItems="left">
<FormControl marginY="2%">
<FormLabel htmlFor="telescopeOrg">
Organization
</FormLabel>
<InputGroup size="sm">
<Input
id="telescopeOrg"
placeholder="Astronomer, LLC"
errorBorderColor="red.300"
value={state.telescopeOrganizationId}
onChange={(e) => dispatch({
type: 'set-telescope-org',
telescopeOrganizationId: e.target.value,
})}
/>
</InputGroup>
<FormHelperText>
Organization name
</FormHelperText>
</FormControl>
<FormControl marginY="2%">
<FormLabel htmlFor="telescopePresignedUrl">
Pre-signed URL
</FormLabel>
<InputGroup size="sm">
<Input
id="telescopePresignedUrl"
placeholder="https://storage.googleapis.com/astronomer-telescope/..."
errorBorderColor="red.300"
value={state.telescopePresignedUrl}
onChange={(e) => dispatch({
type: 'set-telescope-presigned-url',
telescopePresignedUrl: e.target.value,
})}
/>
</InputGroup>
<FormHelperText>
(Optional) Enter a pre-signed URL to submit the report,
or contact an Astronomer Representative to receive one.
</FormHelperText>
<FormErrorMessage>Please fill both parts.</FormErrorMessage>
</FormControl>
<Tooltip hasArrow label="Upload the report via pre-signed url to Astronomer">
<Button
leftIcon={
isUploading ? <span /> :
isUploadComplete ? <FaCheck /> :
<GoUpload />
}
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 ? <CircularProgress thickness="20px" size="30px" isIndeterminate /> :
isUploadComplete ? '' :
error ? 'Error!' :
'Upload'
}
</Button>
</Tooltip>
<Tooltip hasArrow label="Download the report">
<Button
leftIcon={<GoDownload />}
marginX="5px"
isDisabled={!state.telescopeOrganizationId}
href={route}
as={Link}
download={filename}
target="_blank"
>
Download
</Button>
</Tooltip>
</Box>
</VStack>
</Box>
);
}
101 changes: 97 additions & 4 deletions astronomer_starship/starship_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
Loading