feat(cli): Add session management capabilities

Author: pedro93

metadata-ingestion/src/datahub/cli/session_cli.py

@@ -0,0 +1,144 @@
+import json
+import os
+from typing import Dict
+
+import click
+from tabulate import tabulate
+
+from datahub.cli.cli_utils import fixup_gms_url, generate_access_token
+from datahub.cli.config_utils import load_client_config, persist_raw_datahub_config
+from datahub.ingestion.graph.config import DatahubClientConfig
+
+DATAHUB_SESSIONS_PATH = os.path.expanduser("~/.datahub/sessions.json")
+
+
+def load_sessions() -> Dict[str, DatahubClientConfig]:
+    if not os.path.exists(DATAHUB_SESSIONS_PATH):
+        return {}
+    with open(DATAHUB_SESSIONS_PATH, "r") as f:
+        raw_sessions = json.load(f)
+        return {
+            name: DatahubClientConfig.parse_obj(config)
+            for name, config in raw_sessions.items()
+        }
+
+
+def save_sessions(sessions: Dict[str, DatahubClientConfig]) -> None:
+    os.makedirs(os.path.dirname(DATAHUB_SESSIONS_PATH), exist_ok=True)
+    with open(DATAHUB_SESSIONS_PATH, "w") as f:
+        json.dump(
+            {name: config.dict() for name, config in sessions.items()}, f, indent=2
+        )
+
+
+@click.group()
+def session() -> None:
+    """Manage DataHub session profiles"""
+    pass
+
+
+@session.command()
+@click.option(
+    "--use-password",
+    type=bool,
+    is_flag=True,
+    default=False,
+    help="If passed then uses password to initialise token.",
+)
+def create(use_password: bool) -> None:
+    """Create profile with which to connect to a DataHub instance"""
+
+    sessions = load_sessions()
+
+    click.echo(
+        "Configure which datahub instance to connect to (https://your-instance.acryl.io/gms for Acryl hosted users)"
+    )
+    host = click.prompt(
+        "Enter your DataHub host", type=str, default="http://localhost:8080"
+    )
+    host = fixup_gms_url(host)
+    if use_password:
+        username = click.prompt("Enter your DataHub username", type=str)
+        password = click.prompt(
+            "Enter your DataHub password",
+            type=str,
+        )
+        _, token = generate_access_token(
+            username=username, password=password, gms_url=host
+        )
+    else:
+        token = click.prompt(
+            "Enter your DataHub access token",
+            type=str,
+            default="",
+        )
+
+    profile_name = click.prompt("Enter name for profile", type=str)
+
+    config = DatahubClientConfig(server=host, token=token)
+    sessions[profile_name] = config
+    save_sessions(sessions)
+    click.echo(f"Created profile: {profile_name}")
+
+
+@session.command()
+def list() -> None:
+    """List all session profiles"""
+    sessions = load_sessions()
+    if not sessions:
+        click.echo("No profiles found")
+        return
+
+    headers = ["Profile", "URL"]
+    table_data = [[name, config.server] for name, config in sessions.items()]
+    click.echo(tabulate(table_data, headers=headers))
+
+
+@session.command()
+@click.argument("profile", type=str)
+def delete(profile: str) -> None:
+    """Delete a session profile"""
+    sessions = load_sessions()
+    if profile not in sessions:
+        click.echo(f"Profile {profile} not found")
+        return
+
+    del sessions[profile]
+    save_sessions(sessions)
+    click.echo(f"Deleted profile: {profile}")
+
+
+@session.command()
+@click.argument("profile", type=str)
+def use(profile: str) -> None:
+    """Set the active session"""
+    sessions = load_sessions()
+    session = sessions.get(profile)
+    if session:
+        persist_raw_datahub_config(session.dict())
+        click.echo(f"Using profile {profile}")
+    else:
+        click.echo(f"Profile {profile} not found")
+        return
+
+
+@session.command()
+@click.option(
+    "--profile",
+    type=str,
+    required=True,
+    help="Name of profile under which to save the current datahubenv config",
+)
+def save(profile: str) -> None:
+    """Save the current active datahubenv config as a session"""
+    sessions = load_sessions()
+    if profile in sessions:
+        click.echo(
+            f"Profile {profile} already exists, please make sure to use a unique profile name"
+        )
+        return
+
+    config = load_client_config()
+    sessions[profile] = config
+    save_sessions(sessions)
+    click.echo(f"Saved current datahubenv as profile: {profile}")

metadata-ingestion/src/datahub/entrypoints.py

@@ -26,6 +26,7 @@
 from datahub.cli.ingest_cli import ingest
 from datahub.cli.migrate import migrate
 from datahub.cli.put_cli import put
+from datahub.cli.session_cli import session
 from datahub.cli.specific.assertions_cli import assertions
 from datahub.cli.specific.datacontract_cli import datacontract
 from datahub.cli.specific.dataproduct_cli import dataproduct
@@ -185,6 +186,7 @@ def init(use_password: bool = False) -> None:
 datahub.add_command(datacontract)
 datahub.add_command(assertions)
 datahub.add_command(container)
+datahub.add_command(session)
 
 try:
     from datahub.cli.iceberg_cli import iceberg

metadata-ingestion/tests/unit/cli/test_cli_session.py

