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