[App] Enable to register data connections (#16670)

Co-authored-by: thomas <thomas@thomass-MacBook-Pro.local>

pyproject.toml

@@ -105,13 +105,14 @@ module = [
     "lightning.app.api.http_methods",
     "lightning.app.api.request_types",
     "lightning.app.cli.commands.app_commands",
-    "lightning.app.cli.commands.connection",
     "lightning.app.cli.commands.lightning_cli",
     "lightning.app.cli.commands.cmd_install",
+    "lightning.app.cli.connect.app",
     "lightning.app.cli.commands.pwd",
     "lightning.app.cli.commands.ls",
     "lightning.app.cli.commands.cp",
     "lightning.app.cli.commands.cd",
+    "lightning.app.cli.connect.data",
     "lightning.app.cli.cmd_install",
     "lightning.app.components.database.client",
     "lightning.app.components.database.server",

requirements/app/base.txt

@@ -1,4 +1,4 @@
-lightning-cloud>=0.5.22
+lightning-cloud>=0.5.24
 packaging
 typing-extensions>=4.0.0, <=4.4.0
 deepdiff>=5.7.0, <6.2.4

src/lightning/app/CHANGELOG.md

@@ -22,6 +22,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
  * `pwd`: Return the current folder in your Cloud Platform Filesystem
  * `cp`: Copy files between your Cloud Platform Filesystem and local filesystem
 
+- Added `lightning connect data` to register data connection to s3 buckets ([#16670](https://github.com/Lightning-AI/lightning/pull/16670))
+
+
 ### Changed
 
 - Changed the default `LightningClient(retry=False)` to `retry=True` ([#16382](https://github.com/Lightning-AI/lightning/pull/16382))
@@ -33,6 +36,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Renamed `lightning.app.components.LiteMultiNode` to `lightning.app.components.FabricMultiNode` ([#16505](https://github.com/Lightning-AI/lightning/pull/16505))
 
 
+- Changed the command `lightning connect` to `lightning connect app` for consistency ([#16670](https://github.com/Lightning-AI/lightning/pull/16670))
+
+
 ### Deprecated
 
 -

src/lightning/app/cli/commands/app_commands.py

@@ -18,7 +18,7 @@ from typing import Dict, Optional
 
 import requests
 
-from lightning.app.cli.commands.connection import (
+from lightning.app.cli.connect.app import (
     _clean_lightning_connection,
     _install_missing_requirements,
     _resolve_command_path,

src/lightning/app/cli/commands/cd.py

@@ -21,7 +21,7 @@ from rich.spinner import Spinner
 from rich.text import Text
 
 from lightning.app.cli.commands import ls
-from lightning.app.cli.commands.connection import _LIGHTNING_CONNECTION_FOLDER
+from lightning.app.cli.connect.app import _LIGHTNING_CONNECTION_FOLDER
 from lightning.app.utilities.app_helpers import Logger
 from lightning.app.utilities.cli_helpers import _error_and_exit
 

src/lightning/app/cli/commands/ls.py

@@ -26,7 +26,7 @@ from rich.live import Live
 from rich.spinner import Spinner
 from rich.text import Text
 
-from lightning.app.cli.commands.connection import _LIGHTNING_CONNECTION_FOLDER
+from lightning.app.cli.connect.app import _LIGHTNING_CONNECTION_FOLDER
 from lightning.app.utilities.app_helpers import Logger
 from lightning.app.utilities.cli_helpers import _error_and_exit
 from lightning.app.utilities.network import LightningClient
@@ -44,8 +44,7 @@ def ls(path: Optional[str] = None, print: bool = True, use_live: bool = True) ->
     from lightning.app.cli.commands.cd import _CD_FILE
 
     if sys.platform == "win32":
-        print("`ls` isn't supported on windows. Open an issue on Github.")
-        sys.exit(0)
+        _error_and_exit("`ls` isn't supported on windows. Open an issue on Github.")
 
     root = "/"
 

src/lightning/app/cli/connect/app.py

@@ -37,7 +37,7 @@ _LIGHTNING_CONNECTION_FOLDER = os.path.join(_LIGHTNING_CONNECTION, _PPID)
 
 
 @click.argument("app_name_or_id", required=True)
-def connect(app_name_or_id: str):
+def connect_app(app_name_or_id: str):
     """Connect your local terminal to a running lightning app.
 
     After connecting, the lightning CLI will respond to commands exposed by the app.
@@ -79,8 +79,8 @@ def connect(app_name_or_id: str):
             else:
                 click.echo(f"You are already connected to the cloud Lightning App: {app_name_or_id}.")
         else:
-            disconnect()
-            connect(app_name_or_id)
+            disconnect_app()
+            connect_app(app_name_or_id)
 
     elif app_name_or_id.startswith("localhost"):
 
@@ -217,7 +217,7 @@ def connect(app_name_or_id: str):
         ).wait()
 
 
-def disconnect(logout: bool = False):
+def disconnect_app(logout: bool = False):
     """Disconnect from an App."""
     _clean_lightning_connection()
 

src/lightning/app/cli/connect/data.py

@@ -0,0 +1,92 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+import click
+import rich
+from lightning_cloud.openapi import ProjectIdDataConnectionsBody
+from rich.live import Live
+from rich.spinner import Spinner
+from rich.text import Text
+
+from lightning.app.utilities.app_helpers import Logger
+from lightning.app.utilities.cli_helpers import _error_and_exit
+from lightning.app.utilities.cloud import _get_project
+from lightning.app.utilities.network import LightningClient
+
+logger = Logger(__name__)
+
+
+@click.argument("name", required=True)
+@click.argument("region", required=True)
+@click.argument("source", required=True)
+@click.argument("destination", required=False)
+@click.argument("project_name", required=False)
+def connect_data(
+    name: str,
+    region: str,
+    source: str,
+    destination: str = "",
+    project_name: str = "",
+) -> None:
+    """Create a new data connection."""
+
+    if sys.platform == "win32":
+        _error_and_exit("Data connection isn't supported on windows. Open an issue on Github.")
+
+    with Live(Spinner("point", text=Text("pending...", style="white")), transient=True) as live:
+
+        live.stop()
+
+        client = LightningClient()
+        projects = client.projects_service_list_memberships()
+
+        project_id = None
+
+        for project in projects.memberships:
+
+            if project.name == project_name:
+                project_id = project.project_id
+                break
+
+        if project_id is None:
+            project_id = _get_project(client).project_id
+
+        if not source.startswith("s3://"):
+            return _error_and_exit(
+                "Only public S3 folders are supported for now. Please, open a Github issue with your use case."
+            )
+
+        try:
+            _ = client.data_connection_service_create_data_connection(
+                body=ProjectIdDataConnectionsBody(
+                    name=name,
+                    region=region,
+                    source=source,
+                    destination=destination,
+                ),
+                project_id=project_id,
+            )
+
+            # Note: Expose through lightning show data {DATA_NAME}
+            # response = client.data_connection_service_list_data_connection_artifacts(
+            #     project_id=project_id,
+            #     id=response.id,
+            # )
+            # print(response)
+        except Exception:
+            _error_and_exit("The data connection creation failed.")
+
+    rich.print(f"[green]Succeeded[/green]: You have created a new data connection {name}.")

src/lightning/app/cli/lightning_cli.py

@@ -34,16 +34,17 @@ from lightning.app.cli.cmd_apps import _AppManager
 from lightning.app.cli.cmd_clusters import AWSClusterManager
 from lightning.app.cli.commands.app_commands import _run_app_command
 from lightning.app.cli.commands.cd import cd
-from lightning.app.cli.commands.connection import (
-    _list_app_commands,
-    _retrieve_connection_to_an_app,
-    connect,
-    disconnect,
-)
 from lightning.app.cli.commands.cp import cp
 from lightning.app.cli.commands.logs import logs
 from lightning.app.cli.commands.ls import ls
 from lightning.app.cli.commands.pwd import pwd
+from lightning.app.cli.connect.app import (
+    _list_app_commands,
+    _retrieve_connection_to_an_app,
+    connect_app,
+    disconnect_app,
+)
+from lightning.app.cli.connect.data import connect_data
 from lightning.app.cli.lightning_cli_create import create
 from lightning.app.cli.lightning_cli_delete import delete
 from lightning.app.cli.lightning_cli_list import get_list
@@ -121,8 +122,21 @@ def show() -> None:
     pass
 
 
-_main.command()(connect)
-_main.command()(disconnect)
+@_main.group()
+def connect() -> None:
+    """Connect apps and data."""
+    pass
+
+
+@_main.group()
+def disconnect() -> None:
+    """Disconnect apps."""
+    pass
+
+
+connect.command("app")(connect_app)
+disconnect.command("app")(disconnect_app)
+connect.command("data")(connect_data)
 _main.command()(ls)
 _main.command()(cd)
 _main.command()(cp)
@@ -240,7 +254,7 @@ def login() -> None:
 def logout() -> None:
     """Log out of your lightning.ai account."""
     Auth().clear()
-    disconnect(logout=True)
+    disconnect_app(logout=True)
 
 
 def _run_app(

tests/integrations_app/public/test_commands_and_api.py

@@ -19,11 +19,11 @@ def test_commands_and_api_example_cloud() -> None:
     ):
         # Connect to the App and send the first & second command with the client
         # Requires to be run within the same process.
-        cmd_1 = f"python -m lightning connect {app_name}"
+        cmd_1 = f"python -m lightning connect app {app_name}"
         cmd_2 = "python -m lightning command with client --name=this"
         cmd_3 = "python -m lightning command without client --name=is"
         cmd_4 = "python -m lightning command without client --name=awesome"
-        cmd_5 = "lightning disconnect"
+        cmd_5 = "lightning disconnect app"
         process = Popen(" && ".join([cmd_1, cmd_2, cmd_3, cmd_4, cmd_5]), shell=True)
         process.wait()
 

tests/tests_app/cli/test_connect.py

@@ -7,12 +7,12 @@ import psutil
 import pytest
 
 from lightning.app import _PROJECT_ROOT
-from lightning.app.cli.commands.connection import (
+from lightning.app.cli.connect.app import (
     _list_app_commands,
     _resolve_command_path,
     _retrieve_connection_to_an_app,
-    connect,
-    disconnect,
+    connect_app,
+    disconnect_app,
 )
 from lightning.app.utilities import cli_helpers
 from lightning.app.utilities.commands import base
@@ -20,18 +20,18 @@ from lightning.app.utilities.commands import base
 
 def monkeypatch_connection(monkeypatch, tmpdir, ppid):
     connection_path = os.path.join(tmpdir, ppid)
-    monkeypatch.setattr("lightning.app.cli.commands.connection._clean_lightning_connection", MagicMock())
-    monkeypatch.setattr("lightning.app.cli.commands.connection._PPID", ppid)
-    monkeypatch.setattr("lightning.app.cli.commands.connection._LIGHTNING_CONNECTION", tmpdir)
-    monkeypatch.setattr("lightning.app.cli.commands.connection._LIGHTNING_CONNECTION_FOLDER", connection_path)
+    monkeypatch.setattr("lightning.app.cli.connect.app._clean_lightning_connection", MagicMock())
+    monkeypatch.setattr("lightning.app.cli.connect.app._PPID", ppid)
+    monkeypatch.setattr("lightning.app.cli.connect.app._LIGHTNING_CONNECTION", tmpdir)
+    monkeypatch.setattr("lightning.app.cli.connect.app._LIGHTNING_CONNECTION_FOLDER", connection_path)
     return connection_path
 
 
 def test_connect_disconnect_local(tmpdir, monkeypatch):
-    disconnect()
+    disconnect_app()
 
     with pytest.raises(Exception, match="Connection wasn't successful. Is your app localhost running ?"):
-        connect("localhost")
+        connect_app("localhost")
 
     with open(os.path.join(os.path.dirname(__file__), "jsons/connect_1.json")) as f:
         data = json.load(f)
@@ -43,7 +43,7 @@ def test_connect_disconnect_local(tmpdir, monkeypatch):
 
     messages = []
 
-    disconnect()
+    disconnect_app()
 
     def fn(msg):
         messages.append(msg)
@@ -54,21 +54,21 @@ def test_connect_disconnect_local(tmpdir, monkeypatch):
     response.status_code = 200
     response.json.return_value = data
     monkeypatch.setattr(cli_helpers.requests, "get", MagicMock(return_value=response))
-    connect("localhost")
+    connect_app("localhost")
     assert _retrieve_connection_to_an_app() == ("localhost", None)
     command_path = _resolve_command_path("nested_command")
     assert not os.path.exists(command_path)
     command_path = _resolve_command_path("command_with_client")
     assert os.path.exists(command_path)
     messages = []
-    connect("localhost")
+    connect_app("localhost")
     assert messages == ["You are connected to the local Lightning App."]
 
     messages = []
-    disconnect()
+    disconnect_app()
     assert messages == ["You are disconnected from the local Lightning App."]
     messages = []
-    disconnect()
+    disconnect_app()
     assert messages == [
         "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one."
     ]
@@ -77,7 +77,7 @@ def test_connect_disconnect_local(tmpdir, monkeypatch):
 
 
 def test_connect_disconnect_cloud(tmpdir, monkeypatch):
-    disconnect()
+    disconnect_app()
 
     ppid_1 = str(psutil.Process(os.getpid()).ppid())
     ppid_2 = "222"
@@ -132,7 +132,7 @@ def test_connect_disconnect_cloud(tmpdir, monkeypatch):
     with open(data["paths"]["/command/command_with_client"]["post"]["cls_path"], "rb") as f:
         response.content = f.read()
 
-    connect("example")
+    connect_app("example")
     assert _retrieve_connection_to_an_app() == ("example", "1234")
     commands = _list_app_commands()
     assert commands == ["command with client", "command without client", "nested command"]
@@ -141,27 +141,27 @@ def test_connect_disconnect_cloud(tmpdir, monkeypatch):
     command_path = _resolve_command_path("command_with_client")
     assert os.path.exists(command_path)
     messages = []
-    connect("example")
+    connect_app("example")
     assert messages == ["You are already connected to the cloud Lightning App: example."]
 
     _ = monkeypatch_connection(monkeypatch, tmpdir, ppid=ppid_2)
 
     messages = []
-    connect("example")
+    connect_app("example")
     assert "The lightning CLI now responds to app commands" in messages[0]
 
     messages = []
-    disconnect()
+    disconnect_app()
     assert messages == ["You are disconnected from the cloud Lightning App: example."]
 
     _ = monkeypatch_connection(monkeypatch, tmpdir, ppid=ppid_1)
 
     messages = []
-    disconnect()
+    disconnect_app()
     assert "You aren't connected to any Lightning App" in messages[0]
 
     messages = []
-    disconnect()
+    disconnect_app()
     assert messages == [
         "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one."
     ]

tests/tests_app/cli/test_connect_data.py

@@ -0,0 +1,60 @@
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+from lightning_cloud.openapi import ProjectIdDataConnectionsBody, V1ListMembershipsResponse, V1Membership
+
+from lightning.app.cli.connect import data
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="not supported on windows yet")
+def test_connect_data_no_project(monkeypatch):
+
+    client = MagicMock()
+    client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(memberships=[])
+    monkeypatch.setattr(data, "LightningClient", MagicMock(return_value=client))
+
+    _error_and_exit = MagicMock()
+    monkeypatch.setattr(data, "_error_and_exit", _error_and_exit)
+
+    _get_project = MagicMock()
+    _get_project.return_value = V1Membership(name="project-0", project_id="project-id-0")
+    monkeypatch.setattr(data, "_get_project", _get_project)
+
+    data.connect_data("imagenet", region="us-east-1", source="imagenet", destination="", project_name="project-0")
+
+    _get_project.assert_called()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="not supported on windows yet")
+def test_connect_data(monkeypatch):
+
+    client = MagicMock()
+    client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
+        memberships=[
+            V1Membership(name="project-0", project_id="project-id-0"),
+            V1Membership(name="project-1", project_id="project-id-1"),
+            V1Membership(name="project 2", project_id="project-id-2"),
+        ]
+    )
+    monkeypatch.setattr(data, "LightningClient", MagicMock(return_value=client))
+
+    _error_and_exit = MagicMock()
+    monkeypatch.setattr(data, "_error_and_exit", _error_and_exit)
+    data.connect_data("imagenet", region="us-east-1", source="imagenet", destination="", project_name="project-0")
+
+    _error_and_exit.assert_called_with(
+        "Only public S3 folders are supported for now. Please, open a Github issue with your use case."
+    )
+
+    data.connect_data("imagenet", region="us-east-1", source="s3://imagenet", destination="", project_name="project-0")
+
+    client.data_connection_service_create_data_connection.assert_called_with(
+        project_id="project-id-0",
+        body=ProjectIdDataConnectionsBody(
+            destination="",
+            region="us-east-1",
+            name="imagenet",
+            source="s3://imagenet",
+        ),
+    )

tests/tests_app/utilities/test_commands.py

@@ -10,7 +10,7 @@ from pydantic import BaseModel
 
 from lightning.app import LightningApp, LightningFlow
 from lightning.app.cli.commands.app_commands import _run_app_command
-from lightning.app.cli.commands.connection import connect, disconnect
+from lightning.app.cli.connect.app import connect_app, disconnect_app
 from lightning.app.core.constants import APP_SERVER_PORT
 from lightning.app.runners import MultiProcessRuntime
 from lightning.app.testing.helpers import _RunIf
@@ -145,7 +145,7 @@ def test_configure_commands(monkeypatch):
 
     sleep(0.5)
     monkeypatch.setattr(sys, "argv", ["lightning", "user", "command", "--name=something"])
-    connect("localhost")
+    connect_app("localhost")
     _run_app_command("localhost", None)
     sleep(0.5)
     state = AppState()
@@ -160,5 +160,5 @@ def test_configure_commands(monkeypatch):
         sleep(0.1)
         time_left -= 0.1
     assert process.exitcode == 0
-    disconnect()
+    disconnect_app()
     process.kill()