Extend Github Action Workflow and DB connection for multiple dialects (#387)

Author: hadia206

.github/workflows/pr_testing.yml

@@ -3,6 +3,29 @@ name: PR Testing
 on:
   pull_request:
   workflow_dispatch:
+    inputs:
+      py310:
+        description: "Python 3.10"
+        type: boolean
+        default: true
+      py311:
+        description: "Python 3.11"
+        type: boolean
+        default: true
+      py312:
+        description: "Python 3.12"
+        type: boolean
+        default: true
+      run-python:
+        description: "Run Main PyDough Tests"
+        type: boolean
+        required: false
+        default: true
+      run-sf:
+        description: "Run Snowflake Tests"
+        type: boolean
+        required: false
+        default: false
 
 # Limit CI to cancel previous runs in the same PR
 concurrency:
@@ -26,23 +49,70 @@ jobs:
         run: |
           set -xe pipefail
           echo "commitMsg=$(git log -1 --pretty=format:'%s')" >> $GITHUB_OUTPUT
-
-  run-tests:
-    strategy:
-      matrix:
-        python-version:
-          - "3.10"
-          - "3.11"
-          - "3.12"
-    needs: [get-msg]
-    name: python
+  # Get Python versions from inputs if workflow_dispatch is used
+  # Otherwise, use all versions for pull_request events.
+  get-py-ver-matrix:
+    name: Get Python Version Matrix
     runs-on: ubuntu-latest
+    outputs:
+      # Get the matrix output produced by the step below (set-matrix).
+      # This will be used in the run tests job.
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+      # Capture the Python versions selected via dispatch as a JSON array output
+      # For manual runs (workflow_dispatch), it uses the selected
+      # inputs to determine which versions to run.
+      # For other events like pull_request, it defaults to all versions.
+      - name: Set selected Python versions
+        id: set-matrix
+        run: |
+          versions=()
+
+          if [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then
+            if [ "${{ inputs.py310 }}" == "true" ]; then
+              versions+=("\"3.10\"")
+            fi
+            if [ "${{ inputs.py311 }}" == "true" ]; then
+              versions+=("\"3.11\"")
+            fi
+            if [ "${{ inputs.py312 }}" == "true" ]; then
+              versions+=("\"3.12\"")
+            fi
+          else
+            # For pull_request and other events, use all versions by default
+            versions=( "\"3.10\"" "\"3.11\"" "\"3.12\"" )
+          fi
+          # Join the array elements with commas and wrap in brackets to make valid JSON
+          joined=$(IFS=, ; echo "[${versions[*]}]")
+
+          # Output to GitHub Actions expected format
+          echo "matrix=$joined" >> $GITHUB_OUTPUT
+
+
+
+  run-python-tests:
+    name: Main Python Tests
+    needs: [get-msg, get-py-ver-matrix]
     # https://docs.github.com/en/actions/learn-github-actions/expressions#contains
     # contains is case-insensitive
-    if: contains(needs.get-msg.outputs.commitMsg, '[run ci]')
+    if: |
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run ci]')) ||
+      (github.event_name == 'workflow_dispatch' && inputs.run-python)
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ${{ github.event_name == 'workflow_dispatch'
+                        && fromJSON(needs.get-py-ver-matrix.outputs.matrix)
+                        || fromJSON('["3.10", "3.11", "3.12"]') }}
+
     steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
+      - uses: actions/checkout@v4
+
+      - name: Setup Python ${{ matrix.python-version }}
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
 
       - name: Install uv
         uses: astral-sh/setup-uv@v3
@@ -54,7 +124,23 @@ jobs:
         run: ./demos/setup_tpch.sh ./tpch.db
 
       - name: Run Ruff
-        run: uv run --python ${{ matrix.python-version }} ruff check .
+        run: uv run ruff check .
 
       - name: Run Tests
-        run: uv run --python ${{ matrix.python-version }} pytest tests/
+        run: uv run pytest tests/ -m "not snowflake" -rs
+
+  run-sf-tests:
+    name: Snowflake Tests
+    needs: [get-msg, get-py-ver-matrix]
+    if: |
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run sf]')) ||
+      (github.event_name == 'workflow_dispatch' && inputs.run-sf)
+    uses: ./.github/workflows/sf_testing.yml
+    secrets:
+      SF_USERNAME: ${{ secrets.SF_USERNAME }}
+      SF_PASSWORD: ${{ secrets.SF_PASSWORD }}
+      SF_ACCOUNT: ${{ secrets.SF_ACCOUNT }}
+    with:
+      python-versions: ${{ github.event_name == 'workflow_dispatch'
+                      && needs.get-py-ver-matrix.outputs.matrix
+                      || '["3.10", "3.11", "3.12"]' }}

.github/workflows/sf_testing.yml

