Merge pull request #12113 from stichiboi/stichiboi/apache-airflow-plugin-extension

Stichiboi/apache airflow plugin extension

.changeset/lemon-goats-obey.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-apache-airflow': minor
+---
+
+Exposed DagTableComponent as standalone component + added a prop to get only select DAGs instead of the full list

plugins/apache-airflow/README.md

@@ -42,6 +42,26 @@ yarn --cwd packages/app add @backstage/plugin-apache-airflow
  );
 ```
 
+If you just want to embed the DAGs into an existing page, you can use the `ApacheAirflowDagTable`
+
+```tsx
+import { ApacheAirflowDagTable } from '@backstage/plugin-apache-airflow';
+
+export function SomeEntityPage(): JSX.Element {
+  return (
+    <Grid item md={6} xs={12}>
+      <ApacheAirflowDagTable
+        dagIds={[
+          'example_bash_operator',
+          'example_branch_datetime_operator_2',
+          'example_branch_labels',
+        ]}
+      />
+    </Grid>
+  );
+}
+```
+
 ## Configuration
 
 For links to the Airflow instance, the `baseUrl` must be defined in

plugins/apache-airflow/api-report.md

@@ -8,6 +8,13 @@
 import { BackstagePlugin } from '@backstage/core-plugin-api';
 import { RouteRef } from '@backstage/core-plugin-api';
 
+// @public
+export const ApacheAirflowDagTable: ({
+                                       dagIds,
+                                     }: {
+  dagIds?: string[] | undefined;
+}) => JSX.Element;
+
 // Warning: (ae-missing-release-tag) "ApacheAirflowPage" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
 //
 // @public (undocumented)

plugins/apache-airflow/src/api/ApacheAirflowApi.ts

@@ -25,6 +25,7 @@ export type ApacheAirflowApi = {
   discoveryApi: DiscoveryApi;
   baseUrl: string;
   listDags(options?: { objectsPerRequest: number }): Promise<Dag[]>;
+  getDags(dagIds: string[]): Promise<{ dags: Dag[]; dagsNotFound: string[] }>;
   updateDag(dagId: string, isPaused: boolean): Promise<any>;
   getInstanceStatus(): Promise<InstanceStatus>;
   getInstanceVersion(): Promise<InstanceVersion>;

plugins/apache-airflow/src/api/ApacheAirflowClient.test.ts

@@ -107,6 +107,15 @@ describe('ApacheAirflowClient', () => {
         );
       }),
 
+      rest.get(`${mockBaseUrl}/airflow/dags/:dag_id`, (req, res, ctx) => {
+        const { dag_id } = req.params;
+        const dag = dags.find(d => d.dag_id === dag_id);
+        if (dag) {
+          return res(ctx.json(dag));
+        }
+        return res(ctx.status(404));
+      }),
+
       rest.patch(`${mockBaseUrl}/airflow/dags/:dag_id`, (req, res, ctx) => {
         const { dag_id } = req.params;
         const body = JSON.parse(req.body as string);
@@ -160,4 +169,33 @@ describe('ApacheAirflowClient', () => {
     expect(response.dag_id).toEqual(dagId);
     expect(response.is_paused).toEqual(true);
   });
+
+  it('get only some dags', async () => {
+    setupHandlers();
+    const client = new ApacheAirflowClient({
+      discoveryApi: discoveryApi,
+      baseUrl: 'localhost:8080/',
+    });
+    const dagIds = ['mock_dag_1', 'mock_dag_3'];
+    const response = await client.getDags(dagIds);
+    expect(response.dags.length).toEqual(dagIds.length);
+    response.dags.forEach((dag, index) =>
+      expect(dag.dag_id).toEqual(dagIds[index]),
+    );
+    expect(response.dagsNotFound.length).toEqual(0);
+  });
+
+  it('get dags but ignore NOT FOUND errors', async () => {
+    setupHandlers();
+    const client = new ApacheAirflowClient({
+      discoveryApi: discoveryApi,
+      baseUrl: 'localhost:8080/',
+    });
+    const dagIds = ['mock_dag_1', 'a-random-DAG-id'];
+    const response = await client.getDags(dagIds);
+    expect(response.dags.length).toEqual(1);
+    expect(response.dags[0].dag_id).toEqual('mock_dag_1');
+    expect(response.dagsNotFound.length).toEqual(1);
+    expect(response.dagsNotFound[0]).toEqual('a-random-DAG-id');
+  });
 });

plugins/apache-airflow/src/api/ApacheAirflowClient.ts

@@ -75,6 +75,25 @@ export class ApacheAirflowClient implements ApacheAirflowApi {
     return dags;
   }
 
+  async getDags(
+    dagIds: string[],
+  ): Promise<{ dags: Dag[]; dagsNotFound: string[] }> {
+    const dagsNotFound: string[] = [];
+    const response = await Promise.all(
+      dagIds.map(id => {
+        return this.fetch<Dag>(`/dags/${id}`).catch(e => {
+          if (e.message.toUpperCase('en-US') === 'NOT FOUND') {
+            dagsNotFound.push(id);
+          } else {
+            throw e;
+          }
+        });
+      }),
+    );
+    const dags = response.filter(Boolean) as Dag[];
+    return { dags, dagsNotFound };
+  }
+
   async updateDag(dagId: string, isPaused: boolean): Promise<Dag> {
     const init = {
       method: 'PATCH',

plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.test.tsx

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2022 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { ApacheAirflowApi, apacheAirflowApiRef } from '../../api';
+import { renderInTestApp, TestApiProvider } from '@backstage/test-utils';
+import React from 'react';
+import { DagTableComponent } from './DagTableComponent';
+
+describe('DagTableComponent', () => {
+  const mockApi: jest.Mocked<ApacheAirflowApi> = {
+    listDags: jest.fn().mockResolvedValue([
+      {
+        dag_id: 'mock_dag_1',
+      },
+      {
+        dag_id: 'mock_dag_2',
+      },
+      {
+        dag_id: 'mock_dag_3',
+      },
+    ]),
+    getDags: jest.fn().mockResolvedValue({
+      dags: [{ dag_id: 'mock_dag_1' }],
+      dagsNotFound: ['a-random-id'],
+    }),
+  } as any;
+
+  it('should render all DAGs', async () => {
+    const { getByText } = await renderInTestApp(
+      <TestApiProvider apis={[[apacheAirflowApiRef, mockApi]]}>
+        <DagTableComponent />
+      </TestApiProvider>,
+    );
+
+    ['mock_dag_1', 'mock_dag_2', 'mock_dag_3'].forEach(dagId => {
+      expect(getByText(dagId)).toBeInTheDocument();
+    });
+  });
+
+  it('should render only selected DAGs', async () => {
+    const { getByText } = await renderInTestApp(
+      <TestApiProvider apis={[[apacheAirflowApiRef, mockApi]]}>
+        <DagTableComponent dagIds={['mock_dag_1', 'a-random-id']} />
+      </TestApiProvider>,
+    );
+    expect(getByText('mock_dag_1')).toBeInTheDocument();
+    expect(getByText('Warning: 1 DAGs were not found')).toBeInTheDocument();
+    expect(getByText('a-random-id')).toBeInTheDocument();
+  });
+});

plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.tsx

@@ -20,6 +20,7 @@ import {
   StatusOK,
   Table,
   TableColumn,
+  WarningPanel,
 } from '@backstage/core-components';
 import { useApi } from '@backstage/core-plugin-api';
 import Box from '@material-ui/core/Box';
@@ -129,9 +130,18 @@ export const DenseTable = ({ dags }: DenseTableProps) => {
   );
 };
 
-export const DagTableComponent = () => {
+type DagTableComponentProps = {
+  dagIds?: string[];
+};
+
+export const DagTableComponent = ({ dagIds }: DagTableComponentProps) => {
   const apiClient = useApi(apacheAirflowApiRef);
+
   const { value, loading, error } = useAsync(async (): Promise<Dag[]> => {
+    if (dagIds) {
+      const { dags } = await apiClient.getDags(dagIds);
+      return dags;
+    }
     return await apiClient.listDags();
   }, []);
 
@@ -146,6 +156,22 @@ export const DagTableComponent = () => {
     id: el.dag_id, // table records require `id` attribute
     dagUrl: `${apiClient.baseUrl}dag_details?dag_id=${el.dag_id}`, // construct path to DAG using `baseUrl`
   }));
-
-  return <DenseTable dags={data || []} />;
+  const dagsNotFound =
+    dagIds && value
+      ? dagIds.filter(id => !value.find(d => d.dag_id === id))
+      : [];
+  return (
+    <>
+      {dagsNotFound.length ? (
+        <WarningPanel title={`${dagsNotFound.length} DAGs were not found`}>
+          {dagsNotFound.map(dagId => (
+            <Typography key={dagId}>{dagId}</Typography>
+          ))}
+        </WarningPanel>
+      ) : (
+        ''
+      )}
+      <DenseTable dags={data || []} />
+    </>
+  );
 };

plugins/apache-airflow/src/index.ts

@@ -20,4 +20,8 @@
  * @packageDocumentation
  */
 
-export { apacheAirflowPlugin, ApacheAirflowPage } from './plugin';
+export {
+  apacheAirflowPlugin,
+  ApacheAirflowPage,
+  ApacheAirflowDagTable,
+} from './plugin';

plugins/apache-airflow/src/plugin.ts

@@ -17,11 +17,12 @@
 import { rootRouteRef } from './routes';
 import { apacheAirflowApiRef, ApacheAirflowClient } from './api';
 import {
+  configApiRef,
   createApiFactory,
+  createComponentExtension,
   createPlugin,
   createRoutableExtension,
   discoveryApiRef,
-  configApiRef,
 } from '@backstage/core-plugin-api';
 
 export const apacheAirflowPlugin = createPlugin({
@@ -49,3 +50,20 @@ export const ApacheAirflowPage = apacheAirflowPlugin.provide(
     mountPoint: rootRouteRef,
   }),
 );
+
+/**
+ * Render the DAGs in a table
+ * If the dagIds is specified, only those DAGs are loaded.
+ * Otherwise, it's going to list all the DAGs
+ * @public
+ * @param dagIds - optional string[] of the DAGs to show in the table. If undefined, it will list all DAGs
+ */
+export const ApacheAirflowDagTable = apacheAirflowPlugin.provide(
+  createComponentExtension({
+    name: 'ApacheAirflowDagTable',
+    component: {
+      lazy: () =>
+        import('./components/DagTableComponent').then(m => m.DagTableComponent),
+    },
+  }),
+);