chore: re-organize query variants code

insights/insights/doctype/insights_query/insights_assisted_query.py

@@ -6,7 +6,14 @@
 
 import frappe
 
-from .utils import Column, InsightsTable, Query, get_columns_with_inferred_types
+from .utils import (
+ Column,
+ InsightsDataSource,
+ InsightsTable,
+ Query,
+ get_columns_with_inferred_types,
+ update_sql,
+)
 
 DEFAULT_JSON = {
  "table": {},
@@ -29,6 +36,9 @@ def validate(self):
  if not frappe.parse_json(self.doc.json):
  self.doc.json = frappe.as_json(DEFAULT_JSON)
 
+ def before_save(self):
+ update_sql(self.doc)
+
  @cached_property
  def query_json(self):
  query = frappe.parse_json(self.doc.json)
@@ -39,9 +49,6 @@ def query_json(self):
  query.orders = (c.get("column") for c in query.orders or [])
  return Query(**query)
 
- def get_columns(self):
- return self.get_columns_from_results(self.doc.retrieve_results())
-
  def get_columns_from_results(self, results):
  if not results:
  return []
@@ -107,5 +114,8 @@ def before_fetch(self):
  return
  raise frappe.ValidationError("Query Store data source is not supported for assisted query")
 
- def after_fetch_results(self, results):
+ def after_fetch(self, results):
  return results
+
+ def fetch_results(self):
+ return InsightsDataSource.get_doc(self.doc.data_source).run_query(self.doc)

insights/insights/doctype/insights_query/insights_legacy_query.py

@@ -11,7 +11,12 @@
 from insights.api import fetch_column_values, get_tables
 
 from ..insights_data_source.sources.query_store import sync_query_store
-from .utils import InsightsTable, get_columns_with_inferred_types
+from .utils import (
+ InsightsDataSource,
+ InsightsTable,
+ get_columns_with_inferred_types,
+ update_sql,
+)
 
 DEFAULT_FILTERS = dumps(
  {
@@ -248,12 +253,12 @@ class InsightsLegacyQueryController(InsightsLegacyQueryValidation):
  def __init__(self, doc):
  self.doc = doc
 
+ def before_save(self):
+ update_sql(self.doc)
+
  def after_reset(self):
  self.doc.filters = DEFAULT_FILTERS
 
- def get_columns(self):
- return self.get_columns_from_results(self.doc.retrieve_results())
-
  def get_columns_from_results(self, results):
  if not results:
  return []
@@ -329,7 +334,7 @@ def before_fetch(self):
  sub_stored_queries = [t.table for t in self.doc.tables if t.table != self.doc.name]
  sync_query_store(sub_stored_queries, force=True)
 
- def after_fetch_results(self, results):
+ def after_fetch(self, results):
  if self.has_cumulative_columns():
  results = self.apply_cumulative_sum(results)
  return results
@@ -346,3 +351,6 @@ def apply_cumulative_sum(self, results):
  results_df[column.label] = results_df[column.label].cumsum()
 
  return [results[0]] + results_df.values.tolist()
+
+ def fetch_results(self):
+ return InsightsDataSource.get_doc(self.doc.data_source).run_query(self.doc)

insights/insights/doctype/insights_query/insights_query.py

@@ -27,15 +27,13 @@
 from .utils import (
  CachedResults,
  InsightsChart,
- InsightsDataSource,
  InsightsSettings,
  InsightsTable,
  InsightsTableColumn,
  Status,
  apply_pivot_transform,
  apply_transpose_transform,
  apply_unpivot_transform,
- format_query,
 )
 
 
@@ -45,7 +43,7 @@ def before_validate(self):
  self.title = self.name.replace("-", " ").replace("QRY", "Query")
 
  def before_save(self):
- self.update_sql_query()
+ self.variant_controller.before_save()
 
  def on_update(self):
  self.create_default_chart()
@@ -61,10 +59,6 @@ def is_saved_as_table(self):
  table_name = frappe.db.exists("Insights Table", {"table": self.name, "is_query_based": 1})
  return bool(table_name)
 
- @property
- def _data_source(self):
- return InsightsDataSource.get_doc(self.data_source)
-
  @property
  def results(self):
  fetch_if_not_cached = self.status == Status.SUCCESS.value
@@ -92,22 +86,18 @@ def reset(self):
  new_query.name = self.name
  new_query.title = self.name.replace("-", " ").replace("QRY", "Query")
  new_query.data_source = self.data_source
+ new_query.is_native_query = self.is_native_query
+ new_query.is_assisted_query = self.is_assisted_query
  new_query_dict = new_query.as_dict(no_default_fields=True)
  self.update(new_query_dict)
  self.status = Status.SUCCESS.value
  CachedResults.set(self.name, [])
  self.after_reset()
 
  def after_reset(self):
- self.variant_controller.after_reset()
-
- def update_sql_query(self):
- query = self._data_source.build_query(self)
- query = format_query(query)
- if self.sql == query:
+ if not hasattr(self.variant_controller, "after_reset"):
  return
- self.sql = query
- self.status = Status.PENDING.value
+ self.variant_controller.after_reset()
 
  def create_default_chart(self):
  if frappe.db.exists("Insights Chart", {"query": self.name}):
@@ -131,7 +121,7 @@ def update_insights_table(self, force=False):
  create_insights_table(query_table)
 
  def get_columns(self):
- return self.variant_controller.get_columns()
+ return self.get_columns_from_results(self.retrieve_results())
 
  def update_query_store(self):
  if not self.is_stored:
@@ -172,8 +162,8 @@ def fetch_results(self):
  self._results = []
  start = time.monotonic()
  try:
- self._results = self._data_source.run_query(self)
- self._results = self.after_fetch_results(self._results)
+ self._results = self.variant_controller.fetch_results()
+ self._results = self.after_fetch(self._results)
  self._results = self.process_results_columns(self._results)
  self.execution_time = flt(time.monotonic() - start, 3)
  self.last_execution = frappe.utils.now()
@@ -200,10 +190,10 @@ def process_results_columns(self, results):
  def get_columns_from_results(self, results):
  return self.variant_controller.get_columns_from_results(results)
 
- def after_fetch_results(self, results):
+ def after_fetch(self, results):
  if self.transforms:
  results = self.apply_transforms(results)
- results = self.variant_controller.after_fetch_results(results)
+ results = self.variant_controller.after_fetch(results)
  return results
 
  def apply_transforms(self, results):

insights/insights/doctype/insights_query/insights_query_client.py

@@ -77,10 +77,6 @@ def convert_to_assisted(self):
  self.is_assisted_query = 1
  self.save()
 
- @frappe.whitelist()
- def get_source_schema(self):
- return self._data_source.get_schema()
-
  @frappe.whitelist()
  def get_chart_name(self):
  return InsightsChart.get_name(query=self.name)

insights/insights/doctype/insights_query/insights_raw_query.py

@@ -2,7 +2,7 @@
 # For license information, please see license.txt
 
 
-from .utils import get_columns_with_inferred_types
+from .utils import InsightsDataSource, get_columns_with_inferred_types, update_sql
 
 
 class InsightsRawQueryController:
@@ -12,30 +12,25 @@ def __init__(self, doc):
  def validate(self):
  pass
 
- def after_reset(self):
- self.doc.is_native_query = 1
+ def before_save(self):
+ update_sql(self.doc)
 
- def get_sql(self):
- return self.doc.sql
-
- def get_columns(self, results=None):
- if not results:
- results = self.doc.retrieve_results()
+ def get_columns_from_results(self, results):
  if not results:
  return []
  return get_columns_with_inferred_types(results)
 
- def get_columns_from_results(self, results):
- return self.get_columns(results)
-
  def before_fetch(self):
  pass
 
- def after_fetch_results(self, results):
+ def after_fetch(self, results):
  return results
 
  def get_tables_columns(self):
  return []
 
  def get_selected_tables(self):
  return []
+
+ def fetch_results(self):
+ return InsightsDataSource.get_doc(self.doc.data_source).run_query(self.doc)

insights/insights/doctype/insights_query/utils.py

@@ -86,6 +86,15 @@ class Status(Enum):
  FAILED = "Pending Execution"
 
 
+def update_sql(query):
+ sql = InsightsDataSource.get_doc(query.data_source)
+ sql = format_query(sql)
+ if query.sql == sql:
+ return
+ query.sql = sql
+ query.status = Status.PENDING.value
+
+
 def format_query(query):
  return (
  sqlparse.format(