From 059a356841f585840070b5c7b1a262b0ca5e0cf4 Mon Sep 17 00:00:00 2001
From: agl29 <ayush.goyal@cloudera.com>
Date: Thu, 24 Oct 2024 12:43:49 +0530
Subject: [PATCH] [pytest] add pytest-cov

---
 .github/workflows/commitflow-py3.yml          | 39 +++++++++++++++++++
 desktop/core/base_requirements.txt            |  1 +
 .../notebook/src/notebook/connectors/trino.py |  7 ++++
 3 files changed, 47 insertions(+)

diff --git a/.github/workflows/commitflow-py3.yml b/.github/workflows/commitflow-py3.yml
index a43fd333eb2..be079b5a99c 100644
--- a/.github/workflows/commitflow-py3.yml
+++ b/.github/workflows/commitflow-py3.yml
@@ -63,6 +63,45 @@ jobs:
       run: |
         PYTHONWARNINGS=always ./build/env/bin/pytest
 
+    - name: Check and comment if no unit test files are modified
+      if: matrix.python-version == '3.8'
+      run: |
+        git fetch origin master
+        changed_files=$(git diff --name-only origin/master)
+        
+        if echo "$changed_files" | grep -qE '(^test|_test\.py|^tests|_tests\.py|.test)'; then
+          echo "✅ Unit test files were modified."
+        else
+          echo "⚠️ No unit test files modified."
+          
+          curl -X POST \
+            -H "Authorization: token ${{ secrets.GITHUB_TOKEN }}" \
+            -H "Accept: application/vnd.github.v3+json" \
+            -d '{"body":"⚠️ No unit test files modified. Please ensure that changes are properly tested. ⚠️"}' \
+            "https://api.github.com/repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/comments"
+        fi
+
+    # - name: Get list of changed Python files
+    #   id: changes
+    #   run: |
+    #     git fetch origin master
+    #     echo "::set-output name=changed_files::$(git diff --name-only origin/master | grep '\.py$' | tr '\n' ' ')"
+
+    # - name: Run tests with coverage for changed files
+    #   run: |
+    #     if [ "${{ steps.changes.outputs.changed_files }}" != "" ]; then
+    #       ./build/env/bin/pytest --cov --cov-report=xml --cov-report=term-missing ${{ steps.changes.outputs.changed_files }}
+    #     else
+    #       echo "No Python files were changed, skipping tests."
+    #     fi
+
+    - name: Upload coverage report as artifact
+      if: ${{ steps.changes.outputs.changed_files != '' }}
+      uses: actions/upload-artifact@v4
+      with:
+        name: coverage-report
+        path: coverage.xml
+
     - name: run python lints
       run: |
         PYTHONWARNINGS=always ./build/env/bin/hue runruff check
diff --git a/desktop/core/base_requirements.txt b/desktop/core/base_requirements.txt
index ac61a3abccd..548468e4564 100644
--- a/desktop/core/base_requirements.txt
+++ b/desktop/core/base_requirements.txt
@@ -47,6 +47,7 @@ protobuf==3.20.3
 py==1.11.0
 pyformance==0.3.2
 pytest==8.1.1
+pytest-cov==5.0.0
 pytest-django==4.8.0
 python-dateutil==2.8.2
 python-daemon==2.2.4
diff --git a/desktop/libs/notebook/src/notebook/connectors/trino.py b/desktop/libs/notebook/src/notebook/connectors/trino.py
index 921b1ce0b7c..149090b22bb 100644
--- a/desktop/libs/notebook/src/notebook/connectors/trino.py
+++ b/desktop/libs/notebook/src/notebook/connectors/trino.py
@@ -289,6 +289,13 @@ def close_statement(self, notebook, snippet):
     return {'status': 0}
 
   def close_session(self, session):
+    catalogs = self._show_catalogs()
+    databases = []
+
+    for catalog in catalogs:
+      query_client = TrinoQuery(self.trino_request, 'SHOW SCHEMAS FROM ' + catalog)
+      response = query_client.execute()
+      databases += [f'{catalog}.{item}' for sublist in response.rows for item in sublist]
     # Avoid closing session on page refresh or editor close for now
     pass