def test_query_with_pre_filter_without_filter_taxons(self):
filter_clause = TaxonValueFilterClause({
'type':
FilterClauseType.TAXON_VALUE.value,
'taxon':
'ad_name',
'operator':
SimpleFilterOperator.LIKE.value,
'value':
'%abc%',
})
result, _, _ = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self.scope,
get_specific_select_mocked_taxons(
['impressions', 'ad_id', filter_clause.taxon]),
get_specific_select_mocked_taxons(['impressions', 'ad_id']),
self.graph,
'data-source',
filter_clause=filter_clause,
).get_query()
self.write_test_expectations('query.sql', compile_query(result))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(result).strip())
def test_namespaced_taxons_build_query(self):
result, _, effectively_used_models = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self.scope,
get_specific_select_mocked_taxons([
'impressions', 'ad_id', f'{MOCK_DATA_SOURCE_NAME}|dimension',
f'{MOCK_DATA_SOURCE_NAME}|metric'
]),
get_specific_select_mocked_taxons([
'impressions', 'ad_id', f'{MOCK_DATA_SOURCE_NAME}|dimension',
f'{MOCK_DATA_SOURCE_NAME}|metric'
]),
self.graph,
'data-source',
).get_query()
self.write_test_expectations('query.sql', compile_query(result))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(result))
self.assertDictEqual(
{
'noncached_models': [
'mock_data_source.metric_model',
'mock_data_source.entity_model'
],
},
effectively_used_models.to_primitive(),
)
def test_query_with_pre_taxon_taxon_pre_filter(self):
filter_clause = TaxonTaxonFilterClause({
'type':
FilterClauseType.TAXON_VALUE.value,
'taxon':
'spend',
'right_taxon':
'impressions',
'operator':
SimpleFilterOperator.EQ.value,
})
result, _, _ = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self.scope,
get_specific_select_mocked_taxons(
['ad_id'] + [s for s in filter_clause.get_taxon_slugs()]),
get_specific_select_mocked_taxons(['ad_id']),
self.graph,
'data-source',
filter_clause=filter_clause,
).get_query()
self.write_test_expectations('query.sql', compile_query(result))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(result).strip())
def _run_basic_test(self, projection_taxons, selected_taxons=None):
if selected_taxons is None:
selected_taxons = projection_taxons
query, taxon_model_info_map, _ = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self._scope,
get_specific_select_mocked_taxons(selected_taxons),
get_specific_select_mocked_taxons(projection_taxons),
self._graph,
'data-source',
).get_query()
final_dataframe = ProjectionBuilder.query(
query,
taxon_model_info_map,
get_specific_select_mocked_taxons(projection_taxons),
'data-source',
None,
None,
None,
)
self.write_test_expectations('query.sql',
compile_query(final_dataframe.query))
expected = self.read_test_expectations('query.sql')
assert compile_query(final_dataframe.query) == expected
def test_basic_build_join_query(self):
taxons = get_specific_select_mocked_taxons(
['spend', 'impressions', 'ad_id', 'ad_name', 'week_of_year'])
result, _, _ = SelectBuilder(SNOWFLAKE_HUSKY_CONTEXT, self.scope,
taxons, taxons, self.graph,
'data-source').get_query()
self.write_test_expectations('query.sql', compile_query(result))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(result))
def test_basic_build_query_with_pre_filter(self):
projected_taxons = ['impressions', 'spend']
pre_filter = TaxonValueFilterClause({
'type':
FilterClauseType.TAXON_VALUE.value,
'taxon':
'ad_name',
'operator':
SimpleFilterOperator.LIKE,
'value':
'zombies!',
})
selected_taxons = projected_taxons + ['ad_name']
query, taxon_model_info_map, _ = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self._scope,
get_specific_select_mocked_taxons(selected_taxons),
get_specific_select_mocked_taxons(projected_taxons),
self._graph,
'data-source',
pre_filter,
).get_query()
final_dataframe = ProjectionBuilder.query(
query,
taxon_model_info_map,
get_specific_select_mocked_taxons(projected_taxons),
'data-source',
None,
None,
None,
{'context'},
)
self.write_test_expectations('query.sql',
compile_query(final_dataframe.query))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(final_dataframe.query))
def test_basic_build_join_query(self):
taxons = get_specific_select_mocked_taxons(
['spend', 'gender', 'impressions', 'ad_id', 'ad_name'])
result, _, effectively_used_models = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT, self.scope, taxons, taxons, self.graph,
'data-source').get_query()
self.write_test_expectations('query.sql', compile_query(result))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(result))
self.assertDictEqual(
{
'noncached_models': [
'mock_data_source.metric_gender_model',
'mock_data_source.entity_model'
],
},
effectively_used_models.to_primitive(),
)
def test_basic_build_query_with_order_by(self):
selected_taxons = ['impressions', 'ad_id']
query, taxon_model_info_map, _ = SelectBuilder(
SNOWFLAKE_HUSKY_CONTEXT,
self._scope,
get_specific_select_mocked_taxons(selected_taxons),
get_specific_select_mocked_taxons(selected_taxons),
self._graph,
'data-source',
).get_query()
taxon_order_1 = TaxonDataOrder({
'taxon': 'impressions',
'type': TaxonOrderType.desc.value
})
taxon_order_2 = TaxonDataOrder({
'taxon': 'ad_id',
'type': TaxonOrderType.asc.value
})
order_by = [taxon_order_1, taxon_order_2]
final_dataframe = ProjectionBuilder.query(
query,
taxon_model_info_map,
get_specific_select_mocked_taxons(selected_taxons),
'data-source',
order_by,
1,
2,
{'context'},
)
self.write_test_expectations('query.sql',
compile_query(final_dataframe.query))
expected = self.read_test_expectations('query.sql')
self.assertEqual(expected, compile_query(final_dataframe.query))
def build_query(
ctx: HuskyQueryContext,
subrequest: InternalDataRequest,
query_info: QueryInfo,
preloaded_taxons: TaxonMap,
dimension_templates: Optional[List[SqlFormulaTemplate]] = None,
filter_templates: Optional[TaxonToTemplate] = None,
) -> Dataframe:
"""
Returns Query and Taxons obtained in it
:param dimension_templates Sql column templates to select
:param filter_templates Filter temples keyed by taxon slug, referenced from scope or preagg filters.
"""
dimension_templates = dimension_templates or []
filter_templates = filter_templates or dict()
# Fetch Taxons
simple_taxon_manager = SimpleTaxonManager.initialize(
subrequest, dimension_templates, filter_templates, preloaded_taxons
)
data_sources = set(subrequest.properties.data_sources)
if len(subrequest.properties.data_sources) != 1:
# Joining across data sources is more complex and not implemented yet.
raise MultipleDataSources(data_sources)
data_source = subrequest.properties.data_sources[0]
models = ModelRetriever.load_models(data_sources, subrequest.scope, subrequest.properties.model_name)
# Build Graph
graph = GraphBuilder.create_with_models(models)
# Create Select Query
select_query, taxon_model_info_map, effectively_used_models = SelectBuilder(
ctx,
subrequest.scope,
simple_taxon_manager.graph_select_taxons,
simple_taxon_manager.projection_taxons,
graph,
data_source,
subrequest.preaggregation_filters,
dimension_templates,
filter_templates,
).get_query()
query_info.definition = QueryDefinition({'effectively_used_models': effectively_used_models})
logger.debug('Select Query: %s', compile_query(select_query, ctx.dialect))
# Create Projection Query
final_dataframe = ProjectionBuilder.query(
select_query,
taxon_model_info_map,
simple_taxon_manager.projection_taxons,
subrequest.properties.data_source,
subrequest.order_by,
subrequest.limit,
subrequest.offset,
dimension_templates,
)
logger.debug('Projection Query: %s', compile_query(final_dataframe.query, ctx.dialect))
return final_dataframe