mirror of https://github.com/apache/superset.git
301 lines
11 KiB
Python
301 lines
11 KiB
Python
"""Unit tests for Superset Celery worker"""
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import time
|
|
import unittest
|
|
|
|
import pandas as pd
|
|
|
|
from superset import app, appbuilder, cli, db, models, sql_lab, dataframe
|
|
from superset.security import sync_role_definitions
|
|
|
|
from .base_tests import SupersetTestCase
|
|
|
|
QueryStatus = models.QueryStatus
|
|
|
|
BASE_DIR = app.config.get('BASE_DIR')
|
|
|
|
|
|
class CeleryConfig(object):
|
|
BROKER_URL = 'sqla+sqlite:///' + app.config.get('SQL_CELERY_DB_FILE_PATH')
|
|
CELERY_IMPORTS = ('superset.sql_lab', )
|
|
CELERY_RESULT_BACKEND = 'db+sqlite:///' + app.config.get('SQL_CELERY_RESULTS_DB_FILE_PATH')
|
|
CELERY_ANNOTATIONS = {'sql_lab.add': {'rate_limit': '10/s'}}
|
|
CONCURRENCY = 1
|
|
app.config['CELERY_CONFIG'] = CeleryConfig
|
|
|
|
|
|
class UtilityFunctionTests(SupersetTestCase):
|
|
|
|
# TODO(bkyryliuk): support more cases in CTA function.
|
|
def test_create_table_as(self):
|
|
select_query = "SELECT * FROM outer_space;"
|
|
updated_select_query = sql_lab.create_table_as(
|
|
select_query, "tmp")
|
|
self.assertEqual(
|
|
"CREATE TABLE tmp AS \nSELECT * FROM outer_space;",
|
|
updated_select_query)
|
|
|
|
updated_select_query_with_drop = sql_lab.create_table_as(
|
|
select_query, "tmp", override=True)
|
|
self.assertEqual(
|
|
"DROP TABLE IF EXISTS tmp;\n"
|
|
"CREATE TABLE tmp AS \nSELECT * FROM outer_space;",
|
|
updated_select_query_with_drop)
|
|
|
|
select_query_no_semicolon = "SELECT * FROM outer_space"
|
|
updated_select_query_no_semicolon = sql_lab.create_table_as(
|
|
select_query_no_semicolon, "tmp")
|
|
self.assertEqual(
|
|
"CREATE TABLE tmp AS \nSELECT * FROM outer_space",
|
|
updated_select_query_no_semicolon)
|
|
|
|
multi_line_query = (
|
|
"SELECT * FROM planets WHERE\n"
|
|
"Luke_Father = 'Darth Vader';")
|
|
updated_multi_line_query = sql_lab.create_table_as(
|
|
multi_line_query, "tmp")
|
|
expected_updated_multi_line_query = (
|
|
"CREATE TABLE tmp AS \nSELECT * FROM planets WHERE\n"
|
|
"Luke_Father = 'Darth Vader';")
|
|
self.assertEqual(
|
|
expected_updated_multi_line_query,
|
|
updated_multi_line_query)
|
|
|
|
|
|
class CeleryTestCase(SupersetTestCase):
|
|
def __init__(self, *args, **kwargs):
|
|
super(CeleryTestCase, self).__init__(*args, **kwargs)
|
|
self.client = app.test_client()
|
|
|
|
def get_query_by_name(self, sql):
|
|
session = db.session
|
|
query = session.query(models.Query).filter_by(sql=sql).first()
|
|
session.close()
|
|
return query
|
|
|
|
def get_query_by_id(self, id):
|
|
session = db.session
|
|
query = session.query(models.Query).filter_by(id=id).first()
|
|
session.close()
|
|
return query
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
try:
|
|
os.remove(app.config.get('SQL_CELERY_DB_FILE_PATH'))
|
|
except OSError as e:
|
|
app.logger.warn(str(e))
|
|
try:
|
|
os.remove(app.config.get('SQL_CELERY_RESULTS_DB_FILE_PATH'))
|
|
except OSError as e:
|
|
app.logger.warn(str(e))
|
|
|
|
sync_role_definitions()
|
|
|
|
worker_command = BASE_DIR + '/bin/superset worker'
|
|
subprocess.Popen(
|
|
worker_command, shell=True, stdout=subprocess.PIPE)
|
|
|
|
admin = appbuilder.sm.find_user('admin')
|
|
if not admin:
|
|
appbuilder.sm.add_user(
|
|
'admin', 'admin', ' user', 'admin@fab.org',
|
|
appbuilder.sm.find_role('Admin'),
|
|
password='general')
|
|
cli.load_examples(load_test_data=True)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
subprocess.call(
|
|
"ps auxww | grep 'celeryd' | awk '{print $2}' | xargs kill -9",
|
|
shell=True
|
|
)
|
|
subprocess.call(
|
|
"ps auxww | grep 'superset worker' | awk '{print $2}' | "
|
|
"xargs kill -9",
|
|
shell=True
|
|
)
|
|
|
|
def run_sql(self, db_id, sql, client_id, cta='false', tmp_table='tmp',
|
|
async='false'):
|
|
self.login()
|
|
resp = self.client.post(
|
|
'/superset/sql_json/',
|
|
data=dict(
|
|
database_id=db_id,
|
|
sql=sql,
|
|
async=async,
|
|
select_as_cta=cta,
|
|
tmp_table_name=tmp_table,
|
|
client_id=client_id,
|
|
),
|
|
)
|
|
self.logout()
|
|
return json.loads(resp.data.decode('utf-8'))
|
|
|
|
def test_add_limit_to_the_query(self):
|
|
session = db.session
|
|
main_db = self.get_main_database(db.session)
|
|
eng = main_db.get_sqla_engine()
|
|
|
|
select_query = "SELECT * FROM outer_space;"
|
|
updated_select_query = main_db.wrap_sql_limit(select_query, 100)
|
|
# Different DB engines have their own spacing while compiling
|
|
# the queries, that's why ' '.join(query.split()) is used.
|
|
# In addition some of the engines do not include OFFSET 0.
|
|
self.assertTrue(
|
|
"SELECT * FROM (SELECT * FROM outer_space;) AS inner_qry "
|
|
"LIMIT 100" in ' '.join(updated_select_query.split())
|
|
)
|
|
|
|
select_query_no_semicolon = "SELECT * FROM outer_space"
|
|
updated_select_query_no_semicolon = main_db.wrap_sql_limit(
|
|
select_query_no_semicolon, 100)
|
|
self.assertTrue(
|
|
"SELECT * FROM (SELECT * FROM outer_space) AS inner_qry "
|
|
"LIMIT 100" in
|
|
' '.join(updated_select_query_no_semicolon.split())
|
|
)
|
|
|
|
multi_line_query = (
|
|
"SELECT * FROM planets WHERE\n Luke_Father = 'Darth Vader';"
|
|
)
|
|
updated_multi_line_query = main_db.wrap_sql_limit(multi_line_query, 100)
|
|
self.assertTrue(
|
|
"SELECT * FROM (SELECT * FROM planets WHERE "
|
|
"Luke_Father = 'Darth Vader';) AS inner_qry LIMIT 100" in
|
|
' '.join(updated_multi_line_query.split())
|
|
)
|
|
|
|
def test_run_sync_query(self):
|
|
main_db = self.get_main_database(db.session)
|
|
eng = main_db.get_sqla_engine()
|
|
perm_name = 'can_sql_json'
|
|
|
|
db_id = main_db.id
|
|
# Case 1.
|
|
# Table doesn't exist.
|
|
sql_dont_exist = 'SELECT name FROM table_dont_exist'
|
|
result1 = self.run_sql(db_id, sql_dont_exist, "1", cta='true')
|
|
self.assertTrue('error' in result1)
|
|
|
|
# Case 2.
|
|
# Table and DB exists, CTA call to the backend.
|
|
sql_where = (
|
|
"SELECT name FROM ab_permission WHERE name='{}'".format(perm_name))
|
|
result2 = self.run_sql(
|
|
db_id, sql_where, "2", tmp_table='tmp_table_2', cta='true')
|
|
self.assertEqual(QueryStatus.SUCCESS, result2['query']['state'])
|
|
self.assertEqual([], result2['data'])
|
|
self.assertEqual([], result2['columns'])
|
|
query2 = self.get_query_by_id(result2['query']['serverId'])
|
|
|
|
# Check the data in the tmp table.
|
|
df2 = pd.read_sql_query(sql=query2.select_sql, con=eng)
|
|
data2 = df2.to_dict(orient='records')
|
|
self.assertEqual([{'name': perm_name}], data2)
|
|
|
|
# Case 3.
|
|
# Table and DB exists, CTA call to the backend, no data.
|
|
sql_empty_result = 'SELECT * FROM ab_user WHERE id=666'
|
|
result3 = self.run_sql(
|
|
db_id, sql_empty_result, "3", tmp_table='tmp_table_3', cta='true',)
|
|
self.assertEqual(QueryStatus.SUCCESS, result3['query']['state'])
|
|
self.assertEqual([], result3['data'])
|
|
self.assertEqual([], result3['columns'])
|
|
|
|
query3 = self.get_query_by_id(result3['query']['serverId'])
|
|
self.assertEqual(QueryStatus.SUCCESS, query3.status)
|
|
|
|
def test_run_async_query(self):
|
|
main_db = self.get_main_database(db.session)
|
|
eng = main_db.get_sqla_engine()
|
|
|
|
# Schedule queries
|
|
|
|
# Case 1.
|
|
# Table and DB exists, async CTA call to the backend.
|
|
sql_where = "SELECT name FROM ab_role WHERE name='Admin'"
|
|
result1 = self.run_sql(
|
|
main_db.id, sql_where, "4", async='true', tmp_table='tmp_async_1',
|
|
cta='true')
|
|
assert result1['query']['state'] in (
|
|
QueryStatus.PENDING, QueryStatus.RUNNING, QueryStatus.SUCCESS)
|
|
|
|
time.sleep(1)
|
|
|
|
# Case 1.
|
|
query1 = self.get_query_by_id(result1['query']['serverId'])
|
|
df1 = pd.read_sql_query(query1.select_sql, con=eng)
|
|
self.assertEqual(QueryStatus.SUCCESS, query1.status)
|
|
self.assertEqual([{'name': 'Admin'}], df1.to_dict(orient='records'))
|
|
self.assertEqual(QueryStatus.SUCCESS, query1.status)
|
|
self.assertTrue("FROM tmp_async_1" in query1.select_sql)
|
|
self.assertTrue("LIMIT 666" in query1.select_sql)
|
|
self.assertEqual(
|
|
"CREATE TABLE tmp_async_1 AS \nSELECT name FROM ab_role "
|
|
"WHERE name='Admin'", query1.executed_sql)
|
|
self.assertEqual(sql_where, query1.sql)
|
|
if eng.name != 'sqlite':
|
|
self.assertEqual(1, query1.rows)
|
|
self.assertEqual(666, query1.limit)
|
|
self.assertEqual(False, query1.limit_used)
|
|
self.assertEqual(True, query1.select_as_cta)
|
|
self.assertEqual(True, query1.select_as_cta_used)
|
|
|
|
def test_get_columns_dict(self):
|
|
main_db = self.get_main_database(db.session)
|
|
df = main_db.get_df("SELECT * FROM multiformat_time_series", None)
|
|
cdf = dataframe.SupersetDataFrame(df)
|
|
if main_db.sqlalchemy_uri.startswith('sqlite'):
|
|
self.assertEqual(
|
|
[{'is_date': True, 'type': 'datetime_string', 'name': 'ds',
|
|
'is_dim': False},
|
|
{'is_date': True, 'type': 'datetime_string', 'name': 'ds2',
|
|
'is_dim': False},
|
|
{'agg': 'sum', 'is_date': False, 'type': 'int64',
|
|
'name': 'epoch_ms', 'is_dim': False},
|
|
{'agg': 'sum', 'is_date': False, 'type': 'int64',
|
|
'name': 'epoch_s', 'is_dim': False},
|
|
{'is_date': True, 'type': 'datetime_string', 'name': 'string0',
|
|
'is_dim': False},
|
|
{'is_date': False, 'type': 'object',
|
|
'name': 'string1', 'is_dim': True},
|
|
{'is_date': True, 'type': 'datetime_string', 'name': 'string2',
|
|
'is_dim': False},
|
|
{'is_date': False, 'type': 'object',
|
|
'name': 'string3', 'is_dim': True}]
|
|
, cdf.columns_dict
|
|
)
|
|
else:
|
|
self.assertEqual(
|
|
[{'is_date': True, 'type': 'datetime_string', 'name': 'ds',
|
|
'is_dim': False},
|
|
{'is_date': True, 'type': 'datetime64[ns]',
|
|
'name': 'ds2', 'is_dim': False},
|
|
{'agg': 'sum', 'is_date': False, 'type': 'int64',
|
|
'name': 'epoch_ms', 'is_dim': False},
|
|
{'agg': 'sum', 'is_date': False, 'type': 'int64',
|
|
'name': 'epoch_s', 'is_dim': False},
|
|
{'is_date': True, 'type': 'datetime_string', 'name': 'string0',
|
|
'is_dim': False},
|
|
{'is_date': False, 'type': 'object',
|
|
'name': 'string1', 'is_dim': True},
|
|
{'is_date': True, 'type': 'datetime_string', 'name': 'string2',
|
|
'is_dim': False},
|
|
{'is_date': False, 'type': 'object',
|
|
'name': 'string3', 'is_dim': True}]
|
|
, cdf.columns_dict
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|