@@ -0,0 +1,163 @@
+import json
+import os
+from unittest.mock import patch
+
+import pytest
+from click.testing import CliRunner
+
+from datahub.cli.session_cli import session
+from datahub.ingestion.graph.config import DatahubClientConfig
+
+
+@pytest.fixture
+def mock_sessions_file(tmp_path):
+    # Create a temporary sessions file for testing
+    test_sessions_file = tmp_path / "sessions.json"
+    with patch(
+        "datahub.cli.session_cli.DATAHUB_SESSIONS_PATH", str(test_sessions_file)
+    ):
+        yield test_sessions_file
+    # Cleanup
+    if test_sessions_file.exists():
+        test_sessions_file.unlink()
+
+
+@pytest.fixture
+def sample_config():
+    return DatahubClientConfig(server="http://localhost:8080", token="test-token")
+
+
+def test_create_session(mock_sessions_file):
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        # Test creating a new profile
+        result = runner.invoke(
+            session,
+            ["create"],
+            input="http://localhost:8080\ntest-token\ntest-profile\n",
+        )
+        assert result.exit_code == 0
+        assert "Created profile: test-profile" in result.output
+
+        # Verify the session was saved
+        with open(mock_sessions_file) as f:
+            saved_sessions = json.load(f)
+            assert "test-profile" in saved_sessions
+            assert saved_sessions["test-profile"]["server"] == "http://localhost:8080"
+            assert saved_sessions["test-profile"]["token"] == "test-token"
+
+
+def test_list_sessions_empty_data(mock_sessions_file):
+    runner = CliRunner()
+    # Create an empty sessions file
+    os.makedirs(os.path.dirname(mock_sessions_file), exist_ok=True)
+    with open(mock_sessions_file, "w") as f:
+        json.dump({}, f)
+
+    result = runner.invoke(session, ["list"])
+    assert result.exit_code == 0
+    assert "No profiles found" in result.output
+
+
+def test_list_sessions(mock_sessions_file, sample_config):
+    runner = CliRunner()
+    # Setup test data
+    sessions_data = {"test-profile": sample_config.dict()}
+    os.makedirs(os.path.dirname(mock_sessions_file), exist_ok=True)
+    with open(mock_sessions_file, "w") as f:
+        json.dump(sessions_data, f)
+
+    result = runner.invoke(session, ["list"])
+    assert result.exit_code == 0
+    assert "test-profile" in result.output
+    assert "http://localhost:8080" in result.output
+
+
+def test_delete_session(mock_sessions_file, sample_config):
+    runner = CliRunner()
+    # Setup test data
+    sessions_data = {"test-profile": sample_config.dict()}
+    os.makedirs(os.path.dirname(mock_sessions_file), exist_ok=True)
+    with open(mock_sessions_file, "w") as f:
+        json.dump(sessions_data, f)
+
+    # Test deleting existing profile
+    result = runner.invoke(session, ["delete", "test-profile"])
+    assert result.exit_code == 0
+    assert "Deleted profile: test-profile" in result.output
+
+    # Verify profile was deleted
+    with open(mock_sessions_file) as f:
+        saved_sessions = json.load(f)
+        assert "test-profile" not in saved_sessions
+
+
+def test_delete_unknown_session(mock_sessions_file, sample_config):
+    runner = CliRunner()
+    # Setup test data with a known profile
+    sessions_data = {"test-profile": sample_config.dict()}
+    os.makedirs(os.path.dirname(mock_sessions_file), exist_ok=True)
+    with open(mock_sessions_file, "w") as f:
+        json.dump(sessions_data, f)
+
+    # Test deleting non-existent profile
+    result = runner.invoke(session, ["delete", "unknown-profile"])
+    assert result.exit_code == 0
+    assert "Profile unknown-profile not found" in result.output
+
+    # Verify original profile still exists
+    with open(mock_sessions_file) as f:
+        saved_sessions = json.load(f)
+        assert "test-profile" in saved_sessions
+
+
+def test_use_session(mock_sessions_file, sample_config):
+    runner = CliRunner()
+    # Setup test data
+    sessions_data = {"test-profile": sample_config.dict()}
+    os.makedirs(os.path.dirname(mock_sessions_file), exist_ok=True)
+    with open(mock_sessions_file, "w") as f:
+        json.dump(sessions_data, f)
+
+    with patch("datahub.cli.session_cli.persist_raw_datahub_config") as mock_persist:
+        result = runner.invoke(session, ["use", "test-profile"])
+        assert result.exit_code == 0
+        assert "Using profile test-profile" in result.output
+        mock_persist.assert_called_once_with(sample_config.dict())
+
+
+def test_save_session(mock_sessions_file, sample_config):
+    runner = CliRunner()
+    with patch(
+        "datahub.cli.session_cli.load_client_config", return_value=sample_config
+    ):
+        result = runner.invoke(session, ["save", "--profile", "new-profile"])
+        assert result.exit_code == 0
+        assert "Saved current datahubenv as profile: new-profile" in result.output
+
+        # Verify the session was saved
+        with open(mock_sessions_file) as f:
+            saved_sessions = json.load(f)
+            assert "new-profile" in saved_sessions
+            assert saved_sessions["new-profile"]["server"] == "http://localhost:8080"
+            assert saved_sessions["new-profile"]["token"] == "test-token"
+
+
+def test_create_session_with_password(mock_sessions_file):
+    runner = CliRunner()
+    with runner.isolated_filesystem(), patch(
+        "datahub.cli.session_cli.generate_access_token",
+        return_value=("username", "generated-token"),
+    ):
+        result = runner.invoke(
+            session,
+            ["create", "--use-password"],
+            input="http://localhost:8080\nusername\npassword\ntest-profile\n",
+        )
+        assert result.exit_code == 0
+        assert "Created profile: test-profile" in result.output
+
+        # Verify the session was saved with generated token
+        with open(mock_sessions_file) as f:
+            saved_sessions = json.load(f)
+            assert saved_sessions["test-profile"]["token"] == "generated-token"