mirror of https://github.com/apache/superset.git
feat: Save column data into json_metadata for all Query executions (#20059)
* add save_metadata function to QueryDAO * use set_extra_json_key * added test * Update queries_test.py * fix pylint * add to session * add to session * refactor * forgot the return
This commit is contained in:
parent
1d410eb763
commit
660af409a4
|
@ -16,6 +16,7 @@
|
||||||
# under the License.
|
# under the License.
|
||||||
import logging
|
import logging
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
from typing import Any, Dict
|
||||||
|
|
||||||
from superset.dao.base import BaseDAO
|
from superset.dao.base import BaseDAO
|
||||||
from superset.extensions import db
|
from superset.extensions import db
|
||||||
|
@ -48,3 +49,10 @@ class QueryDAO(BaseDAO):
|
||||||
saved_query.rows = query.rows
|
saved_query.rows = query.rows
|
||||||
saved_query.last_run = datetime.now()
|
saved_query.last_run = datetime.now()
|
||||||
db.session.commit()
|
db.session.commit()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def save_metadata(query: Query, payload: Dict[str, Any]) -> None:
|
||||||
|
# pull relevant data from payload and store in extra_json
|
||||||
|
columns = payload.get("columns", {})
|
||||||
|
db.session.add(query)
|
||||||
|
query.set_extra_json_key("columns", columns)
|
||||||
|
|
|
@ -187,7 +187,7 @@ def get_sql_results( # pylint: disable=too-many-arguments
|
||||||
return handle_query_error(ex, query, session)
|
return handle_query_error(ex, query, session)
|
||||||
|
|
||||||
|
|
||||||
def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-locals,too-many-statements
|
def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statements
|
||||||
sql_statement: str,
|
sql_statement: str,
|
||||||
query: Query,
|
query: Query,
|
||||||
session: Session,
|
session: Session,
|
||||||
|
|
|
@ -34,6 +34,7 @@ from superset.sqllab.exceptions import (
|
||||||
QueryIsForbiddenToAccessException,
|
QueryIsForbiddenToAccessException,
|
||||||
SqlLabException,
|
SqlLabException,
|
||||||
)
|
)
|
||||||
|
from superset.sqllab.execution_context_convertor import ExecutionContextConvertor
|
||||||
from superset.sqllab.limiting_factor import LimitingFactor
|
from superset.sqllab.limiting_factor import LimitingFactor
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
@ -42,6 +43,7 @@ if TYPE_CHECKING:
|
||||||
from superset.sqllab.sql_json_executer import SqlJsonExecutor
|
from superset.sqllab.sql_json_executer import SqlJsonExecutor
|
||||||
from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
|
from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
CommandResult = Dict[str, Any]
|
CommandResult = Dict[str, Any]
|
||||||
|
@ -94,11 +96,19 @@ class ExecuteSqlCommand(BaseCommand):
|
||||||
status = SqlJsonExecutionStatus.QUERY_ALREADY_CREATED
|
status = SqlJsonExecutionStatus.QUERY_ALREADY_CREATED
|
||||||
else:
|
else:
|
||||||
status = self._run_sql_json_exec_from_scratch()
|
status = self._run_sql_json_exec_from_scratch()
|
||||||
|
|
||||||
|
self._execution_context_convertor.set_payload(
|
||||||
|
self._execution_context, status
|
||||||
|
)
|
||||||
|
|
||||||
|
# save columns into metadata_json
|
||||||
|
self._query_dao.save_metadata(
|
||||||
|
self._execution_context.query, self._execution_context_convertor.payload
|
||||||
|
)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": status,
|
"status": status,
|
||||||
"payload": self._execution_context_convertor.to_payload(
|
"payload": self._execution_context_convertor.serialize_payload(),
|
||||||
self._execution_context, status
|
|
||||||
),
|
|
||||||
}
|
}
|
||||||
except (SqlLabException, SupersetErrorsException) as ex:
|
except (SqlLabException, SupersetErrorsException) as ex:
|
||||||
raise ex
|
raise ex
|
||||||
|
@ -209,12 +219,3 @@ class CanAccessQueryValidator:
|
||||||
class SqlQueryRender:
|
class SqlQueryRender:
|
||||||
def render(self, execution_context: SqlJsonExecutionContext) -> str:
|
def render(self, execution_context: SqlJsonExecutionContext) -> str:
|
||||||
raise NotImplementedError()
|
raise NotImplementedError()
|
||||||
|
|
||||||
|
|
||||||
class ExecutionContextConvertor:
|
|
||||||
def to_payload(
|
|
||||||
self,
|
|
||||||
execution_context: SqlJsonExecutionContext,
|
|
||||||
execution_status: SqlJsonExecutionStatus,
|
|
||||||
) -> str:
|
|
||||||
raise NotImplementedError()
|
|
||||||
|
|
|
@ -16,52 +16,53 @@
|
||||||
# under the License.
|
# under the License.
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import TYPE_CHECKING
|
import logging
|
||||||
|
from typing import Any, Dict, TYPE_CHECKING
|
||||||
|
|
||||||
import simplejson as json
|
import simplejson as json
|
||||||
|
|
||||||
import superset.utils.core as utils
|
import superset.utils.core as utils
|
||||||
from superset.sqllab.command import ExecutionContextConvertor
|
|
||||||
from superset.sqllab.command_status import SqlJsonExecutionStatus
|
from superset.sqllab.command_status import SqlJsonExecutionStatus
|
||||||
from superset.sqllab.utils import apply_display_max_row_configuration_if_require
|
from superset.sqllab.utils import apply_display_max_row_configuration_if_require
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from superset.models.sql_lab import Query
|
from superset.models.sql_lab import Query
|
||||||
from superset.sqllab.sql_json_executer import SqlResults
|
from superset.sqllab.sql_json_executer import SqlResults
|
||||||
from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
|
from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
|
||||||
|
|
||||||
|
|
||||||
class ExecutionContextConvertorImpl(ExecutionContextConvertor):
|
class ExecutionContextConvertor:
|
||||||
_max_row_in_display_configuration: int # pylint: disable=invalid-name
|
_max_row_in_display_configuration: int # pylint: disable=invalid-name
|
||||||
|
_exc_status: SqlJsonExecutionStatus
|
||||||
|
payload: Dict[str, Any]
|
||||||
|
|
||||||
def set_max_row_in_display(self, value: int) -> None:
|
def set_max_row_in_display(self, value: int) -> None:
|
||||||
self._max_row_in_display_configuration = value # pylint: disable=invalid-name
|
self._max_row_in_display_configuration = value # pylint: disable=invalid-name
|
||||||
|
|
||||||
def to_payload(
|
def set_payload(
|
||||||
self,
|
self,
|
||||||
execution_context: SqlJsonExecutionContext,
|
execution_context: SqlJsonExecutionContext,
|
||||||
execution_status: SqlJsonExecutionStatus,
|
execution_status: SqlJsonExecutionStatus,
|
||||||
) -> str:
|
) -> None:
|
||||||
|
self._exc_status = execution_status
|
||||||
if execution_status == SqlJsonExecutionStatus.HAS_RESULTS:
|
if execution_status == SqlJsonExecutionStatus.HAS_RESULTS:
|
||||||
return self._to_payload_results_based(
|
self.payload = execution_context.get_execution_result() or {}
|
||||||
execution_context.get_execution_result() or {}
|
else:
|
||||||
)
|
self.payload = execution_context.query.to_dict()
|
||||||
return self._to_payload_query_based(execution_context.query)
|
|
||||||
|
|
||||||
def _to_payload_results_based(self, execution_result: SqlResults) -> str:
|
def serialize_payload(self) -> str:
|
||||||
|
if self._exc_status == SqlJsonExecutionStatus.HAS_RESULTS:
|
||||||
return json.dumps(
|
return json.dumps(
|
||||||
apply_display_max_row_configuration_if_require(
|
apply_display_max_row_configuration_if_require(
|
||||||
execution_result, self._max_row_in_display_configuration
|
self.payload, self._max_row_in_display_configuration
|
||||||
),
|
),
|
||||||
default=utils.pessimistic_json_iso_dttm_ser,
|
default=utils.pessimistic_json_iso_dttm_ser,
|
||||||
ignore_nan=True,
|
ignore_nan=True,
|
||||||
encoding=None,
|
encoding=None,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _to_payload_query_based( # pylint: disable=no-self-use
|
|
||||||
self, query: Query
|
|
||||||
) -> str:
|
|
||||||
return json.dumps(
|
return json.dumps(
|
||||||
{"query": query.to_dict()}, default=utils.json_int_dttm_ser, ignore_nan=True
|
{"query": self.payload}, default=utils.json_int_dttm_ser, ignore_nan=True
|
||||||
)
|
)
|
||||||
|
|
|
@ -113,7 +113,7 @@ from superset.sqllab.exceptions import (
|
||||||
QueryIsForbiddenToAccessException,
|
QueryIsForbiddenToAccessException,
|
||||||
SqlLabException,
|
SqlLabException,
|
||||||
)
|
)
|
||||||
from superset.sqllab.execution_context_convertor import ExecutionContextConvertorImpl
|
from superset.sqllab.execution_context_convertor import ExecutionContextConvertor
|
||||||
from superset.sqllab.limiting_factor import LimitingFactor
|
from superset.sqllab.limiting_factor import LimitingFactor
|
||||||
from superset.sqllab.query_render import SqlQueryRenderImpl
|
from superset.sqllab.query_render import SqlQueryRenderImpl
|
||||||
from superset.sqllab.sql_json_executer import (
|
from superset.sqllab.sql_json_executer import (
|
||||||
|
@ -1331,7 +1331,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
|
||||||
@has_access_api
|
@has_access_api
|
||||||
@event_logger.log_this
|
@event_logger.log_this
|
||||||
@expose("/testconn", methods=["POST", "GET"])
|
@expose("/testconn", methods=["POST", "GET"])
|
||||||
def testconn(self) -> FlaskResponse: # pylint: disable=no-self-use
|
def testconn(self) -> FlaskResponse:
|
||||||
"""Tests a sqla connection"""
|
"""Tests a sqla connection"""
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"%s.testconn "
|
"%s.testconn "
|
||||||
|
@ -2306,7 +2306,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
|
||||||
@event_logger.log_this
|
@event_logger.log_this
|
||||||
@expose("/validate_sql_json/", methods=["POST", "GET"])
|
@expose("/validate_sql_json/", methods=["POST", "GET"])
|
||||||
def validate_sql_json(
|
def validate_sql_json(
|
||||||
# pylint: disable=too-many-locals,no-self-use
|
# pylint: disable=too-many-locals
|
||||||
self,
|
self,
|
||||||
) -> FlaskResponse:
|
) -> FlaskResponse:
|
||||||
"""Validates that arbitrary sql is acceptable for the given database.
|
"""Validates that arbitrary sql is acceptable for the given database.
|
||||||
|
@ -2406,7 +2406,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
|
||||||
sql_json_executor = Superset._create_sql_json_executor(
|
sql_json_executor = Superset._create_sql_json_executor(
|
||||||
execution_context, query_dao
|
execution_context, query_dao
|
||||||
)
|
)
|
||||||
execution_context_convertor = ExecutionContextConvertorImpl()
|
execution_context_convertor = ExecutionContextConvertor()
|
||||||
execution_context_convertor.set_max_row_in_display(
|
execution_context_convertor.set_max_row_in_display(
|
||||||
int(config.get("DISPLAY_MAX_ROW")) # type: ignore
|
int(config.get("DISPLAY_MAX_ROW")) # type: ignore
|
||||||
)
|
)
|
||||||
|
|
|
@ -0,0 +1,55 @@
|
||||||
|
# Licensed to the Apache Software Foundation (ASF) under one
|
||||||
|
# or more contributor license agreements. See the NOTICE file
|
||||||
|
# distributed with this work for additional information
|
||||||
|
# regarding copyright ownership. The ASF licenses this file
|
||||||
|
# to you 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 json
|
||||||
|
from typing import Iterator
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from sqlalchemy.orm.session import Session
|
||||||
|
|
||||||
|
|
||||||
|
def test_query_dao_save_metadata(app_context: None, session: Session) -> None:
|
||||||
|
from superset.models.core import Database
|
||||||
|
from superset.models.sql_lab import Query
|
||||||
|
|
||||||
|
engine = session.get_bind()
|
||||||
|
Query.metadata.create_all(engine) # pylint: disable=no-member
|
||||||
|
|
||||||
|
db = Database(database_name="my_database", sqlalchemy_uri="sqlite://")
|
||||||
|
|
||||||
|
query_obj = Query(
|
||||||
|
client_id="foo",
|
||||||
|
database=db,
|
||||||
|
tab_name="test_tab",
|
||||||
|
sql_editor_id="test_editor_id",
|
||||||
|
sql="select * from bar",
|
||||||
|
select_sql="select * from bar",
|
||||||
|
executed_sql="select * from bar",
|
||||||
|
limit=100,
|
||||||
|
select_as_cta=False,
|
||||||
|
rows=100,
|
||||||
|
error_message="none",
|
||||||
|
results_key="abc",
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(db)
|
||||||
|
session.add(query_obj)
|
||||||
|
|
||||||
|
from superset.queries.dao import QueryDAO
|
||||||
|
|
||||||
|
query = session.query(Query).one()
|
||||||
|
QueryDAO.save_metadata(query=query, payload={"columns": []})
|
||||||
|
assert query.extra.get("columns", None) == []
|
Loading…
Reference in New Issue