Repository: dbt-athena/dbt-athena
Pull request #353 (2023-07-12 12:51:05)

feat: add data_scanned_in_bytes to adapter response
### Description

This PR adds the field `data_scanned_in_bytes` to the adapter response, which allows for monitoring of Athena costs per DBT project/model/etc., see #335.

## Models used to test - Optional

Tested by running some models from my own DBT project and comparing the output of 

```
cat target/run_results.json |  jq '.results[].adapter_response.data_scanned_in_bytes'
```

with the query stats from Athena console -> results checked out.


## Checklist
- [x] You followed [contributing section](https://github.com/dbt-athena/dbt-athena#contributing)
- [x] You kept your Pull Request small and focused on a single feature or bug fix.
- [x] You added unit testing when necessary
- [x] You added functional testing when necessary



Patch:
diff --git a/dbt/adapters/athena/connections.py b/dbt/adapters/athena/connections.py
index e69cf1d..40d4c85 100644
--- a/dbt/adapters/athena/connections.py
+++ b/dbt/adapters/athena/connections.py
@@ -37,6 +37,11 @@ from dbt.exceptions import ConnectionError, DbtRuntimeError
 logger = AdapterLogger("Athena")
 
 
+@dataclass
+class AthenaAdapterResponse(AdapterResponse):
+    data_scanned_in_bytes: Optional[int] = None
+
+
 @dataclass
 class AthenaCredentials(Credentials):
     s3_staging_dir: str
@@ -223,9 +228,14 @@ class AthenaConnectionManager(SQLConnectionManager):
         return connection
 
     @classmethod
-    def get_response(cls, cursor: AthenaCursor) -> AdapterResponse:
+    def get_response(cls, cursor: AthenaCursor) -> AthenaAdapterResponse:
         code = "OK" if cursor.state == AthenaQueryExecution.STATE_SUCCEEDED else "ERROR"
-        return AdapterResponse(_message=f"{code} {cursor.rowcount}", rows_affected=cursor.rowcount, code=code)
+        return AthenaAdapterResponse(
+            _message=f"{code} {cursor.rowcount}",
+            rows_affected=cursor.rowcount,
+            code=code,
+            data_scanned_in_bytes=cursor.data_scanned_in_bytes,
+        )
 
     def cancel(self, connection: Connection) -> None:
         connection.handle.cancel()


Test patch:
diff --git a/tests/unit/test_connection_manager.py b/tests/unit/test_connection_manager.py
index a0ede75..c37a479 100644
--- a/tests/unit/test_connection_manager.py
+++ b/tests/unit/test_connection_manager.py
@@ -4,7 +4,7 @@ import pytest
 from pyathena.model import AthenaQueryExecution
 
 from dbt.adapters.athena import AthenaConnectionManager
-from dbt.contracts.connection import AdapterResponse
+from dbt.adapters.athena.connections import AthenaAdapterResponse
 
 
 class TestAthenaConnectionManager:
@@ -19,11 +19,13 @@ class TestAthenaConnectionManager:
         cursor = mock.MagicMock()
         cursor.rowcount = 1
         cursor.state = state
+        cursor.data_scanned_in_bytes = 123
         cm = AthenaConnectionManager(mock.MagicMock())
         response = cm.get_response(cursor)
-        assert isinstance(response, AdapterResponse)
+        assert isinstance(response, AthenaAdapterResponse)
         assert response.code == result
         assert response.rows_affected == 1
+        assert response.data_scanned_in_bytes == 123
 
     def test_data_type_code_to_name(self):
         cm = AthenaConnectionManager(mock.MagicMock())