diff --git a/django_ormql/query.py b/django_ormql/query.py index 863ced4..0d9ea80 100644 --- a/django_ormql/query.py +++ b/django_ormql/query.py @@ -185,6 +185,99 @@ class Generator(Generator): } +# Maps Django's Field.get_internal_type() to the ORMQL sql_type vocabulary +# emitted by BaseColumn.sql_type subclasses (see columns.py). Keeping one +# vocabulary means /tables/ introspection and per-query result metadata speak +# the same language. +INTERNAL_TYPE_TO_SQL_TYPE = { + "DateField": "DATE", + "DateTimeField": "DATETIME", + "TimeField": "TIME", + "DurationField": "DURATION", + "IntegerField": "INT", + "BigIntegerField": "INT", + "SmallIntegerField": "INT", + "PositiveIntegerField": "INT", + "PositiveSmallIntegerField": "INT", + "PositiveBigIntegerField": "INT", + "AutoField": "INT", + "BigAutoField": "INT", + "SmallAutoField": "INT", + "DecimalField": "DECIMAL", + "FloatField": "FLOAT", + "BooleanField": "BOOLEAN", + "JSONField": "JSONB", + "CharField": "TEXT", + "TextField": "TEXT", + "EmailField": "TEXT", + "URLField": "TEXT", + "SlugField": "TEXT", + "UUIDField": "TEXT", + "GenericIPAddressField": "TEXT", + "FileField": "TEXT", + "FilePathField": "TEXT", + "ImageField": "TEXT", + "BinaryField": "TEXT", +} + + +def _describe_expression(expr, query=None): + """Extract (sql_type, nullable, hint) from a Django expression's output_field. + Hint is dependent on type and currently only available for truncated datetimes. + + Aggregates and Cast() / typed Func() have `.output_field` directly. Plain + F() references only know their field once resolved against a query; when + `query` is supplied, we resolve on a cloned query (so we don't mutate the + original with extra joins) and read `output_field` off the resolved node. + + Returns ("", None, None) for anything we can't pin down — the frontend treats + that as "unknown, render opaquely" which matches the pre-metadata + behavior. + """ + out = None + try: + out = expr.output_field + except: # noqa + out = None + if out is None and query is not None: + try: + resolved = expr.resolve_expression(query=query.chain(), allow_joins=True) + out = resolved.output_field + except: # noqa + out = None + if out is None: + return "", None + try: + internal = out.get_internal_type() + except: # noqa + return "", None + sql_type = INTERNAL_TYPE_TO_SQL_TYPE.get(internal, "") + nullable = getattr(out, "null", None) + hint = None + + if isinstance(expr, functions.datetime.TruncBase): + hint = {"truncated": expr.kind} + + return sql_type, nullable, hint + + +class Result: + """Iterable query result carrying per-query column metadata. + + Iterating yields row dicts just like the previous evaluate() generator, so + existing `for row in result:` / `list(result)` callers keep working. The + `columns` attribute is a list of {"name", "type", "nullable"} dicts using + the sql_type vocabulary from INTERNAL_TYPE_TO_SQL_TYPE. + """ + + def __init__(self, rows, columns): + self._rows = rows + self.columns = columns + + def __iter__(self): + return iter(self._rows) + + class Query: def __init__(self, sql, tables, placeholders, timezone, default_limit): self.sql = sql @@ -598,7 +691,7 @@ def _resolve(e, parent_stack, depth): elif isinstance(expression, expressions.Subquery): if not isinstance(expression.this, expressions.Select): raise QueryNotSupported("Only SELECT subqueries are supported") - qs, _ = self._select_to_qs( + qs, _, _ = self._select_to_qs( expression.this, parent_table_stack=parent_table_stack + [table] ) return db_func.AutoTypedSubquery( @@ -607,7 +700,7 @@ def _resolve(e, parent_stack, depth): elif isinstance(expression, expressions.Exists): if not isinstance(expression.this, expressions.Select): raise QueryNotSupported("Only SELECT subqueries are supported") - qs, _ = self._select_to_qs( + qs, _, _ = self._select_to_qs( expression.this, parent_table_stack=parent_table_stack + [table] ) return models.Exists(qs) @@ -719,6 +812,7 @@ def _select_to_qs(self, root, parent_table_stack): values_names = {} aggregations = {} name_to_aggregation = {} + column_types = {} for i, e in enumerate(root.args["expressions"]): if isinstance(e, expressions.Star): raise QueryNotSupported("SELECT * is not supported") @@ -750,6 +844,9 @@ def _select_to_qs(self, root, parent_table_stack): parent_table_stack=parent_table_stack, ) values_names[f"expr{i}"] = n + column_types[f"expr{i}"] = _describe_expression( + django_e, query=qs.query + ) if root.args.get("distinct"): qs = qs.distinct() @@ -858,7 +955,7 @@ def _select_to_qs(self, root, parent_table_stack): elif offset is not None or limit is not None: qs = qs[offset:limit] - return qs, values_names + return qs, values_names, column_types def _flatten_unions(self, root): if isinstance(root, expressions.Select): @@ -895,7 +992,15 @@ def parse(self): try: queries = self._flatten_unions(ast) results = [self._select_to_qs(query, []) for query in queries] - if len({len(values_names.keys()) for qs, values_names in results}) != 1: + if ( + len( + { + len(values_names.keys()) + for qs, values_names, column_types in results + } + ) + != 1 + ): raise QueryError( "All parts of UNION query must return same number of columns" ) @@ -905,24 +1010,39 @@ def parse(self): raise QueryError("Invalid combination of types") from e except Exception as e: raise QueryError("Query parsing failed") from e - - return [qs for qs, values_names in results], results[0][1] + return ( + [qs for qs, values_names, column_types in results], + results[0][1], + results[0][2], + ) def evaluate(self): - querysets, values_names = self.parse() + querysets, values_names, column_types = self.parse() + columns = [ + { + "name": values_names[k], + "type": column_types.get(k, ("", None))[0], + "nullable": column_types.get(k, ("", None))[1], + "hint": column_types.get(k, ("", None))[2], + } + for k in values_names + ] - for qs in querysets: - if isinstance(qs, dict): - yield {values_names[k]: v for k, v in qs.items()} - else: - try: - if settings.DEBUG: - print(f"Generated statement: {qs.query!s}") - for row in qs: - yield { - values_names[k]: v - for k, v in row.items() - if k in values_names - } - except (FieldError, ValueError) as e: - raise QueryError("Invalid combination of types") from e + def _iter(): + for qs in querysets: + if isinstance(qs, dict): + yield {values_names[k]: v for k, v in qs.items()} + else: + try: + if settings.DEBUG: + print(f"Generated statement: {qs.query!s}") + for row in qs: + yield { + values_names[k]: v + for k, v in row.items() + if k in values_names + } + except (FieldError, ValueError) as e: + raise QueryError("Invalid combination of types") from e + + return Result(_iter(), columns) diff --git a/tests/test_result_columns.py b/tests/test_result_columns.py new file mode 100644 index 0000000..d3ae1c7 --- /dev/null +++ b/tests/test_result_columns.py @@ -0,0 +1,141 @@ +import pytest + +from django_ormql.query import Result + + +@pytest.mark.django_db +def test_result_is_iterable_wrapper(engine_t1): + res = engine_t1.query("SELECT title FROM categories") + assert isinstance(res, Result) + rows = list(res) + assert rows == [{"title": "Books"}, {"title": "DVDs"}] + + +@pytest.mark.django_db +def test_columns_simple_fields(engine_t1): + res = engine_t1.query("SELECT title, price, publication_date FROM products") + assert res.columns == [ + {"name": "title", "type": "TEXT", "nullable": False, "hint": None}, + {"name": "price", "type": "DECIMAL", "nullable": False, "hint": None}, + {"name": "publication_date", "type": "DATE", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_datetime_field(engine_t1): + res = engine_t1.query("SELECT created FROM orders") + assert res.columns == [ + {"name": "created", "type": "DATETIME", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_nullable(engine_t1): + res = engine_t1.query("SELECT name FROM orderpositions") + assert res.columns == [ + {"name": "name", "type": "TEXT", "nullable": True, "hint": None}, + ] + + +@pytest.mark.django_db +@pytest.mark.xfail(reason="Known bug, we cannot see the nullability of foreign keys") +def test_columns_nullable_foreignkey(engine_t1): + # Customer is nullable on Order, but name not on customer + res = engine_t1.query("SELECT customer FROM orders") + assert res.columns == [ + {"name": "customer", "type": "TEXT", "nullable": True, "hint": None}, + ] + res = engine_t1.query("SELECT customer.name FROM orders") + assert res.columns == [ + {"name": "customer.name", "type": "TEXT", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_boolean_field(engine_t1): + res = engine_t1.query("SELECT enabled FROM customers") + assert res.columns == [ + {"name": "enabled", "type": "BOOLEAN", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_json_field(engine_t1): + res = engine_t1.query("SELECT address FROM customers") + assert res.columns == [ + {"name": "address", "type": "JSONB", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_json_field_extraction(engine_t1): + res = engine_t1.query("SELECT address->city->state->code as code FROM customers") + assert res.columns == [ + {"name": "code", "type": "JSONB", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_count_aggregate(engine_t1): + res = engine_t1.query("SELECT COUNT(*) AS n FROM categories") + assert res.columns == [ + {"name": "n", "type": "INT", "nullable": False, "hint": None} + ] + + +@pytest.mark.django_db +def test_columns_sum_aggregate(engine_t1): + res = engine_t1.query("SELECT SUM(price) AS total FROM products") + assert len(res.columns) == 1 + assert res.columns[0]["name"] == "total" + assert res.columns[0]["type"] == "DECIMAL" + + +@pytest.mark.django_db +def test_columns_group_by_with_aggregate(engine_t1): + res = engine_t1.query( + "SELECT category.title, COUNT(*) AS n FROM products GROUP BY category.title" + ) + assert res.columns == [ + {"name": "category.title", "type": "TEXT", "nullable": False, "hint": None}, + {"name": "n", "type": "INT", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_preserved_with_alias(engine_t1): + res = engine_t1.query("SELECT title AS name, price AS amount FROM products") + assert res.columns == [ + {"name": "name", "type": "TEXT", "nullable": False, "hint": None}, + {"name": "amount", "type": "DECIMAL", "nullable": False, "hint": None}, + ] + + +@pytest.mark.django_db +def test_columns_datetime_trunc(engine_t1): + for k in ("year", "quarter", "month", "week", "day", "hour", "minute", "second"): + res = engine_t1.query(f"SELECT DATETRUNC('{k}', created) as col FROM orders") + assert res.columns == [ + { + "name": "col", + "type": "DATETIME", + "nullable": False, + "hint": {"truncated": k}, + }, + ] + + +@pytest.mark.django_db +def test_columns_date_trunc(engine_t1): + for k in ("year", "quarter", "month", "week"): + res = engine_t1.query( + f"SELECT DATETRUNC('{k}', publication_date) as col FROM products" + ) + assert res.columns == [ + { + "name": "col", + "type": "DATE", + "nullable": False, + "hint": {"truncated": k}, + }, + ] diff --git a/tests/testapp/models.py b/tests/testapp/models.py index 3fb894d..dac6833 100644 --- a/tests/testapp/models.py +++ b/tests/testapp/models.py @@ -77,6 +77,7 @@ class OrderPosition(models.Model): quantity = models.IntegerField() single_price = models.DecimalField(max_digits=10, decimal_places=2) tax_rate = models.DecimalField(max_digits=10, decimal_places=2) + name = models.CharField(max_length=250, null=True) class Meta: ordering = ("id",) diff --git a/tests/testapp/ormql.py b/tests/testapp/ormql.py index 12a87e3..74083a8 100644 --- a/tests/testapp/ormql.py +++ b/tests/testapp/ormql.py @@ -116,6 +116,7 @@ class Meta: "quantity", "single_price", "tax_rate", + "name", ]