From 389f6602bcb4e988366e105725f5d616bbdb2a69 Mon Sep 17 00:00:00 2001
From: "Daniele.Mazzotta" <daniele.mazzotta@tii.ae>
Date: Fri, 10 Jun 2022 17:38:19 +0400
Subject: [PATCH] added new getDags method to fetch specific dagIds

Signed-off-by: Daniele.Mazzotta <daniele.mazzotta@tii.ae>
---
 .../src/api/ApacheAirflowApi.ts               |  1 +
 .../src/api/ApacheAirflowClient.test.ts       | 27 +++++++++++++++++++
 .../src/api/ApacheAirflowClient.ts            | 19 +++++++++++++
 .../DagTableComponent/DagTableComponent.tsx   | 17 +++++++++++-
 4 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/plugins/apache-airflow/src/api/ApacheAirflowApi.ts b/plugins/apache-airflow/src/api/ApacheAirflowApi.ts
index dde8c57802..838292524a 100644
--- a/plugins/apache-airflow/src/api/ApacheAirflowApi.ts
+++ b/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>;
diff --git a/plugins/apache-airflow/src/api/ApacheAirflowClient.test.ts b/plugins/apache-airflow/src/api/ApacheAirflowClient.test.ts
index 704b17eb34..e56c622ff2 100644
--- a/plugins/apache-airflow/src/api/ApacheAirflowClient.test.ts
+++ b/plugins/apache-airflow/src/api/ApacheAirflowClient.test.ts
@@ -160,4 +160,31 @@ 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[0].dag_id).toEqual('mock_dag_1');
+    expect(response.dagsNotFound[0]).toEqual('a-random-DAG-id');
+  });
 });
diff --git a/plugins/apache-airflow/src/api/ApacheAirflowClient.ts b/plugins/apache-airflow/src/api/ApacheAirflowClient.ts
index 847a1b1ca6..bf095a96ba 100644
--- a/plugins/apache-airflow/src/api/ApacheAirflowClient.ts
+++ b/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 === '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',
diff --git a/plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.tsx b/plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.tsx
index cf343a596e..18439e1902 100644
--- a/plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.tsx
+++ b/plugins/apache-airflow/src/components/DagTableComponent/DagTableComponent.tsx
@@ -129,9 +129,24 @@ 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, dagsNotFound } = await apiClient.getDags(dagIds);
+      if (dagsNotFound.length) {
+        throw new Error(
+          `${dagsNotFound.length} DAGs were not found:\n${dagsNotFound.join(
+            ';\n',
+          )}`,
+        );
+      }
+      return dags;
+    }
     return await apiClient.listDags();
   }, []);