Return cached data source schema when available (#4934)

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

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)

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