committing database connection along with pytest setup.

.github/workflows/pytest.yml

@@ -0,0 +1,40 @@
+name: Python package
+
+on:
+  pull_request:
+  push:
+    branches: [ $default-branch ]
+
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          pip install .[test]
+      - name: Test with pytest
+        run: |
+          pytest
+
+  ruff:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Ruff Lint
+        uses: astral-sh/ruff-action@v1
+        with:
+          args: "check --config pyproject.toml"
+      - name: Ruff Format
+        uses: astral-sh/ruff-action@v1
+        with:
+          args: "format --diff"

src/analytics/DatabaseConnection.py

@@ -0,0 +1,34 @@
+import duckdb 
+
+class DatabaseConnection:
+    def __init__(self, db_url: str, min_elapsed=600):
+        self.db_url = db_url
+        self.connection = None
+        self.connect()
+        self.df = self.connection.query(
+            "select GPUs, GPUMemUsage, GPUComputeUsage, GPUType, Elapsed, "
+            "StartTime,"
+            "StartTime-SubmitTime as Queued, TimeLimit, Interactive, "
+            "IsArray, JobID, ArrayID, Status, Constraints, Partition, User, Account from Jobs "
+            f"where GPUs > 0 and Elapsed>{int(min_elapsed)} and GPUType is not null "
+            " and Status != 'CANCELLED' and Status != 'FAILED'"
+        ).to_df()
+
+    def connect(self):
+        # Simulate a database connection
+        self.connection = duckdb.connect(self.db_url)
+        print(f"Connected to {self.db_url}")
+        return self.connection
+    def as_dataframe(self):
+        """Return the dataframe containing job metrics."""
+        return self.df
+    def disconnect(self):
+        # Simulate closing the database connection
+        self.connection.close()
+        self.connection = None
+
+    def is_connected(self) -> bool:
+        return self.connection is not None
+
+    def get_connection_info(self) -> str:
+        return self.connection if self.is_connected() else "No active connection"

tests/test_example.py

@@ -0,0 +1,13 @@
+import pytest 
+from analytics.DatabaseConnection import DatabaseConnection
+import pandas as pd
+def createDBConnection():
+    return DatabaseConnection("C:/Users/Nitya Karthik A/ds4cg-job-analytics/data/slurm_data_small.db")
+
+def test_db_connection():
+    db_conn = createDBConnection()
+    assert db_conn.is_connected() == True
+    assert isinstance(db_conn.as_dataframe(), pd.DataFrame)
+    db_conn.disconnect()
+    assert db_conn.is_connected() == False
+    assert db_conn.get_connection_info() == "No active connection"