count on union crashes

[snstanton]

On the head of main.

With a simplified query like:

CompanyModel.query.union().count()

The sqlalchemy_utils.py module raises an error trying to dereference a nonexistent 'entity' property:

ERROR	graphql.execution.utils:utils.py:155	Traceback (most recent call last):
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/graphql/execution/executor.py", line 452, in resolve_or_error
    return executor.execute(resolve_fn, source, info, **args)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/graphql/execution/executors/sync.py", line 16, in execute
    return fn(*args, **kwargs)
  File "/Users/scotts/workspace/bug/schema.py", line 10, in resolve_fail
    return CompanyModel.query.union().count()
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 3163, in count
    return self._from_self(col).enable_eagerloads(False).scalar()
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 2888, in scalar
    ret = self.one()
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 2865, in one
    return self._iter().one()
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 2903, in _iter
    result = self.session.execute(
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 1693, in execute
    result = fn(orm_exec_state)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/session.py", line 299, in do_orm_execute
    entities = all_entities_in_statement(execute_state.statement)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 34, in all_entities_in_statement
    entities = get_column_entities(statement)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 65, in get_column_entities
    entities |= get_column_entities(child)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 65, in get_column_entities
    entities |= get_column_entities(child)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 65, in get_column_entities
    entities |= get_column_entities(child)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 59, in get_column_entities
    entities = _entities_in_statement(statement)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 55, in _entities_in_statement
    return set(e for e in entities if e is not None)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 55, in <genexpr>
    return set(e for e in entities if e is not None)
  File "/Users/scotts/.pyenv/versions/bug/lib/python3.10/site-packages/sqlalchemy_oso/sqlalchemy_utils.py", line 54, in <genexpr>
    entities = (cd["entity"] for cd in statement.column_descriptions)
graphql.error.located_error.GraphQLLocatedError: 'entity'

I have attached a project that demonstrates the issue.

bug.zip

The problem appears to be on line 54 of sqlalchemy_utils.py where it dereferences cd["entity"] instead of doing a get.

[kkirsche]

Would you be able to use markdown to make your code more readable by using triple backticks? Also, what version of SQLAlchemy and the oso libraries are you using?

That'll make it easier for others to assist you. Personally, as a contributor not a member of the Oso team, I wouldn't review this as formatted and without additional context (which I may have missed).

Thanks for your time and the bug report.

[snstanton]

Sorry about that. I was using Jira syntax instead of markdown. I've submitted a pull request for the fix, so hopefully it doesn't take much effort to evaluate.