async def test_edgeql_group_tuple_01(self): await self.con.execute(''' create type tup { create multi property tup -> tuple<int64, int64> ; }; insert tup { tup := {(1, 1), (1, 2), (1, 1), (2, 1)} }; ''') await self.assert_query_result( ''' with X := tup.tup, group X using z := X by z; ''', tb.bag([{ "elements": [[1, 2]], "key": { "z": [1, 2] } }, { "elements": [[2, 1]], "key": { "z": [2, 1] } }, { "elements": tb.bag([[1, 1], [1, 1]]), "key": { "z": [1, 1] } }]))
async def test_edgeql_for_in_computable_05(self): await self.assert_query_result( r''' SELECT User { select_deck := ( FOR letter IN {'X'} UNION ( (SELECT .deck.name) ) ) } FILTER .name = 'Alice'; ''', [{ "select_deck": tb.bag(["Bog monster", "Dragon", "Giant turtle", "Imp"]) }], ) # This one caused a totally nonsense type error. await self.assert_query_result( r''' SELECT User { select_deck := ( FOR letter IN 'X' UNION ( ((SELECT .deck).name) ) ) } FILTER .name = 'Alice'; ''', [{ "select_deck": tb.bag(["Bog monster", "Dragon", "Giant turtle", "Imp"]) }], )
async def test_edgeql_group_agg_multi_03(self): await self.assert_query_result( ''' for g in (group BooleanTest by .val) union ( array_agg(g.elements.tags) ); ''', tb.bag([ ["red"], [], tb.bag(["red", "green"]), tb.bag(["red", "black"]), ]), )
async def test_edgeql_group_simple_agg_01(self): await self.assert_query_result( r''' with module cards select (group Card by .element) { el := .key.element, cs := array_agg(.elements) }; ''', tb.bag([ { 'el': "Water", 'cs': [{ 'id': str }] * 2 }, { 'el': "Fire", 'cs': [{ 'id': str }] * 2 }, { 'el': "Earth", 'cs': [{ 'id': str }] * 2 }, { 'el': "Air", 'cs': [{ 'id': str }] * 3 }, ]), )
async def test_edgeql_group_process_for_01a(self): await self.assert_query_result( r''' WITH MODULE cards FOR g IN (GROUP Card BY .element) UNION ( element := g.key.element, cnt := count(g.elements), ); ''', tb.bag([ { "cnt": 2, "element": "Water" }, { "cnt": 2, "element": "Fire" }, { "cnt": 2, "element": "Earth" }, { "cnt": 3, "element": "Air" }, ]))
async def test_edgeql_group_process_for_01c(self): await self.assert_query_result( r''' with module cards for h in (group Card by .element) union (for g in h union ( element := g.key.element, cnt := count(g.elements), )); ''', tb.bag([ { "cnt": 2, "element": "Water" }, { "cnt": 2, "element": "Fire" }, { "cnt": 2, "element": "Earth" }, { "cnt": 3, "element": "Air" }, ]))
async def test_edgeql_for_in_computable_16(self): await self.assert_query_result( r''' SELECT User { select_deck := assert_exists(assert_distinct(( FOR letter IN {'I', 'B'} UNION ( SELECT User.deck { name, letter := letter } FILTER User.deck.name[0] = letter ) ))) } FILTER .name = 'Alice'; ''', [{ 'select_deck': tb.bag([ { 'name': 'Bog monster', 'letter': 'B' }, { 'name': 'Imp', 'letter': 'I' }, ]) }], )
async def test_edgeql_group_process_for_01d(self): await self.assert_query_result( r''' with module cards for g in (group Card by .element) union (for gi in 0 union ( element := g.key.element, cst := sum(g.elements.cost + gi), )); ''', tb.bag([ { "cst": 5, "element": "Water" }, { "cst": 6, "element": "Fire" }, { "cst": 4, "element": "Earth" }, { "cst": 7, "element": "Air" }, ]))
async def test_edgeql_group_by_tuple_01(self): await self.assert_query_result( r""" GROUP Issue USING B := (Issue.status.name, Issue.time_estimate) # This tuple will be {} for Issues lacking # time_estimate. So effectively we're expecting only 2 # subsets, grouped by: # - {} # - ('Open', 3000) BY B """, tb.bag([ { 'key': { 'B': ["Open", 3000] }, 'elements': [{}] * 1, }, { 'key': { 'B': None }, 'elements': [{}] * 3, }, ]), )
async def test_edgeql_for_in_computable_07(self): await self.assert_query_result( r''' SELECT User { select_deck := assert_distinct(( WITH ps := (FOR x IN {"!", "?"} UNION ( SELECT { z := x }).z), FOR letter IN {'I', 'B'} UNION ( SELECT .deck { name, letter := letter ++ "!" ++ ps, } FILTER User.deck.name[0] = letter ) )) } FILTER .name = 'Alice'; ''', [{ "select_deck": tb.bag([ { "letter": ["B!!", "B!?"], "name": "Bog monster" }, { "letter": ["I!!", "I!?"], "name": "Imp" }, ]) }], )
async def test_edgeql_group_simple_02(self): await self.assert_query_result( r''' SELECT (GROUP cards::Card {name} BY .element) ''', tb.bag([{ "elements": tb.bag([{ "name": "Bog monster" }, { "name": "Giant turtle" }]), "key": { "element": "Water" } }, { "elements": tb.bag([{ "name": "Imp" }, { "name": "Dragon" }]), "key": { "element": "Fire" } }, { "elements": tb.bag([{ "name": "Dwarf" }, { "name": "Golem" }]), "key": { "element": "Earth" } }, { "elements": tb.bag([{ "name": "Sprite" }, { "name": "Giant eagle" }, { "name": "Djinn" }]), "key": { "element": "Air" } }]))
async def test_edgeql_group_agg_multi_02(self): await self.assert_query_result( ''' with module cards for g in (group Card BY .element) union ( count((Award { multi z := g.elements.name }.z)) ); ''', tb.bag([6, 6, 6, 9]), )
async def test_edgeql_group_simple_agg_02(self): await self.assert_query_result( r''' with module cards select (group Card by .element) { el := .key.element, cs := array_agg(.elements { name }) }; ''', tb.bag([{ "cs": tb.bag([{ "name": "Bog monster" }, { "name": "Giant turtle" }]), "el": "Water" }, { "cs": tb.bag([{ "name": "Imp" }, { "name": "Dragon" }]), "el": "Fire", }, { "cs": tb.bag([{ "name": "Dwarf" }, { "name": "Golem" }]), "el": "Earth", }, { "cs": tb.bag([{ "name": "Sprite" }, { "name": "Giant eagle" }, { "name": "Djinn" }]), "el": "Air", }]))
async def test_edgeql_group_simple_03(self): # the compilation here is kind of a bummer; could we avoid an # unnest? await self.assert_query_result( r''' SELECT (GROUP cards::Card {name} BY .element) FILTER .key.element != 'Air'; ''', tb.bag([ { "elements": tb.bag([{ "name": "Bog monster" }, { "name": "Giant turtle" }]), "key": { "element": "Water" } }, { "elements": tb.bag([{ "name": "Imp" }, { "name": "Dragon" }]), "key": { "element": "Fire" } }, { "elements": tb.bag([{ "name": "Dwarf" }, { "name": "Golem" }]), "key": { "element": "Earth" } }, ]))
async def test_edgeql_group_tuple_02(self): await self.assert_query_result( ''' with X := {(1, 1), (1, 2), (1, 1), (2, 1)}, group X using z := X by z; ''', tb.bag([{ "elements": [[1, 2]], "key": { "z": [1, 2] } }, { "elements": [[2, 1]], "key": { "z": [2, 1] } }, { "elements": tb.bag([[1, 1], [1, 1]]), "key": { "z": [1, 1] } }]))
async def test_edgeql_group_process_select_04(self): await self.assert_query_result( r''' WITH MODULE cards SELECT (GROUP Card BY .element) { cnt := count(.elements), }; ''', tb.bag([{ "cnt": 2 }, { "cnt": 2 }, { "cnt": 2 }, { "cnt": 3 }]))
async def test_edgeql_group_agg_multi_01(self): await self.assert_query_result( ''' with module cards for g in (group Card BY .element) union ( array_agg(g.elements.name ++ {"!", "?"}) ); ''', tb.bag([{ "Bog monster!", "Bog monster?", "Giant turtle!", "Giant turtle?" }, {"Imp!", "Imp?", "Dragon!", "Dragon?"}, {"Dwarf!", "Dwarf?", "Golem!", "Golem?"}, { "Sprite!", "Sprite?", "Giant eagle!", "Giant eagle?", "Djinn!", "Djinn?" }]))
async def test_edgeql_for_in_computable_10(self): # This is basically test_edgeql_for_in_computable_01 but with # a WITH binding inside the computable and no link prop # If we just drop the WITH Z binding, we get # `letter does not exist`. # If we replace the WITH Z part with an extra SELECT, # we get the same buggy behavior. await self.assert_query_result( r''' SELECT (SELECT User { select_deck := ( WITH Z := ( FOR letter IN {'I', 'B'} UNION ( SELECT .deck { name, # just define an ad-hoc link prop letter := letter } FILTER User.deck.name[0] = letter ) ), SELECT assert_distinct(Z) ) } FILTER .name = 'Alice') { select_deck: {name, letter} }; ''', [{ 'select_deck': tb.bag([ { 'name': 'Bog monster', 'letter': 'B' }, { 'name': 'Imp', 'letter': 'I' }, ]) }], )
async def test_edgeql_group_for_01(self): await self.assert_query_result( r''' WITH MODULE cards FOR g in (GROUP Card BY .element) UNION ( WITH U := g.elements, SELECT U { name, cost_ratio := .cost / math::mean(g.elements.cost) }); ''', tb.bag([{ "cost_ratio": 0.42857142857142855, "name": "Sprite" }, { "cost_ratio": 0.8571428571428571, "name": "Giant eagle" }, { "cost_ratio": 1.7142857142857142, "name": "Djinn" }, { "cost_ratio": 0.5, "name": "Dwarf" }, { "cost_ratio": 1.5, "name": "Golem" }, { "cost_ratio": 0.3333333333333333, "name": "Imp" }, { "cost_ratio": 1.6666666666666667, "name": "Dragon" }, { "cost_ratio": 0.8, "name": "Bog monster" }, { "cost_ratio": 1.2, "name": "Giant turtle" }]))
async def test_edgeql_for_in_computable_08(self): await self.assert_query_result( r''' SELECT User { select_deck := assert_distinct(( WITH ps := (FOR x in {"!", "?"} UNION (x++"")) FOR letter IN {'I', 'B'} UNION ( SELECT .deck { name, letter := letter ++ "!" ++ ps, correlated := (ps, ps), uncorrelated := ((SELECT ps), (SELECT ps)), } FILTER User.deck.name[0] = letter ) )) } FILTER .name = 'Alice'; ''', [{ "select_deck": tb.bag([ { "name": "Bog monster", "letter": {"B!!", "B!?"}, "correlated": {("!", "!"), ("?", "?")}, "uncorrelated": {("!", "!"), ("!", "?"), ("?", "!"), ("?", "?")} }, { "name": "Imp", "letter": {"I!!", "I!?"}, "correlated": {("!", "!"), ("?", "?")}, "uncorrelated": {("!", "!"), ("!", "?"), ("?", "!"), ("?", "?")} }, ]) }], )
async def test_edgeql_group_process_select_01(self): await self.assert_query_result( r''' WITH MODULE cards SELECT (GROUP Card BY .element) { element := .key.element, cnt := count(.elements), }; ''', tb.bag([{ "cnt": 2, "element": "Water" }, { "cnt": 2, "element": "Fire" }, { "cnt": 2, "element": "Earth" }, { "cnt": 3, "element": "Air" }]))
async def test_edgeql_for_in_computable_09(self): # This is basically test_edgeql_for_in_computable_01 but with # a WITH binding in front of the whole shape await self.assert_query_result( r''' WITH U := ( SELECT User { select_deck := ( FOR letter IN {'I', 'B'} UNION ( SELECT User.deck { name, # just define an ad-hoc link prop @letter := letter } FILTER User.deck.name[0] = letter ) ) } FILTER .name = 'Alice' ), SELECT U { name, select_deck: { name, @letter } }; ''', [{ 'select_deck': tb.bag([ { 'name': 'Bog monster', '@letter': 'B' }, { 'name': 'Imp', '@letter': 'I' }, ]) }], )
async def test_edgeql_group_semijoin_group_01(self): await self.assert_query_result( ''' with module cards group ( select (group Card{name, cost} by .element) order by .key.element limit 1 ).elements by .cost; ''', tb.bag([{ "elements": [{ "cost": 1, "name": "Sprite" }], "grouping": ["cost"], "key": { "cost": 1 } }, { "elements": [{ "cost": 2, "name": "Giant eagle" }], "grouping": ["cost"], "key": { "cost": 2 } }, { "elements": [{ "cost": 4, "name": "Djinn" }], "grouping": ["cost"], "key": { "cost": 4 } }]))
async def test_edgeql_globals_09(self): # Test that overloaded functions work await self.con.execute(''' create function is_active(x: Named) -> bool using (false); create function is_active(x: User) -> bool using ( select x.name ?= global cur_user); create function is_active(x: Card) -> bool using ( select x.name not in array_unpack(global banned_cards)); create function is_active2(x: Named) -> bool using (is_active(x)); ''') await self.con.execute(''' set global cur_user := "******"; ''') await self.con.execute(''' set global banned_cards := ["Dragon"]; ''') results = tb.bag([ { "active": True, "name": "Imp" }, { "active": False, "name": "Dragon" }, { "active": True, "name": "Bog monster" }, { "active": True, "name": "Giant turtle" }, { "active": True, "name": "Dwarf" }, { "active": True, "name": "Golem" }, { "active": True, "name": "Sprite" }, { "active": True, "name": "Giant eagle" }, { "active": False, "name": "Alice" }, { "active": True, "name": "Bob" }, { "active": False, "name": "Carol" }, { "active": False, "name": "Dave" }, { "active": True, "name": "Djinn" }, { "active": False, "name": "1st" }, { "active": False, "name": "2nd" }, { "active": False, "name": "3rd" }, ]) await self.assert_query_result( r''' select Named { name, required active := is_active(Named) } ''', results) await self.assert_query_result( r''' select Named { name, required active := is_active2(Named) } ''', results) # swap the function to use def_cur_user and make sure it still works await self.con.execute(''' alter function is_active(x: User) using ( select x.name = global def_cur_user); ''') await self.con.execute(''' set global cur_user := "******"; ''') await self.con.execute(''' set global def_cur_user := "******"; ''') await self.assert_query_result( r''' select Named { name, required active := is_active(Named) } ''', results) # An indirect call, to make sure it gets update that way await self.assert_query_result( r''' select Named { name, required active := is_active2(Named) } ''', results)
async def test_edgeql_group_by_group_by_02(self): res = tb.bag([{ "elements": tb.bag([ { "key": { "cost": 1, "element": None }, "n": 3 }, { "key": { "cost": 2, "element": None }, "n": 2 }, { "key": { "cost": 3, "element": None }, "n": 2 }, { "key": { "cost": 4, "element": None }, "n": 1 }, { "key": { "cost": 5, "element": None }, "n": 1 }, ]), "key": { "grouping": ["cost"] } }, { "elements": tb.bag([ { "key": { "cost": None, "element": "Water" }, "n": 2 }, { "key": { "cost": None, "element": "Earth" }, "n": 2 }, { "key": { "cost": None, "element": "Fire" }, "n": 2 }, { "key": { "cost": None, "element": "Air" }, "n": 3 }, ]), "key": { "grouping": ["element"] } }]) await self.assert_query_result( ''' WITH MODULE cards, G := ( GROUP ( GROUP Card BY {.element, .cost} ) USING grouping := array_agg(.grouping) BY grouping), SELECT G { key: {grouping}, elements: { n := count(.elements), key: {element, cost}} } ''', res, ) await self.assert_query_result( ''' WITH MODULE cards, SELECT ( GROUP ( GROUP Card BY {.element, .cost} ) USING grouping := array_agg(.grouping) BY grouping) { key: {grouping}, elements: { n := count(.elements), key: {element, cost}} } ''', res, )
async def test_edgeql_group_by_group_by_01(self): res = tb.bag([{ "elements": tb.bag([{ "agrouping": ["element"], "key": { "element": "Water", "nowners": None }, "num": 2 }, { "agrouping": ["element"], "key": { "element": "Fire", "nowners": None }, "num": 2 }, { "agrouping": ["element"], "key": { "element": "Earth", "nowners": None }, "num": 2 }, { "agrouping": ["element"], "key": { "element": "Air", "nowners": None }, "num": 3 }]), "grouping": ["agrouping"], "key": { "agrouping": ["element"] } }, { "elements": tb.bag([{ "agrouping": ["nowners"], "key": { "element": None, "nowners": 3 }, "num": 1 }, { "agrouping": ["nowners"], "key": { "element": None, "nowners": 4 }, "num": 2 }, { "agrouping": ["nowners"], "key": { "element": None, "nowners": 2 }, "num": 5 }, { "agrouping": ["nowners"], "key": { "element": None, "nowners": 1 }, "num": 1 }]), "grouping": ["agrouping"], "key": { "agrouping": ["nowners"] } }]) qry = r''' WITH MODULE cards GROUP ( SELECT ( GROUP Card USING nowners := count(.owners) BY {.element, nowners} ) { num := count(.elements), key: {element, nowners}, agrouping := array_agg((SELECT _ := .grouping ORDER BY _)) } ) BY .agrouping ''' await self.assert_query_result(qry, res) # Wrapping in a select caused trouble await self.assert_query_result(f'SELECT ({qry})', res)
async def test_edgeql_group_sets_02(self): await self.assert_query_result( r''' WITH MODULE cards GROUP Card USING nowners := count(.owners) BY {.element, nowners}; ''', tb.bag([{ "elements": [{ "id": str }] * 2, "grouping": ["element"], "key": { "element": "Water", "nowners": None } }, { "elements": [{ "id": str }] * 2, "grouping": ["element"], "key": { "element": "Fire", "nowners": None } }, { "elements": [{ "id": str }] * 2, "grouping": ["element"], "key": { "element": "Earth", "nowners": None } }, { "elements": [{ "id": str }] * 3, "grouping": ["element"], "key": { "element": "Air", "nowners": None } }, { "elements": [{ "id": str }] * 1, "grouping": ["nowners"], "key": { "element": None, "nowners": 3 } }, { "elements": [{ "id": str }] * 2, "grouping": ["nowners"], "key": { "element": None, "nowners": 4 } }, { "elements": [{ "id": str }] * 5, "grouping": ["nowners"], "key": { "element": None, "nowners": 2 } }, { "elements": [{ "id": str }] * 1, "grouping": ["nowners"], "key": { "element": None, "nowners": 1 } }]), )
async def test_edgeql_group_sets_01(self): await self.assert_query_result( r''' WITH MODULE cards GROUP Card {name} USING nowners := count(.owners) BY {.element, nowners}; ''', tb.bag([{ "elements": [{ "name": "Bog monster" }, { "name": "Giant turtle" }], "grouping": ["element"], "key": { "element": "Water", "nowners": None } }, { "elements": [{ "name": "Dragon" }, { "name": "Imp" }], "grouping": ["element"], "key": { "element": "Fire", "nowners": None } }, { "elements": [{ "name": "Dwarf" }, { "name": "Golem" }], "grouping": ["element"], "key": { "element": "Earth", "nowners": None } }, { "elements": [ { "name": "Djinn" }, { "name": "Giant eagle" }, { "name": "Sprite" }, ], "grouping": ["element"], "key": { "element": "Air", "nowners": None } }, { "elements": [{ "name": "Golem" }], "grouping": ["nowners"], "key": { "element": None, "nowners": 3 } }, { "elements": [{ "name": "Bog monster" }, { "name": "Giant turtle" }], "grouping": ["nowners"], "key": { "element": None, "nowners": 4 } }, { "elements": [ { "name": "Djinn" }, { "name": "Dragon" }, { "name": "Dwarf" }, { "name": "Giant eagle" }, { "name": "Sprite" }, ], "grouping": ["nowners"], "key": { "element": None, "nowners": 2 } }, { "elements": [{ "name": "Imp" }], "grouping": ["nowners"], "key": { "element": None, "nowners": 1 } }]), sort={'elements': lambda x: x['name']}, )
async def _test_edgeql_group_by_group_by_03(self, qry): res = tb.bag([{ "el": "Water", "groups": tb.bag([{ "elements": [{ "cost": 2, "name": "Bog monster" }], "even": 0 }, { "elements": [{ "cost": 3, "name": "Giant turtle" }], "even": 1 }]) }, { "el": "Fire", "groups": [{ "elements": tb.bag([{ "cost": 1, "name": "Imp" }, { "cost": 5, "name": "Dragon" }]), "even": 1 }] }, { "el": "Earth", "groups": [{ "elements": tb.bag([{ "cost": 1, "name": "Dwarf" }, { "cost": 3, "name": "Golem" }]), "even": 1 }] }, { "el": "Air", "groups": tb.bag([{ "elements": tb.bag([{ "cost": 2, "name": "Giant eagle" }, { "cost": 4, "name": "Djinn" }]), "even": 0 }, { "elements": [{ "cost": 1, "name": "Sprite" }], "even": 1 }]) }]) await self.assert_query_result(qry, res)