@@ -0,0 +1,60 @@
+name: Run Snowflake Tests
+
+on:
+  workflow_call:
+    inputs:
+      python-versions:
+        description: "JSON string of Python versions"
+        type: string
+        required: true
+    secrets:
+      SF_USERNAME:
+        required: true
+      SF_PASSWORD:
+        required: true
+      SF_ACCOUNT:
+        required: true
+
+jobs:
+  sf-tests:
+    name: Snowflake Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ${{ fromJSON(inputs.python-versions) }}
+
+    env:
+      SF_USERNAME: ${{ secrets.SF_USERNAME }}
+      SF_PASSWORD: ${{ secrets.SF_PASSWORD }}
+      SF_ACCOUNT: ${{ secrets.SF_ACCOUNT }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Python ${{ matrix.python-version }}
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "0.4.23"
+
+      - name: Create virtual environment
+        # uv requires an existing virtual environment to install packages.
+        # Running `uv venv` creates the `.venv` directory so that subsequent
+        # `uv pip install` commands install dependencies inside this environment.
+        # Without this step, `uv pip install` fails with 
+        # "No virtual environment found".
+        run: uv venv
+
+      - name: Install dependencies
+        run: uv pip install -e ".[snowflake]"
+
+      - name: Confirm Snowflake connector is installed
+        run: uv run python -c "import snowflake.connector; print(snowflake.connector.__version__)"
+
+      - name: Run Snowflake Tests
+        run: uv run pytest -m snowflake tests/ -rs

documentation/usage.md

@@ -351,7 +351,7 @@ pydough.active_session.database = education_context
 shakespeare_context  = pydough.active_session.load_database("sqlite", database="db_files/education.db")
 ```
 
-Notice that both APIs `load_database_context` and `sesion.load_database` take in the name of the databse type first and all the connection keyword arguments, and also return the context object.
+Notice that both APIs `load_database_context` and `sesion.load_database` take in the name of the database type first and all the connection keyword arguments, and also return the context object.
 
 It is important to ensure that the correct database context is being used for several reasons:
 - It controls what SQL dialect is used when translating from PyDough to SQL.

pydough/database_connectors/database_connector.py

@@ -5,12 +5,13 @@
 """
 # Copyright (C) 2024 Bodo Inc. All rights reserved.
 
-import sqlite3
 from dataclasses import dataclass
 from enum import Enum
 
 import pandas as pd
 
+from .db_types import DBConnection, DBCursor
+
 __all__ = ["DatabaseConnection", "DatabaseContext", "DatabaseDialect"]
 
 
@@ -24,9 +25,9 @@ class DatabaseConnection:
     # Database connection that follows DB API 2.0 specification.
     # sqlite3 contains the connection specification and is packaged
     # with Python.
-    _connection: sqlite3.Connection
+    _connection: DBConnection
 
-    def __init__(self, connection: sqlite3.Connection) -> None:
+    def __init__(self, connection: DBConnection) -> None:
         self._connection = connection
 
     def execute_query_df(self, sql: str) -> pd.DataFrame:
@@ -42,10 +43,10 @@ def execute_query_df(self, sql: str) -> pd.DataFrame:
         Returns:
             list[pt.Any]: A list of rows returned by the query.
         """
-        cursor: sqlite3.Cursor = self._connection.cursor()
+        cursor: DBCursor = self._connection.cursor()
         try:
             cursor.execute(sql)
-        except sqlite3.OperationalError as e:
+        except Exception as e:
             print(f"ERROR WHILE EXECUTING QUERY:\n{sql}")
             raise e
         column_names: list[str] = [description[0] for description in cursor.description]
@@ -59,13 +60,13 @@ def execute_query_df(self, sql: str) -> pd.DataFrame:
     # how this will be available at a user API level.
 
     @property
-    def connection(self) -> sqlite3.Connection:
+    def connection(self) -> DBConnection:
         """
         Get the database connection. This API may be removed if all
         the functionality can be encapsulated in the DatabaseConnection.
 
         Returns:
-            sqlite3.Connection: The connection PyDough is managing.
+            The database connection PyDough is managing.
         """
         return self._connection
 

pydough/database_connectors/db_types.py

@@ -0,0 +1,52 @@
+"""
+Type aliases for database connections and cursors used in type hints.
+use `if TYPE_CHECKING:` to import database-specific modules in a way
+that allows static type checkers to understand the types without triggering
+runtime imports. This avoids runtime errors when some optional dependencies
+are not installed.
+"""
+
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+if TYPE_CHECKING:
+    # Importing database-specific modules only for type checking
+    # This allows us to use type hints for SQL dialect connections
+    # (SQLite, ..etc.)
+    # without requiring these modules at runtime unless they are actually used.
+    import sqlite3
+
+    SQLiteConn: TypeAlias = sqlite3.Connection
+    SQLiteCursor: TypeAlias = sqlite3.Cursor
+
+    # TBD: Placeholder lines to add other dialects.
+    # 1. Replace with actual dialect module
+    # import dialect1_module
+    # 2. Replace with other dialect connections
+    # Dialect1_Conn: TypeAlias = dialect1_module.Connection
+    # 3. Replace with other dialect cursors
+    # Dialect1_Cursor: TypeAlias = dialect1_module.Cursor
+
+    # 4. Define the type aliases for database connections and cursors
+    DBConnection: TypeAlias = SQLiteConn  # | Dialect1_Conn
+    DBCursor: TypeAlias = SQLiteCursor  # | Dialect1_Cursor
+else:
+    DBConnection: TypeAlias = Any
+    DBCursor: TypeAlias = Any
+    SQLiteConn: TypeAlias = Any
+    SQLiteCursor: TypeAlias = Any
+    # Dialect1_Conn: TypeAlias = Any
+    # Dialect1_Cursor: TypeAlias = Any
+
+# This allows us to use these type aliases in the rest of the code
+# without worrying about whether the specific database modules are available.
+__all__ = [
+    "DBConnection",
+    "DBCursor",
+    "SQLiteConn",
+    "SQLiteCursor",
+    # "Dialect1_Conn",
+    # "Dialect1_Cursor",
+]
+# The type aliases are used to provide a consistent interface for database connections
+# and cursors across different database backends, allowing for easier
+# type hinting and code readability without requiring the actual database modules.