Add test for queries reaching completed state

Author: guyco33

tests/integration/test_sqlalchemy_integration.py

@@ -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_
@@ -177,6 +178,23 @@ def test_conjunctions(trino_connection):
     assert len(rows) == 1
 
 
+@pytest.mark.parametrize('trino_connection', ['system'], indirect=True)
+def test_scalar_query_completed_state(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:
+            # TODO: Query state should be FINISHED instead (will be fixed in https://github.com/trinodb/trino-python-client/pull/220)
+            assert row['state'] == 'FAILED'
+
+
 @pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
 def test_textual_sql(trino_connection):
     _, conn = trino_connection