Skip to content

Fix issue of queries not reaching completed state #210

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions tests/integration/test_sqlalchemy_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License
import os
import pytest
import sqlalchemy as sqla
from sqlalchemy.sql import and_, or_, not_
Expand Down Expand Up @@ -177,6 +178,22 @@ def test_conjunctions(trino_connection):
assert len(rows) == 1


@pytest.mark.parametrize('trino_connection', ['system'], indirect=True)
def test_completed_states(trino_connection):
_, conn = trino_connection
metadata = sqla.MetaData()
queries = sqla.Table('queries', metadata, schema='runtime', autoload_with=conn)
s = sqla.select(queries.c.state).where(queries.c.query == "SELECT version()")
result = conn.execute(s)
rows = result.fetchall()
assert len(rows) > 0
for row in rows:
if os.environ.get("TRINO_VERSION") == '351':
assert row['state'] == 'FAILED'
else:
assert row['state'] == 'FINISHED'


@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
def test_textual_sql(trino_connection):
_, conn = trino_connection
Expand Down
2 changes: 1 addition & 1 deletion trino/sqlalchemy/dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,7 @@ def _get_server_version_info(self, connection: Connection) -> Any:
query = "SELECT version()"
try:
res = connection.execute(sql.text(query))
version = res.scalar()
version = res.scalar_one()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add a test?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems that 351 behaves differently with query state ..

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc: @arhimondr @losipiuk Is there consensus on whether the new behaviour is working as expected? I think this is related to trinodb/trino#13055

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be a non-op change as the only difference between scalar and scalar_one is that scalar_one would additionally throw if no rows are returned.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SELECT version() should always return exactly one row and we want to raise exception when it doesn't.
Alternative solution might be to execute res.fetchall() after version = res.scalar() but it's cumbersome.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue is not limited to sqlalchemy. Following dbapi code will also result in cancelled query while the result is scalar.

cur.execute('SELECT VERSION()')
cur.fetchone()
cur.cancel()

It is indeed related to the PR on Trino that @hashhar refers to. The issue is that the python client doesn't acknowledge the reception of the data by calling the next_uri on the Trino API.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe #220 is the way to go. Thanks @mdesmet for the thorough investigation and the fix.

return tuple([version])
except exc.ProgrammingError as e:
logger.debug(f"Failed to get server version: {e.orig.message}")
Expand Down