Merge pull request #87 from astronomer/dev

v2.0.1

astronomer_starship/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.0.0"
+__version__ = "2.0.1"
 
 
 def get_provider_info():

astronomer_starship/compat/starship_compatability.py

@@ -153,7 +153,7 @@ def get_test_data(attrs: dict, method: "Union[str, None]" = None) -> "Dict[str,
     >>> get_test_data(method="PATCH", attrs=StarshipAirflow.dag_attrs())
     {'dag_id': 'dag_0', 'is_paused': False}
     >>> get_test_data(attrs=StarshipAirflow.dag_attrs()) # doctest: +ELLIPSIS
-    {'dag_id': 'dag_0', 'schedule_interval': '@once', 'is_paused': False, ... 'task_count': 0}
+    {'dag_id': 'dag_0', 'schedule_interval': '@once', 'is_paused': False, ... 'dag_run_count': 0}
     """
 
     if method:
@@ -335,7 +335,7 @@ def dag_attrs(cls) -> "Dict[str, AttrDesc]":
                 "test_value": "baz",
             },
             "tags": {
-                "attr": "tags",
+                "attr": None,
                 "methods": [],
                 "test_value": ["bar", "foo"],
             },
@@ -344,51 +344,35 @@ def dag_attrs(cls) -> "Dict[str, AttrDesc]":
                 "methods": [],
                 "test_value": 0,
             },
-            "task_count": {
-                "attr": None,
-                "methods": [],
-                "test_value": 0,
-            },
         }
 
     def get_dags(self):
         """Get all DAGs"""
-        from airflow.models import DagModel, TaskInstance, DagRun
-        from sqlalchemy.sql.functions import count
-        from sqlalchemy import distinct
+        from airflow.models import DagModel
 
         try:
+            fields = [
+                getattr(DagModel, attr_desc["attr"])
+                for attr_desc in self.dag_attrs().values()
+                if attr_desc["attr"] is not None
+            ]
             # py36/sqlalchemy1.3 doesn't like label?
             # noinspection PyUnresolvedReferences
-            results = (
-                self.session.query(
-                    DagModel,
-                    count(distinct(TaskInstance.task_id)).label("task_count"),
-                    count(distinct(DagRun.run_id)).label("dag_run_count"),
-                )
-                .join(
-                    TaskInstance, TaskInstance.dag_id == DagModel.dag_id, isouter=True
-                )
-                .join(DagRun, DagRun.dag_id == DagModel.dag_id, isouter=True)
-                .group_by(DagModel)
-                .all()
-            )
             return json.loads(
                 json.dumps(
                     [
                         {
                             attr: (
-                                # e.g. result.DagModel.dag_id
-                                getattr(
-                                    getattr(result, "DagModel"), attr_desc["attr"], None
-                                )
-                                if attr_desc["attr"] is not None
-                                else getattr(result, attr)
-                                # e.g. result.task_count
+                                self._get_tags(result.dag_id)
+                                if attr == "tags"
+                                else self._get_dag_run_count(result.dag_id)
+                                if attr == "dag_run_count"
+                                else getattr(result, attr_desc["attr"], None)
+                                # e.g. result.dag_id
                             )
                             for attr, attr_desc in self.dag_attrs().items()
                         }
-                        for result in results
+                        for result in self.session.query(*fields).all()
                     ],
                     default=str,
                 )
@@ -417,6 +401,23 @@ def set_dag_is_paused(self, dag_id: str, is_paused: bool):
             self.session.rollback()
             raise e
 
+    def _get_tags(self, dag_id: str):
+        try:
+            from airflow.models import DagTag
+
+            # noinspection PyTypeChecker
+            return [
+                tag[0]
+                for tag in self.session.query(DagTag.name)
+                .filter(DagTag.dag_id == dag_id)
+                .all()
+            ]
+        except ImportError:
+            return []
+        except Exception as e:
+            self.session.rollback()
+            raise e
+
     def _get_dag_run_count(self, dag_id: str):
         from airflow.models import DagRun
         from sqlalchemy.sql.functions import count

astronomer_starship/src/index.jsx

@@ -1,17 +1,11 @@
 import React from 'react';
 import ReactDOM from 'react-dom/client';
-import { ChakraProvider, extendTheme } from '@chakra-ui/react';
+import { ChakraProvider } from '@chakra-ui/react';
 import App from './App';
 
-const theme = extendTheme({
-  colors: {},
-  initialColorMode: 'system',
-  useSystemColorMode: true,
-});
-
 ReactDOM.createRoot(document.getElementById('root')).render(
   <React.StrictMode>
-    <ChakraProvider theme={theme}>
+    <ChakraProvider>
       <App />
     </ChakraProvider>
   </React.StrictMode>,

astronomer_starship/src/pages/SetupPage.jsx

@@ -24,24 +24,15 @@ import PropTypes from 'prop-types';
 import {
   CheckIcon, ExternalLinkIcon, RepeatIcon,
 } from '@chakra-ui/icons';
-import { MdNightlight } from 'react-icons/md';
-import { NavLink } from 'react-router-dom';
-import { IoTelescopeOutline } from 'react-icons/io5';
 import { getTargetUrlFromParts, tokenUrlFromAirflowUrl } from '../util';
 import ValidatedUrlCheckbox from '../component/ValidatedUrlCheckbox';
 
 export default function SetupPage({ state, dispatch }) {
-  const { colorMode, toggleColorMode } = useColorMode();
   return (
     <Box>
       <HStack>
         <Text fontSize="xl">Starship is a utility to migrate Airflow metadata between instances</Text>
         <Spacer />
-        <Button size="sm" leftIcon={<MdNightlight />} onClick={toggleColorMode}>
-          {colorMode === 'light' ? 'Dark' : 'Light'}
-          {' '}
-          Mode
-        </Button>
         <Button
           size="sm"
           leftIcon={<RepeatIcon />}

astronomer_starship/starship_api.py

@@ -352,7 +352,6 @@ def dags(self):
                 "owners": "user",
                 "tags": ["tag1", "tag2"],
                 "dag_run_count": 2,
-                "task_count": 3
             },
             ...
         ]

justfile

@@ -169,6 +169,12 @@ deploy-test-no-symlink TESTPATH:
 serve-frontend: build
     cd astronomer_starship && npx vite
 
+# Restart just the webserver container (e.g. to reload the plugin)
+restart-webserver:
+    docker restart $(docker container ls --filter name=webserver --format="{{{{.ID}}") \
+      && docker logs -f $(docker container ls --filter name=webserver --format="{{{{.ID}}") --since 1m
+
+
 # Update the baseline for detect-secrets / pre-commit
 update-secrets:
     detect-secrets scan  > .secrets.baseline  # pragma: allowlist secret

tests/api_integration_test.py

@@ -198,4 +198,10 @@ def test_integration_dag_runs_and_task_instances(url_and_token_and_starship):
     }
     # gets blanked out
     test_input["task_instances"][0]["executor_config"] = None
+
+    if "trigger_timeout" in actual_task_instance:
+        del actual_task_instance["trigger_timeout"]
+    if "trigger_timeout" in test_input["task_instances"][0]:
+        del test_input["task_instances"][0]["trigger_timeout"]
+
     assert actual_task_instance == test_input["task_instances"][0], actual_task_instance