2019-01-15 18:53:27 -05:00
|
|
|
# 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.
|
2020-02-06 19:45:37 -05:00
|
|
|
# isort:skip_file
|
2018-03-15 20:53:34 -04:00
|
|
|
import textwrap
|
2019-06-24 01:37:41 -04:00
|
|
|
import unittest
|
2021-02-22 08:01:33 -05:00
|
|
|
from unittest import mock
|
2021-09-16 03:55:57 -04:00
|
|
|
|
|
|
|
from superset.exceptions import SupersetException
|
2021-07-01 11:03:07 -04:00
|
|
|
from tests.integration_tests.fixtures.birth_names_dashboard import (
|
|
|
|
load_birth_names_dashboard_with_slices,
|
|
|
|
)
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2020-12-09 15:02:29 -05:00
|
|
|
import pytest
|
2017-02-27 01:54:48 -05:00
|
|
|
from sqlalchemy.engine.url import make_url
|
2021-09-07 06:50:24 -04:00
|
|
|
from sqlalchemy.types import DateTime
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2021-07-01 11:03:07 -04:00
|
|
|
import tests.integration_tests.test_app
|
2020-03-13 15:18:22 -04:00
|
|
|
from superset import app, db as metadata_db
|
2021-09-07 06:50:24 -04:00
|
|
|
from superset.db_engine_specs.postgres import PostgresEngineSpec
|
2021-09-26 14:15:57 -04:00
|
|
|
from superset.common.db_query_status import QueryStatus
|
2017-03-10 12:11:51 -05:00
|
|
|
from superset.models.core import Database
|
2020-03-13 15:18:22 -04:00
|
|
|
from superset.models.slice import Slice
|
2021-09-07 06:50:24 -04:00
|
|
|
from superset.models.sql_types.base import literal_dttm_type_factory
|
2021-09-26 14:15:57 -04:00
|
|
|
from superset.utils.core import get_example_database
|
2019-10-18 17:44:27 -04:00
|
|
|
|
2018-04-15 19:21:33 -04:00
|
|
|
from .base_tests import SupersetTestCase
|
2020-12-09 15:02:29 -05:00
|
|
|
from .fixtures.energy_dashboard import load_energy_table_with_slice
|
2017-02-27 01:54:48 -05:00
|
|
|
|
|
|
|
|
2020-06-29 18:36:06 -04:00
|
|
|
class TestDatabaseModel(SupersetTestCase):
|
2019-08-02 13:01:28 -04:00
|
|
|
@unittest.skipUnless(
|
|
|
|
SupersetTestCase.is_module_installed("requests"), "requests not installed"
|
|
|
|
)
|
2017-04-10 18:36:58 -04:00
|
|
|
def test_database_schema_presto(self):
|
2019-06-25 16:34:48 -04:00
|
|
|
sqlalchemy_uri = "presto://presto.airbnb.io:8080/hive/default"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=sqlalchemy_uri)
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2017-04-10 18:36:58 -04:00
|
|
|
db = make_url(model.get_sqla_engine().url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("hive/default", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
db = make_url(model.get_sqla_engine(schema="core_db").url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("hive/core_db", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
sqlalchemy_uri = "presto://presto.airbnb.io:8080/hive"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=sqlalchemy_uri)
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2017-04-10 18:36:58 -04:00
|
|
|
db = make_url(model.get_sqla_engine().url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("hive", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
db = make_url(model.get_sqla_engine(schema="core_db").url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("hive/core_db", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
|
|
|
def test_database_schema_postgres(self):
|
2019-06-25 16:34:48 -04:00
|
|
|
sqlalchemy_uri = "postgresql+psycopg2://postgres.airbnb.io:5439/prod"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=sqlalchemy_uri)
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2017-04-10 18:36:58 -04:00
|
|
|
db = make_url(model.get_sqla_engine().url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("prod", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
db = make_url(model.get_sqla_engine(schema="foo").url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("prod", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-24 01:37:41 -04:00
|
|
|
@unittest.skipUnless(
|
2019-06-25 16:34:48 -04:00
|
|
|
SupersetTestCase.is_module_installed("thrift"), "thrift not installed"
|
|
|
|
)
|
2019-06-24 01:37:41 -04:00
|
|
|
@unittest.skipUnless(
|
2019-06-25 16:34:48 -04:00
|
|
|
SupersetTestCase.is_module_installed("pyhive"), "pyhive not installed"
|
|
|
|
)
|
2017-04-10 18:36:58 -04:00
|
|
|
def test_database_schema_hive(self):
|
2019-06-25 16:34:48 -04:00
|
|
|
sqlalchemy_uri = "hive://hive@hive.airbnb.io:10000/default?auth=NOSASL"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=sqlalchemy_uri)
|
2017-04-10 18:36:58 -04:00
|
|
|
db = make_url(model.get_sqla_engine().url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("default", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
db = make_url(model.get_sqla_engine(schema="core_db").url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("core_db", db)
|
2017-02-27 01:54:48 -05:00
|
|
|
|
2019-06-24 01:37:41 -04:00
|
|
|
@unittest.skipUnless(
|
2019-06-25 16:34:48 -04:00
|
|
|
SupersetTestCase.is_module_installed("MySQLdb"), "mysqlclient not installed"
|
|
|
|
)
|
2017-04-10 18:36:58 -04:00
|
|
|
def test_database_schema_mysql(self):
|
2019-06-25 16:34:48 -04:00
|
|
|
sqlalchemy_uri = "mysql://root@localhost/superset"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=sqlalchemy_uri)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
|
|
|
db = make_url(model.get_sqla_engine().url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("superset", db)
|
2017-04-10 18:36:58 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
db = make_url(model.get_sqla_engine(schema="staging").url).database
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual("staging", db)
|
2017-09-18 12:52:29 -04:00
|
|
|
|
2019-06-24 01:37:41 -04:00
|
|
|
@unittest.skipUnless(
|
2019-06-25 16:34:48 -04:00
|
|
|
SupersetTestCase.is_module_installed("MySQLdb"), "mysqlclient not installed"
|
|
|
|
)
|
2017-09-18 12:52:29 -04:00
|
|
|
def test_database_impersonate_user(self):
|
2019-06-25 16:34:48 -04:00
|
|
|
uri = "mysql://root@localhost"
|
|
|
|
example_user = "giuseppe"
|
2019-10-15 19:51:04 -04:00
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=uri)
|
2017-09-18 12:52:29 -04:00
|
|
|
|
|
|
|
model.impersonate_user = True
|
|
|
|
user_name = make_url(model.get_sqla_engine(user_name=example_user).url).username
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(example_user, user_name)
|
2017-09-18 12:52:29 -04:00
|
|
|
|
|
|
|
model.impersonate_user = False
|
|
|
|
user_name = make_url(model.get_sqla_engine(user_name=example_user).url).username
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertNotEqual(example_user, user_name)
|
2018-03-15 20:53:34 -04:00
|
|
|
|
2021-02-22 08:01:33 -05:00
|
|
|
@mock.patch("superset.models.core.create_engine")
|
|
|
|
def test_impersonate_user_presto(self, mocked_create_engine):
|
|
|
|
uri = "presto://localhost"
|
|
|
|
principal_user = "logged_in_user"
|
|
|
|
extra = """
|
|
|
|
{
|
|
|
|
"metadata_params": {},
|
|
|
|
"engine_params": {
|
|
|
|
"connect_args":{
|
|
|
|
"protocol": "https",
|
|
|
|
"username":"original_user",
|
|
|
|
"password":"original_user_password"
|
|
|
|
}
|
|
|
|
},
|
|
|
|
"metadata_cache_timeout": {},
|
2021-10-25 06:53:06 -04:00
|
|
|
"schemas_allowed_for_file_upload": []
|
2021-02-22 08:01:33 -05:00
|
|
|
}
|
|
|
|
"""
|
|
|
|
|
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=uri, extra=extra)
|
|
|
|
|
|
|
|
model.impersonate_user = True
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "presto://logged_in_user@localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {
|
|
|
|
"protocol": "https",
|
|
|
|
"username": "original_user",
|
|
|
|
"password": "original_user_password",
|
|
|
|
"principal_username": "logged_in_user",
|
|
|
|
}
|
|
|
|
|
|
|
|
model.impersonate_user = False
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "presto://localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {
|
|
|
|
"protocol": "https",
|
|
|
|
"username": "original_user",
|
|
|
|
"password": "original_user_password",
|
|
|
|
}
|
|
|
|
|
2021-05-29 02:54:18 -04:00
|
|
|
@mock.patch("superset.models.core.create_engine")
|
|
|
|
def test_impersonate_user_trino(self, mocked_create_engine):
|
|
|
|
uri = "trino://localhost"
|
|
|
|
principal_user = "logged_in_user"
|
|
|
|
|
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=uri)
|
|
|
|
|
|
|
|
model.impersonate_user = True
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "trino://localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {
|
|
|
|
"user": "logged_in_user",
|
|
|
|
}
|
|
|
|
|
|
|
|
uri = "trino://original_user:original_user_password@localhost"
|
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=uri)
|
|
|
|
model.impersonate_user = True
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "trino://original_user@localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {"user": "logged_in_user"}
|
|
|
|
|
2021-02-22 08:01:33 -05:00
|
|
|
@mock.patch("superset.models.core.create_engine")
|
|
|
|
def test_impersonate_user_hive(self, mocked_create_engine):
|
|
|
|
uri = "hive://localhost"
|
|
|
|
principal_user = "logged_in_user"
|
|
|
|
extra = """
|
|
|
|
{
|
|
|
|
"metadata_params": {},
|
|
|
|
"engine_params": {
|
|
|
|
"connect_args":{
|
|
|
|
"protocol": "https",
|
|
|
|
"username":"original_user",
|
|
|
|
"password":"original_user_password"
|
|
|
|
}
|
|
|
|
},
|
|
|
|
"metadata_cache_timeout": {},
|
2021-10-25 06:53:06 -04:00
|
|
|
"schemas_allowed_for_file_upload": []
|
2021-02-22 08:01:33 -05:00
|
|
|
}
|
|
|
|
"""
|
|
|
|
|
|
|
|
model = Database(database_name="test_database", sqlalchemy_uri=uri, extra=extra)
|
|
|
|
|
|
|
|
model.impersonate_user = True
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "hive://localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {
|
|
|
|
"protocol": "https",
|
|
|
|
"username": "original_user",
|
|
|
|
"password": "original_user_password",
|
|
|
|
"configuration": {"hive.server2.proxy.user": "logged_in_user"},
|
|
|
|
}
|
|
|
|
|
|
|
|
model.impersonate_user = False
|
|
|
|
model.get_sqla_engine(user_name=principal_user)
|
|
|
|
call_args = mocked_create_engine.call_args
|
|
|
|
|
|
|
|
assert str(call_args[0][0]) == "hive://localhost"
|
|
|
|
|
|
|
|
assert call_args[1]["connect_args"] == {
|
|
|
|
"protocol": "https",
|
|
|
|
"username": "original_user",
|
|
|
|
"password": "original_user_password",
|
|
|
|
}
|
|
|
|
|
2020-12-09 15:02:29 -05:00
|
|
|
@pytest.mark.usefixtures("load_energy_table_with_slice")
|
2018-03-15 20:53:34 -04:00
|
|
|
def test_select_star(self):
|
2019-09-05 15:44:34 -04:00
|
|
|
db = get_example_database()
|
2019-06-25 16:34:48 -04:00
|
|
|
table_name = "energy_usage"
|
2019-09-05 15:44:34 -04:00
|
|
|
sql = db.select_star(table_name, show_cols=False, latest_partition=False)
|
2020-08-27 12:49:18 -04:00
|
|
|
quote = db.inspector.engine.dialect.identifier_preparer.quote_identifier
|
2020-08-06 15:07:22 -04:00
|
|
|
expected = (
|
|
|
|
textwrap.dedent(
|
|
|
|
f"""\
|
2018-03-15 20:53:34 -04:00
|
|
|
SELECT *
|
2020-08-27 12:49:18 -04:00
|
|
|
FROM {quote(table_name)}
|
2019-06-25 16:34:48 -04:00
|
|
|
LIMIT 100"""
|
2020-08-06 15:07:22 -04:00
|
|
|
)
|
2020-08-27 12:49:18 -04:00
|
|
|
if db.backend in {"presto", "hive"}
|
2020-08-06 15:07:22 -04:00
|
|
|
else textwrap.dedent(
|
|
|
|
f"""\
|
|
|
|
SELECT *
|
2020-08-27 12:49:18 -04:00
|
|
|
FROM {table_name}
|
2020-08-06 15:07:22 -04:00
|
|
|
LIMIT 100"""
|
|
|
|
)
|
2019-06-25 16:34:48 -04:00
|
|
|
)
|
2020-08-06 15:07:22 -04:00
|
|
|
assert expected in sql
|
2019-09-05 15:44:34 -04:00
|
|
|
sql = db.select_star(table_name, show_cols=True, latest_partition=False)
|
2020-08-27 12:49:18 -04:00
|
|
|
# TODO(bkyryliuk): unify sql generation
|
|
|
|
if db.backend == "presto":
|
|
|
|
assert (
|
|
|
|
textwrap.dedent(
|
|
|
|
"""\
|
|
|
|
SELECT "source" AS "source",
|
|
|
|
"target" AS "target",
|
|
|
|
"value" AS "value"
|
|
|
|
FROM "energy_usage"
|
|
|
|
LIMIT 100"""
|
|
|
|
)
|
|
|
|
== sql
|
2020-08-06 15:07:22 -04:00
|
|
|
)
|
2020-08-27 12:49:18 -04:00
|
|
|
elif db.backend == "hive":
|
|
|
|
assert (
|
|
|
|
textwrap.dedent(
|
|
|
|
"""\
|
|
|
|
SELECT `source`,
|
|
|
|
`target`,
|
|
|
|
`value`
|
|
|
|
FROM `energy_usage`
|
|
|
|
LIMIT 100"""
|
|
|
|
)
|
|
|
|
== sql
|
|
|
|
)
|
|
|
|
else:
|
|
|
|
assert (
|
|
|
|
textwrap.dedent(
|
|
|
|
"""\
|
|
|
|
SELECT source,
|
|
|
|
target,
|
|
|
|
value
|
|
|
|
FROM energy_usage
|
|
|
|
LIMIT 100"""
|
|
|
|
)
|
|
|
|
in sql
|
2020-08-06 15:07:22 -04:00
|
|
|
)
|
2019-09-05 15:44:34 -04:00
|
|
|
|
|
|
|
def test_select_star_fully_qualified_names(self):
|
|
|
|
db = get_example_database()
|
|
|
|
schema = "schema.name"
|
|
|
|
table_name = "table/name"
|
|
|
|
sql = db.select_star(
|
|
|
|
table_name, schema=schema, show_cols=False, latest_partition=False
|
|
|
|
)
|
|
|
|
fully_qualified_names = {
|
|
|
|
"sqlite": '"schema.name"."table/name"',
|
|
|
|
"mysql": "`schema.name`.`table/name`",
|
|
|
|
"postgres": '"schema.name"."table/name"',
|
|
|
|
}
|
|
|
|
fully_qualified_name = fully_qualified_names.get(db.db_engine_spec.engine)
|
|
|
|
if fully_qualified_name:
|
|
|
|
expected = textwrap.dedent(
|
|
|
|
f"""\
|
|
|
|
SELECT *
|
|
|
|
FROM {fully_qualified_name}
|
|
|
|
LIMIT 100"""
|
|
|
|
)
|
|
|
|
assert sql.startswith(expected)
|
2018-04-18 19:17:28 -04:00
|
|
|
|
2018-05-29 17:20:17 -04:00
|
|
|
def test_single_statement(self):
|
2019-09-08 13:18:09 -04:00
|
|
|
main_db = get_example_database()
|
2018-05-29 17:20:17 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
if main_db.backend == "mysql":
|
|
|
|
df = main_db.get_df("SELECT 1", None)
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(df.iat[0, 0], 1)
|
2018-05-29 17:20:17 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
df = main_db.get_df("SELECT 1;", None)
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(df.iat[0, 0], 1)
|
2018-05-29 17:20:17 -04:00
|
|
|
|
|
|
|
def test_multi_statement(self):
|
2019-09-08 13:18:09 -04:00
|
|
|
main_db = get_example_database()
|
2018-05-29 17:20:17 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
if main_db.backend == "mysql":
|
|
|
|
df = main_db.get_df("USE superset; SELECT 1", None)
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(df.iat[0, 0], 1)
|
2018-05-29 17:20:17 -04:00
|
|
|
|
|
|
|
df = main_db.get_df("USE superset; SELECT ';';", None)
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(df.iat[0, 0], ";")
|
2018-05-29 17:20:17 -04:00
|
|
|
|
2021-09-16 03:55:57 -04:00
|
|
|
@mock.patch("superset.models.core.create_engine")
|
|
|
|
def test_get_sqla_engine(self, mocked_create_engine):
|
|
|
|
model = Database(
|
|
|
|
database_name="test_database", sqlalchemy_uri="mysql://root@localhost",
|
|
|
|
)
|
|
|
|
model.db_engine_spec.get_dbapi_exception_mapping = mock.Mock(
|
|
|
|
return_value={Exception: SupersetException}
|
|
|
|
)
|
|
|
|
mocked_create_engine.side_effect = Exception()
|
|
|
|
with self.assertRaises(SupersetException):
|
|
|
|
model.get_sqla_engine()
|
|
|
|
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2020-06-29 18:36:06 -04:00
|
|
|
class TestSqlaTableModel(SupersetTestCase):
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2018-04-27 00:13:52 -04:00
|
|
|
def test_get_timestamp_expression(self):
|
2021-08-10 12:18:46 -04:00
|
|
|
col_type = (
|
|
|
|
"VARCHAR"
|
|
|
|
if get_example_database().backend == "presto"
|
|
|
|
else "TemporalWrapperType"
|
|
|
|
)
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col = tbl.get_column("ds")
|
2018-04-27 00:13:52 -04:00
|
|
|
sqla_literal = ds_col.get_timestamp_expression(None)
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(str(sqla_literal.compile()), "ds")
|
2021-08-10 12:18:46 -04:00
|
|
|
assert type(sqla_literal.type).__name__ == col_type
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
sqla_literal = ds_col.get_timestamp_expression("P1D")
|
2021-08-10 12:18:46 -04:00
|
|
|
assert type(sqla_literal.type).__name__ == col_type
|
2019-06-25 16:34:48 -04:00
|
|
|
compiled = "{}".format(sqla_literal.compile())
|
|
|
|
if tbl.database.backend == "mysql":
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(compiled, "DATE(ds)")
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2019-02-04 15:34:24 -05:00
|
|
|
prev_ds_expr = ds_col.expression
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col.expression = "DATE_ADD(ds, 1)"
|
|
|
|
sqla_literal = ds_col.get_timestamp_expression("P1D")
|
2021-08-10 12:18:46 -04:00
|
|
|
assert type(sqla_literal.type).__name__ == col_type
|
2019-06-25 16:34:48 -04:00
|
|
|
compiled = "{}".format(sqla_literal.compile())
|
|
|
|
if tbl.database.backend == "mysql":
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(compiled, "DATE(DATE_ADD(ds, 1))")
|
2019-02-04 15:34:24 -05:00
|
|
|
ds_col.expression = prev_ds_expr
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2018-04-27 00:13:52 -04:00
|
|
|
def test_get_timestamp_expression_epoch(self):
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col = tbl.get_column("ds")
|
2018-04-27 00:13:52 -04:00
|
|
|
|
|
|
|
ds_col.expression = None
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col.python_date_format = "epoch_s"
|
2018-04-27 00:13:52 -04:00
|
|
|
sqla_literal = ds_col.get_timestamp_expression(None)
|
2019-06-25 16:34:48 -04:00
|
|
|
compiled = "{}".format(sqla_literal.compile())
|
|
|
|
if tbl.database.backend == "mysql":
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(compiled, "from_unixtime(ds)")
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col.python_date_format = "epoch_s"
|
|
|
|
sqla_literal = ds_col.get_timestamp_expression("P1D")
|
|
|
|
compiled = "{}".format(sqla_literal.compile())
|
|
|
|
if tbl.database.backend == "mysql":
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(compiled, "DATE(from_unixtime(ds))")
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2019-02-04 15:34:24 -05:00
|
|
|
prev_ds_expr = ds_col.expression
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col.expression = "DATE_ADD(ds, 1)"
|
|
|
|
sqla_literal = ds_col.get_timestamp_expression("P1D")
|
|
|
|
compiled = "{}".format(sqla_literal.compile())
|
|
|
|
if tbl.database.backend == "mysql":
|
2019-10-21 10:49:12 -04:00
|
|
|
self.assertEqual(compiled, "DATE(from_unixtime(DATE_ADD(ds, 1)))")
|
2019-02-04 15:34:24 -05:00
|
|
|
ds_col.expression = prev_ds_expr
|
2018-04-27 00:13:52 -04:00
|
|
|
|
2019-02-04 15:34:24 -05:00
|
|
|
def query_with_expr_helper(self, is_timeseries, inner_join=True):
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2019-06-25 16:34:48 -04:00
|
|
|
ds_col = tbl.get_column("ds")
|
2019-02-04 15:34:24 -05:00
|
|
|
ds_col.expression = None
|
|
|
|
ds_col.python_date_format = None
|
|
|
|
spec = self.get_database_by_id(tbl.database_id).db_engine_spec
|
2019-07-23 15:13:58 -04:00
|
|
|
if not spec.allows_joins and inner_join:
|
2019-02-04 15:34:24 -05:00
|
|
|
# if the db does not support inner joins, we cannot force it so
|
|
|
|
return None
|
2019-07-23 15:13:58 -04:00
|
|
|
old_inner_join = spec.allows_joins
|
|
|
|
spec.allows_joins = inner_join
|
2019-02-04 15:34:24 -05:00
|
|
|
arbitrary_gby = "state || gender || '_test'"
|
2019-06-25 16:34:48 -04:00
|
|
|
arbitrary_metric = dict(
|
2021-01-08 17:13:20 -05:00
|
|
|
label="arbitrary", expressionType="SQL", sqlExpression="SUM(num_boys)"
|
2019-06-25 16:34:48 -04:00
|
|
|
)
|
2019-02-04 15:34:24 -05:00
|
|
|
query_obj = dict(
|
2019-06-25 16:34:48 -04:00
|
|
|
groupby=[arbitrary_gby, "name"],
|
2019-02-04 15:34:24 -05:00
|
|
|
metrics=[arbitrary_metric],
|
|
|
|
filter=[],
|
|
|
|
is_timeseries=is_timeseries,
|
|
|
|
columns=[],
|
2019-06-25 16:34:48 -04:00
|
|
|
granularity="ds",
|
2019-02-04 15:34:24 -05:00
|
|
|
from_dttm=None,
|
|
|
|
to_dttm=None,
|
2019-06-25 16:34:48 -04:00
|
|
|
extras=dict(time_grain_sqla="P1Y"),
|
2021-09-16 05:09:08 -04:00
|
|
|
series_limit=15 if inner_join and is_timeseries else None,
|
2019-02-04 15:34:24 -05:00
|
|
|
)
|
|
|
|
qr = tbl.query(query_obj)
|
|
|
|
self.assertEqual(qr.status, QueryStatus.SUCCESS)
|
|
|
|
sql = qr.query
|
|
|
|
self.assertIn(arbitrary_gby, sql)
|
2019-06-25 16:34:48 -04:00
|
|
|
self.assertIn("name", sql)
|
2019-02-04 15:34:24 -05:00
|
|
|
if inner_join and is_timeseries:
|
2019-06-25 16:34:48 -04:00
|
|
|
self.assertIn("JOIN", sql.upper())
|
2019-02-04 15:34:24 -05:00
|
|
|
else:
|
2019-06-25 16:34:48 -04:00
|
|
|
self.assertNotIn("JOIN", sql.upper())
|
2019-07-23 15:13:58 -04:00
|
|
|
spec.allows_joins = old_inner_join
|
2020-01-08 14:50:26 -05:00
|
|
|
self.assertFalse(qr.df.empty)
|
2019-02-04 15:34:24 -05:00
|
|
|
return qr.df
|
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2019-02-04 15:34:24 -05:00
|
|
|
def test_query_with_expr_groupby_timeseries(self):
|
2020-08-06 15:07:22 -04:00
|
|
|
if get_example_database().backend == "presto":
|
|
|
|
# TODO(bkyryliuk): make it work for presto.
|
|
|
|
return
|
|
|
|
|
2019-02-04 15:34:24 -05:00
|
|
|
def cannonicalize_df(df):
|
|
|
|
ret = df.sort_values(by=list(df.columns.values), inplace=False)
|
|
|
|
ret.reset_index(inplace=True, drop=True)
|
|
|
|
return ret
|
|
|
|
|
|
|
|
df1 = self.query_with_expr_helper(is_timeseries=True, inner_join=True)
|
2020-08-03 12:08:49 -04:00
|
|
|
name_list1 = cannonicalize_df(df1).name.values.tolist()
|
2019-02-04 15:34:24 -05:00
|
|
|
df2 = self.query_with_expr_helper(is_timeseries=True, inner_join=False)
|
2020-08-03 12:08:49 -04:00
|
|
|
name_list2 = cannonicalize_df(df1).name.values.tolist()
|
2020-01-08 14:50:26 -05:00
|
|
|
self.assertFalse(df2.empty)
|
2020-08-03 12:08:49 -04:00
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
assert name_list2 == name_list1
|
2019-02-04 15:34:24 -05:00
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2019-02-04 15:34:24 -05:00
|
|
|
def test_query_with_expr_groupby(self):
|
|
|
|
self.query_with_expr_helper(is_timeseries=False)
|
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2018-07-26 18:20:23 -04:00
|
|
|
def test_sql_mutator(self):
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2018-07-26 18:20:23 -04:00
|
|
|
query_obj = dict(
|
|
|
|
groupby=[],
|
2021-03-17 00:00:03 -04:00
|
|
|
metrics=None,
|
2018-07-26 18:20:23 -04:00
|
|
|
filter=[],
|
|
|
|
is_timeseries=False,
|
2019-06-25 16:34:48 -04:00
|
|
|
columns=["name"],
|
2018-07-26 18:20:23 -04:00
|
|
|
granularity=None,
|
2019-06-25 16:34:48 -04:00
|
|
|
from_dttm=None,
|
|
|
|
to_dttm=None,
|
2018-07-26 18:20:23 -04:00
|
|
|
extras={},
|
|
|
|
)
|
|
|
|
sql = tbl.get_query_str(query_obj)
|
2019-07-18 16:17:26 -04:00
|
|
|
self.assertNotIn("-- COMMENT", sql)
|
2018-07-26 18:20:23 -04:00
|
|
|
|
|
|
|
def mutator(*args):
|
2019-07-18 16:17:26 -04:00
|
|
|
return "-- COMMENT\n" + args[0]
|
2019-06-25 16:34:48 -04:00
|
|
|
|
|
|
|
app.config["SQL_QUERY_MUTATOR"] = mutator
|
2018-07-26 18:20:23 -04:00
|
|
|
sql = tbl.get_query_str(query_obj)
|
2019-07-18 16:17:26 -04:00
|
|
|
self.assertIn("-- COMMENT", sql)
|
2018-07-26 18:20:23 -04:00
|
|
|
|
2019-06-25 16:34:48 -04:00
|
|
|
app.config["SQL_QUERY_MUTATOR"] = None
|
2019-07-18 11:51:41 -04:00
|
|
|
|
|
|
|
def test_query_with_non_existent_metrics(self):
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2019-07-18 11:51:41 -04:00
|
|
|
|
|
|
|
query_obj = dict(
|
|
|
|
groupby=[],
|
|
|
|
metrics=["invalid"],
|
|
|
|
filter=[],
|
|
|
|
is_timeseries=False,
|
|
|
|
columns=["name"],
|
|
|
|
granularity=None,
|
|
|
|
from_dttm=None,
|
|
|
|
to_dttm=None,
|
|
|
|
extras={},
|
|
|
|
)
|
|
|
|
|
|
|
|
with self.assertRaises(Exception) as context:
|
|
|
|
tbl.get_query_str(query_obj)
|
|
|
|
|
|
|
|
self.assertTrue("Metric 'invalid' does not exist", context.exception)
|
2020-03-13 15:18:22 -04:00
|
|
|
|
2021-01-11 08:57:55 -05:00
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
2021-11-05 11:05:48 -04:00
|
|
|
def test_data_for_slices_with_no_query_context(self):
|
2021-08-02 15:45:55 -04:00
|
|
|
tbl = self.get_table(name="birth_names")
|
2020-03-13 15:18:22 -04:00
|
|
|
slc = (
|
|
|
|
metadata_db.session.query(Slice)
|
2021-01-11 08:57:55 -05:00
|
|
|
.filter_by(
|
2021-08-03 22:01:39 -04:00
|
|
|
datasource_id=tbl.id, datasource_type=tbl.type, slice_name="Genders",
|
2021-01-11 08:57:55 -05:00
|
|
|
)
|
2020-03-13 15:18:22 -04:00
|
|
|
.first()
|
|
|
|
)
|
|
|
|
data_for_slices = tbl.data_for_slices([slc])
|
2021-08-03 22:01:39 -04:00
|
|
|
assert len(data_for_slices["metrics"]) == 1
|
|
|
|
assert len(data_for_slices["columns"]) == 1
|
|
|
|
assert data_for_slices["metrics"][0]["metric_name"] == "sum__num"
|
|
|
|
assert data_for_slices["columns"][0]["column_name"] == "gender"
|
2021-11-05 11:05:48 -04:00
|
|
|
assert set(data_for_slices["verbose_map"].keys()) == {
|
|
|
|
"__timestamp",
|
|
|
|
"sum__num",
|
|
|
|
"gender",
|
|
|
|
}
|
|
|
|
|
|
|
|
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
|
|
|
|
def test_data_for_slices_with_query_context(self):
|
|
|
|
tbl = self.get_table(name="birth_names")
|
|
|
|
slc = (
|
|
|
|
metadata_db.session.query(Slice)
|
|
|
|
.filter_by(
|
|
|
|
datasource_id=tbl.id,
|
|
|
|
datasource_type=tbl.type,
|
|
|
|
slice_name="Pivot Table v2",
|
|
|
|
)
|
|
|
|
.first()
|
2021-08-03 22:01:39 -04:00
|
|
|
)
|
2021-11-05 11:05:48 -04:00
|
|
|
data_for_slices = tbl.data_for_slices([slc])
|
|
|
|
assert len(data_for_slices["metrics"]) == 1
|
|
|
|
assert len(data_for_slices["columns"]) == 2
|
|
|
|
assert data_for_slices["metrics"][0]["metric_name"] == "sum__num"
|
|
|
|
assert data_for_slices["columns"][0]["column_name"] == "name"
|
|
|
|
assert set(data_for_slices["verbose_map"].keys()) == {
|
|
|
|
"__timestamp",
|
|
|
|
"sum__num",
|
|
|
|
"name",
|
|
|
|
"state",
|
|
|
|
}
|
2021-09-07 06:50:24 -04:00
|
|
|
|
|
|
|
|
|
|
|
def test_literal_dttm_type_factory():
|
|
|
|
orig_type = DateTime()
|
2021-11-25 04:58:44 -05:00
|
|
|
new_type = literal_dttm_type_factory(
|
|
|
|
orig_type, PostgresEngineSpec, "TIMESTAMP", db_extra={}
|
|
|
|
)
|
2021-09-07 06:50:24 -04:00
|
|
|
assert type(new_type).__name__ == "TemporalWrapperType"
|
|
|
|
assert str(new_type) == str(orig_type)
|