From 6b6f60a325027d6ef0fefe3e82dfb510daabff8c Mon Sep 17 00:00:00 2001
From: Damian Owsianny <damian3031@gmail.com>
Date: Mon, 12 Aug 2024 14:17:49 +0200
Subject: [PATCH] Add update_type and rowcount to dbt commands SQL status in
 CLI output

---
 .../unreleased/Features-20240812-161625.yaml  |  7 +++
 dbt/adapters/trino/connections.py             | 12 +++-
 .../adapter/test_sql_status_output.py         | 60 +++++++++++++++++++
 3 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 .changes/unreleased/Features-20240812-161625.yaml
 create mode 100644 tests/functional/adapter/test_sql_status_output.py

diff --git a/.changes/unreleased/Features-20240812-161625.yaml b/.changes/unreleased/Features-20240812-161625.yaml
new file mode 100644
index 00000000..51d21829
--- /dev/null
+++ b/.changes/unreleased/Features-20240812-161625.yaml
@@ -0,0 +1,7 @@
+kind: Features
+body: Added update_type and rowcount to dbt commands SQL status in CLI output
+time: 2024-08-12T16:16:25.330055+02:00
+custom:
+  Author: damian3031
+  Issue: "428"
+  PR: "429"
diff --git a/dbt/adapters/trino/connections.py b/dbt/adapters/trino/connections.py
index 28b0e432..f2d49ad9 100644
--- a/dbt/adapters/trino/connections.py
+++ b/dbt/adapters/trino/connections.py
@@ -493,12 +493,20 @@ def open(cls, connection):
 
     @classmethod
     def get_response(cls, cursor) -> TrinoAdapterResponse:
-        message = "SUCCESS"
+        code = cursor._cursor.update_type
+        if code is None:
+            code = "SUCCESS"
+
+        rows_affected = cursor._cursor.rowcount
+        if rows_affected == -1:
+            message = f"{code}"
+        else:
+            message = f"{code} ({rows_affected:_} rows)"
         return TrinoAdapterResponse(
             _message=message,
             query=cursor._cursor.query,
             query_id=cursor._cursor.query_id,
-            rows_affected=cursor._cursor.rowcount,
+            rows_affected=rows_affected,
         )  # type: ignore
 
     def cancel(self, connection):
diff --git a/tests/functional/adapter/test_sql_status_output.py b/tests/functional/adapter/test_sql_status_output.py
new file mode 100644
index 00000000..13f3bc98
--- /dev/null
+++ b/tests/functional/adapter/test_sql_status_output.py
@@ -0,0 +1,60 @@
+import pytest
+from dbt.tests.util import run_dbt, run_dbt_and_capture
+
+seed_csv = """
+id,name,some_date
+1,Easton,1981-05-20 06:46:51
+2,Lillian,1978-09-03 18:10:33
+3,Jeremiah,1982-03-11 03:59:51
+4,Nolan,1976-05-06 20:21:35
+""".lstrip()
+
+model_sql = """
+select * from {{ ref('seed') }}
+"""
+
+
+class TestSqlStatusOutput:
+    """
+    Testing if SQL status output contains update_type and rowcount
+    """
+
+    @pytest.fixture(scope="class")
+    def seeds(self):
+        return {
+            "seed.csv": seed_csv,
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "materialization_table.sql": model_sql,
+            "materialization_view.sql": model_sql,
+        }
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "name": "sql_status_output",
+            "models": {
+                "sql_status_output": {
+                    "materialization_table": {"+materialized": "table"},
+                    "materialization_view": {"+materialized": "view"},
+                }
+            },
+        }
+
+    def test_run_seed_test(self, project):
+        results = run_dbt(["seed"], expect_pass=True)
+        assert len(results) == 1
+
+        results, logs = run_dbt_and_capture(["--no-use-colors", "run"], expect_pass=True)
+        assert len(results) == 2
+        assert (
+            f" of 2 OK created sql table model {project.test_schema}.materialization_table  [CREATE TABLE (4 rows) in "
+            in logs
+        )
+        assert (
+            f" of 2 OK created sql view model {project.test_schema}.materialization_view  [CREATE VIEW in "
+            in logs
+        )