From 100c7be5e0a23f5497e24a6f24630a086bc638a8 Mon Sep 17 00:00:00 2001
From: Gabriel Dutra <nesk.frz@gmail.com>
Date: Tue, 2 Jun 2020 05:26:53 -0300
Subject: [PATCH] Return cached data source schema when available (#4934)

---
 .../queries/hooks/useDataSourceSchema.js      | 41 +++++++++++--------
 redash/handlers/data_sources.py               |  6 +++
 redash/models/__init__.py                     | 16 ++++----
 3 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/client/app/pages/queries/hooks/useDataSourceSchema.js b/client/app/pages/queries/hooks/useDataSourceSchema.js
index 2c1b0ea4..4b221943 100644
--- a/client/app/pages/queries/hooks/useDataSourceSchema.js
+++ b/client/app/pages/queries/hooks/useDataSourceSchema.js
@@ -1,11 +1,11 @@
-import { reduce } from "lodash";
+import { reduce, has } from "lodash";
 import { useCallback, useEffect, useRef, useState, useMemo } from "react";
 import { axios } from "@/services/axios";
 import DataSource, { SCHEMA_NOT_SUPPORTED } from "@/services/data-source";
 import notification from "@/services/notification";
 
 function sleep(ms) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
+  return new Promise(resolve => setTimeout(resolve, ms));
 }
 
 function getSchema(dataSource, refresh = undefined) {
@@ -13,24 +13,27 @@ function getSchema(dataSource, refresh = undefined) {
     return Promise.resolve([]);
   }
 
-  const fetchSchemaFromJob = (data) => {
-    return sleep(1000).then(() => {
-      return axios.get(`api/jobs/${data.job.id}`).then((data) => {
-        if (data.job.status < 3) {
-          return fetchSchemaFromJob(data);
-        } else if (data.job.status === 3) {
-          return data.job.result;
-        } else if (data.job.status === 4 && data.job.error.code === SCHEMA_NOT_SUPPORTED) {
-          return [];
-        } else {
-          return Promise.reject(new Error(data.job.error));
-        }
-      });
+  const fetchSchemaFromJob = data => {
+    return axios.get(`api/jobs/${data.job.id}`).then(data => {
+      if (data.job.status < 3) {
+        return sleep(1000).then(() => fetchSchemaFromJob(data));
+      } else if (data.job.status === 3) {
+        return data.job.result;
+      } else if (data.job.status === 4 && data.job.error.code === SCHEMA_NOT_SUPPORTED) {
+        return [];
+      } else {
+        return Promise.reject(new Error(data.job.error));
+      }
     });
   };
 
   return DataSource.fetchSchema(dataSource, refresh)
-    .then(fetchSchemaFromJob)
+    .then(data => {
+      if (has(data, "job")) {
+        return fetchSchemaFromJob(data);
+      }
+      return has(data, "schema") ? data.schema : Promise.reject();
+    })
     .catch(() => {
       notification.error("Schema refresh failed.", "Please try again later.");
       return Promise.resolve([]);
@@ -48,9 +51,11 @@ export default function useDataSourceSchema(dataSource) {
 
   const reloadSchema = useCallback(
     (refresh = undefined) => {
-      const refreshToken = Math.random().toString(36).substr(2);
+      const refreshToken = Math.random()
+        .toString(36)
+        .substr(2);
       refreshSchemaTokenRef.current = refreshToken;
-      getSchema(dataSource, refresh).then((data) => {
+      getSchema(dataSource, refresh).then(data => {
         if (refreshSchemaTokenRef.current === refreshToken) {
           setSchema(prepareSchema(data));
         }
diff --git a/redash/handlers/data_sources.py b/redash/handlers/data_sources.py
index b60c204b..b3521ca0 100644
--- a/redash/handlers/data_sources.py
+++ b/redash/handlers/data_sources.py
@@ -187,6 +187,12 @@ class DataSourceSchemaResource(BaseResource):
         require_access(data_source, self.current_user, view_only)
         refresh = request.args.get("refresh") is not None
 
+        if not refresh:
+            cached_schema = data_source.get_cached_schema()
+
+            if cached_schema is not None:
+                return {"schema": cached_schema}
+
         job = get_schema.delay(data_source.id, refresh)
 
         return serialize_job(job)
diff --git a/redash/models/__init__.py b/redash/models/__init__.py
index 5d1a0e49..a7f74270 100644
--- a/redash/models/__init__.py
+++ b/redash/models/__init__.py
@@ -182,12 +182,16 @@ class DataSource(BelongsToOrgMixin, db.Model):
 
         return res
 
-    def get_schema(self, refresh=False):
-        cache = None
-        if not refresh:
-            cache = redis_connection.get(self._schema_key)
+    def get_cached_schema(self):
+        cache = redis_connection.get(self._schema_key)
+        return json_loads(cache) if cache else None
 
-        if cache is None:
+    def get_schema(self, refresh=False):
+        out_schema = None
+        if not refresh:
+            out_schema = self.get_cached_schema()
+
+        if out_schema is None:
             query_runner = self.query_runner
             schema = query_runner.get_schema(get_stats=refresh)
 
@@ -200,8 +204,6 @@ class DataSource(BelongsToOrgMixin, db.Model):
                 out_schema = schema
             finally:
                 redis_connection.set(self._schema_key, json_dumps(out_schema))
-        else:
-            out_schema = json_loads(cache)
 
         return out_schema