diff --git a/querybuilder/atoms/columns.py b/querybuilder/atoms/columns.py index d97f8cc5274256e2a90ffa9c6049ecf5680cde1e..3e9d3e2ae2f79f7528c13257aaa0bd190032249c 100644 --- a/querybuilder/atoms/columns.py +++ b/querybuilder/atoms/columns.py @@ -138,7 +138,17 @@ class Star(Atom): This is NOT a `Column`. """ - __slots__ = () + __slots__ = ("relation",) + + def __init__(self, relation: Optional[qbrelations.Named] = None): + self.relation = relation + + def _get_subtokenize_kwargs(self, tokenizer): + relation = None + if self.relation: + relation = self.relation.subtokenize(tokenizer) + + return {"relation": relation} # Abstracts diff --git a/querybuilder/atoms/relations.py b/querybuilder/atoms/relations.py index 68ded344cf43861d3852ca10175bd61ec0bdaab2..ec9f18286e66f6e1b8e1f6c32b06dca9c34515fe 100644 --- a/querybuilder/atoms/relations.py +++ b/querybuilder/atoms/relations.py @@ -249,6 +249,9 @@ class Named(Prenamed, Fromable): ) super()._init_columns(columns) + def star(self): + return qbcolumns.Star(relation=self) + # Actual classes class Aliased(Named): diff --git a/querybuilder/drivers/sql/tokenizer.py b/querybuilder/drivers/sql/tokenizer.py index e2a7c94ee44f5c5eb784862d355b9381d2e950aa..7dc311ae7f0f0e7d71f7836f7a7ae8da610a4e4c 100644 --- a/querybuilder/drivers/sql/tokenizer.py +++ b/querybuilder/drivers/sql/tokenizer.py @@ -1039,9 +1039,14 @@ class Tokenizer: ########### @__call__.register(qbcolumns.Star) - def _(self, obj: qbcolumns.Star, /) -> TkTree: - # tokenized as an Operator to mirror pygments - return self.tokenize_operator("*") + def _(self, obj: qbcolumns.Star, /, *, relation: Optional[TkSeq]) -> TkTree: + # * tokenized as an Operator to mirror pygments + star = self.tokenize_operator("*") + + if not relation: + return star + + return relation + TkStr(qbtoken.Punctuation, ".").to_seq() + star @__call__.register(qbcolumns.Named) def _(self, obj: qbcolumns.Named, /, *, name: TkTree) -> TkTree: diff --git a/querybuilder/queries/dql.py b/querybuilder/queries/dql.py index c399bcf2d631dca9a33e6488a343cfcc8818dcd8..20ac84e1cf319e707406959cac12d5c26c8229fe 100644 --- a/querybuilder/queries/dql.py +++ b/querybuilder/queries/dql.py @@ -198,7 +198,11 @@ class Select(DQLQuery): columns = [] for i, c in enumerate(self.selected_columns): if isinstance(c, qbcolumns.Star): - columns.extend(self.from_.columns) + if c.relation: + star_columns = c.relation.columns + else: + star_columns = self.from_.columns + columns.extend(star_columns) else: columns.append( qbcolumns.name_column(c, self.aliases[i]) diff --git a/querybuilder/tests/atoms/test_columns.py b/querybuilder/tests/atoms/test_columns.py index dbe90a6f001720d497a59aedcc1c4ea7a404b009..e7f240c6faf1c9963c2176f4c8158bc246382502 100644 --- a/querybuilder/tests/atoms/test_columns.py +++ b/querybuilder/tests/atoms/test_columns.py @@ -612,3 +612,27 @@ class TestExists: result = exists.substitute({pre_query: post_query}) assert post_query == result.query + + +class TestStar: + def test_get_subtokenize_kwargs_with_relation(self): + relation_tok = "rel_name" + relation = Mock() + relation.subtokenize = Mock(return_value=relation_tok) + + star = qbcolumns.Star(relation=relation) + + kwargs = star._get_subtokenize_kwargs(None) + + expected = {"relation": "rel_name"} + + assert expected == kwargs + + def test_get_subtokenize_kwargs_with_relation(self): + star = qbcolumns.Star() + + kwargs = star._get_subtokenize_kwargs(None) + + expected = {"relation": None} + + assert expected == kwargs diff --git a/querybuilder/tests/atoms/test_relations.py b/querybuilder/tests/atoms/test_relations.py index 7e9087e9773212b1dac089579fad9d02022a99e2..6edfd4731dd6f721e2e6f81e2fb736d15e4e151b 100644 --- a/querybuilder/tests/atoms/test_relations.py +++ b/querybuilder/tests/atoms/test_relations.py @@ -262,6 +262,14 @@ class TestNamed: assert tuple(post_rel.columns) == (post_col,) + def test_star(self): + rel = qbrelations.Named("foo") + + star = rel.star() + + assert isinstance(star, qbcolumns.Star) + assert rel == star.relation + class TestAliased: def test_substitute(self): diff --git a/querybuilder/tests/drivers/sql/test_tokenizer.py b/querybuilder/tests/drivers/sql/test_tokenizer.py index 6f7c78ba4bb048762a52187f82d5009507bc78c9..b6d5b34b8f337bd79d8745c48743f460d4aecfc4 100644 --- a/querybuilder/tests/drivers/sql/test_tokenizer.py +++ b/querybuilder/tests/drivers/sql/test_tokenizer.py @@ -1465,12 +1465,26 @@ class TestSQLTokenizer: assert expected == result - def test_tokenize_star(self): - result = self.tk(self.get_empty_instance(qb.atoms.columns.Star)) + def test_tokenize_star_without_relation(self): + result = self.tk(self.get_empty_instance(qb.atoms.columns.Star), relation=()) expected = TkStr(qbtoken.Operator, "*").to_seq() assert expected == result + def test_tokenize_star_with_relation(self): + relation = self.get_dummy_tkseq("relation") + result = self.tk( + self.get_empty_instance(qb.atoms.columns.Star), relation=relation + ) + + expected = ( + relation + + TkStr(qbtoken.Punctuation, ".").to_seq() + + TkStr(qbtoken.Operator, "*").to_seq() + ) + + assert expected == result + def test_aliased_relation_with_tkseq_without_column_aliases(self): name = self.get_dummy_tkseq(value="bar") subrelation = self.get_dummy_tkseq(value="foo") diff --git a/querybuilder/tests/queries/test_dql.py b/querybuilder/tests/queries/test_dql.py index 0ce3cfbe960d129d4abe1b30f0ae15ba373c50a6..4da5ba7f05537fc92b4812c438650e61422ad602 100644 --- a/querybuilder/tests/queries/test_dql.py +++ b/querybuilder/tests/queries/test_dql.py @@ -285,6 +285,38 @@ class TestSelect(TestDQL): for i, c in enumerate(sel.columns): assert expected_names[i] == c.name + def test_column_property_with_star(self): + columns = [qbcolumns.Named(int, f"c{i}") for i in range(4)] + aliases = {1: "a0"} + + rel_1 = qbrelations.Named("foo", columns=columns[:2]) + rel_2 = qbrelations.Named("bar", columns=columns[2:]) + + sel = dql.Select( + (qbcolumns.Star(), columns[0]), aliases, from_=rel_1.product(rel_2) + ) + + expected_names = ["c0", "c1", "c2", "c3", "a0"] + for i, c in enumerate(sel.columns): + assert expected_names[i] == c.name + + def test_column_property_with_relation_star(self): + columns = [qbcolumns.Named(int, f"c{i}") for i in range(4)] + aliases = {1: "a0"} + + rel_1 = qbrelations.Named("foo", columns=columns[:2]) + rel_2 = qbrelations.Named("bar", columns=columns[2:]) + + sel = dql.Select( + (qbcolumns.Star(relation=rel_2), columns[0]), + aliases, + from_=rel_1.product(rel_2), + ) + + expected_names = ["c2", "c3", "a0"] + for i, c in enumerate(sel.columns): + assert expected_names[i] == c.name + def test_orderby(self): columns = [qbcolumns.Named(int, f"c{i}") for i in range(5)] sel = dql.Select(columns, orderby=columns[2])