def test_exception(get_table_by_id_mock, db_engine, fake_table):
get_table_by_id_mock.return_value = fake_table
# Invalid order by column
with pytest.raises(SamplesError):
make_samples_query(table_id=1234, limit=1001, order_by="employee_id")
# Invalid partition
with pytest.raises(SamplesError):
make_samples_query(table_id=1234,
limit=1001,
partition="dt=1000-01-01")
# Invalid filter column
with pytest.raises(SamplesError):
make_samples_query(table_id=1234,
limit=1001,
where=[["employee_id", "=", "5"]])
# Invalid filter op
with pytest.raises(SamplesError):
make_samples_query(table_id=1234,
limit=1001,
where=[["id", "==", "5"]])
# Invalid filter value
with pytest.raises(AttributeError):
make_samples_query(table_id=1234, limit=1001, where=[["id", "=", 5]])
def test_basic(get_table_by_id_mock, db_engine, fake_table):
get_table_by_id_mock.return_value = fake_table
assert """
SELECT
*
FROM data.session_data
WHERE
dt='2019-11-09'
LIMIT 1""" == make_samples_query(table_id=1234, limit=1)
assert """
SELECT
*
FROM data.session_data
WHERE
dt='2019-11-08'
LIMIT 2""" == make_samples_query(table_id=1234,
limit=2,
partition="dt=2019-11-08")
assert """
SELECT
*
FROM data.session_data
WHERE
dt='2019-11-09'
ORDER BY id ASC
LIMIT 3""" == make_samples_query(table_id=1234, limit=3, order_by="id")
def test_where_filter(get_table_by_id_mock, db_engine, fake_table):
get_table_by_id_mock.return_value = fake_table
assert """
SELECT
*
FROM data.session_data
WHERE
dt='2019-11-09' AND id = 5
LIMIT 4""" == make_samples_query(table_id=1234,
limit=4,
where=[["id", "=", "5"]])
assert """
SELECT
*
FROM data.session_data
WHERE
dt='2019-11-09' AND id = 5 AND id IS NOT NULL
LIMIT 4""" == make_samples_query(table_id=1234,
limit=4,
where=[["id", "=", "5"],
["id", "IS NOT NULL", ""]])
def run_sample_query(
self,
table_id,
engine_id,
uid,
limit,
partition,
where,
order_by,
order_by_asc,
):
# Initialize progress to 0 for polling purposes
self.update_state(state="PROGRESS", meta=0)
with DBSession() as session:
query = make_samples_query(
table_id,
limit=limit,
partition=partition,
where=where,
order_by=order_by,
order_by_asc=order_by_asc,
session=session,
)
async_execute_query = ExecuteQuery(True)
async_execute_query(query, engine_id, uid=uid, session=session)
while not async_execute_query.poll():
self.update_state(state="PROGRESS",
meta=async_execute_query.progress)
results = {
"created_at": DATETIME_TO_UTC(datetime.now()),
"value": async_execute_query.result,
"engine_id": engine_id,
"created_by": uid,
}
mysql_cache.set_key(
f"table_samples_{table_id}_{uid}",
results,
expires_after=seconds_in_a_day,
session=session,
)
def get_table_samples_raw_query(
table_id,
partition=None,
where=None,
order_by=None,
order_by_asc=True,
limit=100,
):
with DBSession() as session:
verify_data_table_permission(table_id, session=session)
return make_samples_query(
table_id,
limit=limit,
partition=partition,
where=where,
order_by=order_by,
order_by_asc=order_by_asc,
session=session,
)