Update tests

- Don’t import the whole module
- Correct repr
- Increase coverage of trino_database.py

koku/koku/test_trino_db_utils.py

@@ -1,21 +1,26 @@
 import textwrap
 import uuid
 
+from django.test import TestCase
 from jinjasql import JinjaSql
 from trino.dbapi import Connection
+from trino.exceptions import TrinoQueryError
 
-from . import trino_database as trino_db
 from api.iam.test.iam_test_case import FakeTrinoConn
 from api.iam.test.iam_test_case import FakeTrinoCur
 from api.iam.test.iam_test_case import IamTestCase
+from koku.trino_database import connect
+from koku.trino_database import executescript
+from koku.trino_database import PreprocessStatementError
+from koku.trino_database import TrinoStatementExecError
 
 
 class TestTrinoDatabaseUtils(IamTestCase):
     def test_connect(self):
         """
         Test connection to trino returns trino.dbapi.Connection instance
         """
-        conn = trino_db.connect(schema=self.schema_name, catalog="hive")
+        conn = connect(schema=self.schema_name, catalog="hive")
         self.assertTrue(isinstance(conn, Connection))
         self.assertEqual(conn.schema, self.schema_name)
         self.assertEqual(conn.catalog, "hive")
@@ -50,12 +55,10 @@ def test_executescript(self):
             "int_data": 255,
             "txt_data": "This is a test",
         }
-        results = trino_db.executescript(
-            conn, textwrap.dedent(sqlscript), params=params, preprocessor=JinjaSql().prepare_query
-        )
+        results = executescript(conn, textwrap.dedent(sqlscript), params=params, preprocessor=JinjaSql().prepare_query)
         self.assertEqual(results, [["eek"], ["eek"], ["eek"], ["eek"], ["eek"], ["eek"]])
 
-    def test_executescript_err(self):
+    def test_executescript_preprocessor_error(self):
         """
         Test executescript will raise a preprocessor error
         """
@@ -64,35 +67,58 @@ def test_executescript_err(self):
             select * from eek where val1 in {{val_list}};
         """
         params = {"val_list": (1, 2, 3, 4, 5)}
-        with self.assertRaises(trino_db.PreprocessStatementError):
-            trino_db.executescript(
-                conn, textwrap.dedent(sqlscript), params=params, preprocessor=JinjaSql().prepare_query
-            )
+        with self.assertRaises(PreprocessStatementError):
+            executescript(conn, textwrap.dedent(sqlscript), params=params, preprocessor=JinjaSql().prepare_query)
 
-    def test_executescript_no_preprocess(self):
+    def test_executescript_no_preprocessor_error(self):
         """
-        Test executescript will raise a preprocessor error
+        Test executescript will not raise a preprocessor error
         """
         sqlscript = """
             select x from y;
             select a from b;
         """
         conn = FakeTrinoConn()
-        res = trino_db.executescript(conn, textwrap.dedent(sqlscript))
+        res = executescript(conn, textwrap.dedent(sqlscript))
         self.assertEqual(res, [["eek"], ["eek"]])
 
-    def test_preprocessor_err(self):
+    def test_executescript_trino_error(self):
+        """
+        Test that executescirpt will raise a TrinoStatementExecError
+        """
+
+        class FakeTrinoConn:
+            @property
+            def cursor(self):
+                raise TrinoQueryError(
+                    {
+                        "errorName": "REMOTE_TASK_ERROR",
+                        "errorType": "INTERNAL_ERROR",
+                        "message": "Expected response code",
+                    }
+                )
+
+        with (
+            self.assertRaisesRegex(TrinoStatementExecError, "type=INTERNAL_ERROR"),
+            self.assertLogs("koku.trino_database", level="WARN") as logger,
+        ):
+            executescript(FakeTrinoConn(), "SELECT x from y")
+
+        self.assertIn("WARNING:koku.trino_database:Trino Query Error", logger.output[0])
+
+    def test_preprocessor_error(self):
         def t_preprocessor(*args):
             raise TypeError("This is a test")
 
         sqlscript = """
+            \n
             select x from y;
             select a from b;
         """
         params = {"eek": 1}
         conn = FakeTrinoConn()
-        with self.assertRaises(trino_db.PreprocessStatementError):
-            trino_db.executescript(conn, textwrap.dedent(sqlscript), params=params, preprocessor=t_preprocessor)
+        with self.assertRaises(PreprocessStatementError):
+            executescript(conn, textwrap.dedent(sqlscript), params=params, preprocessor=t_preprocessor)
 
     def test_executescript_error(self):
         def t_exec_error(*args, **kwargs):
@@ -112,4 +138,42 @@ def cursor(self):
         """
         with self.assertRaises(ValueError):
             conn = FakerFakeTrinoConn()
-            trino_db.executescript(conn, sqlscript)
+            executescript(conn, sqlscript)
+
+
+class TestTrinoStatementExecError(TestCase):
+    def test_trino_statement_exec_error(self):
+        """Test TestTrinoStatementExecError behavior"""
+
+        trino_query_error = TrinoQueryError(
+            {
+                "errorName": "REMOTE_TASK_ERROR",
+                "errorCode": 99,
+                "errorType": "INTERNAL_ERROR",
+                "failureInfo": {"type": "CLOUD_TROUBLES"},
+                "message": "Expected response code",
+                "errorLocation": {"lineNumber": 42, "columnNumber": 24},
+            },
+            query_id="20231220_165606_25626_c7r5y",
+        )
+
+        trino_statement_error = TrinoStatementExecError("SELECT x from y", 1, {}, trino_query_error)
+
+        expected_str = (
+            "Trino Query Error (TrinoQueryError) : TrinoQueryError(type=INTERNAL_ERROR, name=REMOTE_TASK_ERROR, "
+            'message="Expected response code", query_id=20231220_165606_25626_c7r5y) statement number 1\n'
+            "Statement: SELECT x from y\n"
+            "Parameters: {}"
+        )
+        expected_repr = (
+            "TrinoStatementExecError("
+            "type=INTERNAL_ERROR, "
+            "name=REMOTE_TASK_ERROR, "
+            "message=Expected response code, "
+            "query_id=20231220_165606_25626_c7r5y)"
+        )
+        self.assertEqual(str(trino_statement_error), expected_str)
+        self.assertEqual(repr(trino_statement_error), expected_repr)
+        self.assertEqual(trino_statement_error.error_code, 99)
+        self.assertEqual(trino_statement_error.error_exception, "CLOUD_TROUBLES")
+        self.assertEqual(trino_statement_error.failure_info, {"type": "CLOUD_TROUBLES"})

koku/koku/trino_database.py

@@ -35,11 +35,11 @@ def __init__(
 
     def __repr__(self):
         return (
-            f"{self.__class__.__name__} "
+            f"{self.__class__.__name__}("
             f"type={self.error_type}, "
-            f"name{self.error_name}, "
-            f"message{self.message}, "
-            f"query_id={self.query_id}"
+            f"name={self.error_name}, "
+            f"message={self.message}, "
+            f"query_id={self.query_id})"
         )
 
     def __str__(self):