use invoke install-sdv-enterprise

Author: R-Palazzo

.github/workflows/run_benchmark_multi_table.yaml

@@ -31,12 +31,15 @@ jobs:
         source venv/bin/activate
 
         python -m pip install --upgrade pip
-        python -m pip install sdv-enterprise --index-url "https://${USERNAME}:${LICENSE_KEY}@pypi.datacebo.com"
+        python -m pip install invoke tomli packaging
+        invoke install-sdv-enterprise
         python -m pip install "sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-516-add-workflows"
 
         echo "VIRTUAL_ENV=$(pwd)/venv" >> $GITHUB_ENV
         echo "$(pwd)/venv/bin" >> $GITHUB_PATH
 
+    - name: Stop workflow after dependency install (temporary)
+      run: exit 0
     - name: Run SDGym Benchmark
       env:
         GCP_SERVICE_ACCOUNT_JSON: ${{ secrets.GCP_SERVICE_ACCOUNT_JSON }}

.github/workflows/run_benchmark_single_table.yml

@@ -28,6 +28,7 @@ jobs:
         source venv/bin/activate
 
         python -m pip install --upgrade pip
+        python -m pip install invoke tomli packaging
         python -m pip install sdv-enterprise --index-url "https://${USERNAME}:${LICENSE_KEY}@pypi.datacebo.com"
         python -m pip install "sdgym[all] @ git+https://github.com/sdv-dev/SDGym.git@issue-516-add-workflows"
 

tasks.py

@@ -219,3 +219,24 @@ def notify_sdgym_benchmark_uploaded(c, folder_name, commit_url=None, modality='s
 
     post_benchmark_uploaded_message(folder_name, commit_url, modality)
 
+
+@task
+def install_sdv_enterprise(c, username=None, license_key=None):
+    """Install sdv-enterprise using sdv-installer if credentials are available."""
+    # 🔹 Import here to avoid triggering sdgym imports at load time
+    from sdgym._benchmark.credentials_utils import sdv_install_cmd
+
+    username = username or os.getenv("SDV_ENTERPRISE_USERNAME")
+    license_key = license_key or os.getenv("SDV_ENTERPRISE_LICENSE_KEY")
+    credentials = {
+        "sdv": {
+            "username": username,
+            "license_key": license_key,
+        }
+    }
+
+    install_cmd = sdv_install_cmd(credentials)
+    if install_cmd:
+        c.run(install_cmd)
+    else:
+        print("No sdv-enterprise credentials found. Skipping installation.")

tests/test_tasks.py

@@ -1,9 +1,12 @@
 """Tests for the ``tasks.py`` file."""
 
+from unittest.mock import Mock, patch
+
 from tasks import (
     _get_extra_dependencies,
     _get_minimum_versions,
     _resolve_version_conflicts,
+    install_sdv_enterprise,
 )
 
 
@@ -209,3 +212,38 @@ def test__resolve_version_conflicts_pointing_to_branch():
         'rdt==1.1.2',
         'copulas==0.12.0',
     ])
+
+
+@patch('sdgym._benchmark.credentials_utils.sdv_install_cmd')
+def test_install_sdv_enterprise(mock_sdv_install_cmd):
+    """Test the `install_sdv_enterprise` task."""
+    # Setup
+    username = 'test_user'
+    license_key = 'test_license_key'
+    mock_context = Mock()
+    mock_sdv_install_cmd.return_value = 'install command'
+
+    # Run
+    install_sdv_enterprise.body(mock_context, username=username, license_key=license_key)
+
+    # Assert
+    mock_sdv_install_cmd.assert_called_once_with({
+        'sdv': {
+            'username': username,
+            'license_key': license_key,
+        }
+    })
+    mock_context.run.assert_called_once_with('install command')
+
+
+def test_install_sdv_enterprise_no_credentials(capsys):
+    """Test the `install_sdv_enterprise` task without credentials."""
+    # Setup
+    mock_context = Mock()
+
+    # Run
+    install_sdv_enterprise.body(mock_context)
+
+    # Assert
+    captured = capsys.readouterr()
+    assert 'No sdv-enterprise credentials found. Skipping installation.' in captured.out