Merge pull request #103 from astronomer/telescope

Add Telescope page

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() {
         <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 (

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 {

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',

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 }) {
       <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 />}

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 (
+    <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>
+  );
+}

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: