Add sgr cloud ephemeral-tunnel command

Author: neumark

splitgraph/cloud/__init__.py

@@ -61,6 +61,7 @@
     INGESTION_JOB_STATUS,
     JOB_LOGS,
     PROFILE_UPSERT,
+    PROVISION_EPHEMERAL_TUNNEL,
     PROVISION_REPOSITORY_TUNNEL,
     REPO_CONDITIONS,
     REPO_PARAMS,
@@ -1121,3 +1122,21 @@ def provision_repository_tunnel(self, namespace: str, repository: str) -> Tuple[
             response.json()["data"]["provisionRepositoryTunnel"]["tunnelConnectHost"],
             response.json()["data"]["provisionRepositoryTunnel"]["tunnelConnectPort"],
         )
+
+    def provision_ephemeral_tunnel(self) -> Tuple[str, str, int, str, int]:
+        response = self._gql(
+            {
+                "query": PROVISION_EPHEMERAL_TUNNEL,
+                "operationName": "ProvisionEphemeralTunnel",
+                "variables": {},
+            },
+            handle_errors=True,
+            anonymous_ok=False,
+        )
+        return (
+            response.json()["data"]["provisionEphemeralTunnel"]["secretToken"],
+            response.json()["data"]["provisionEphemeralTunnel"]["tunnelConnectHost"],
+            response.json()["data"]["provisionEphemeralTunnel"]["tunnelConnectPort"],
+            response.json()["data"]["provisionEphemeralTunnel"]["privateAddressHost"],
+            response.json()["data"]["provisionEphemeralTunnel"]["privateAddressPort"],
+        )

splitgraph/cloud/queries.py

@@ -289,6 +289,17 @@
 }
 """
 
+PROVISION_EPHEMERAL_TUNNEL = """mutation ProvisionEphemeralTunnel {
+  provisionEphemeralTunnel {
+    secretToken,
+    tunnelConnectHost,
+    tunnelConnectPort,
+    privateAddressHost,
+    privateAddressPort
+  }
+}
+"""
+
 EXPORT_JOB_STATUS = """query ExportJobStatus($taskId: UUID!) {
   exportJobStatus(taskId: $taskId) {
     taskId

splitgraph/cloud/tunnel_client.py

@@ -4,13 +4,15 @@
 from os import path
 from typing import IO, Any, Dict, Optional, cast
 
+from splitgraph.config import CONFIG
+from splitgraph.config.config import get_singleton
 from splitgraph.core.repository import Repository
 
 RATHOLE_CLIENT_CONFIG_FILENAME = "rathole-client.toml"
 
 RATHOLE_CLIENT_CONFIG_TEMPLATE = """
 [client]
-remote_addr = "{tunnel_server_management_address}"
+remote_addr = "{tunnel_connect_address}"
 
 [client.transport]
 type = "tls"
@@ -19,55 +21,56 @@
 {trusted_root_line}
 hostname = "{tls_hostname}"
 
-[client.services."{namespace}/{repository}"]
+[client.services."{section_id}"]
 local_addr = "{local_address}"
-# token is provisioner JWT token
-token = "{provisioning_token}"
+token = "{secret_token}"
 
 """
 
 
 def get_rathole_client_config(
-    tunnel_server_management_address: str,
+    tunnel_connect_address: str,
     tls_hostname: str,
     local_address: str,
-    provisioning_token: str,
-    namespace: str,
-    repository: str,
+    secret_token: str,
+    section_id: str,
     trusted_root: Optional[str],
 ) -> str:
     trusted_root_line = f'trusted_root = "{trusted_root}"' if trusted_root else ""
     return RATHOLE_CLIENT_CONFIG_TEMPLATE.format(
-        tunnel_server_management_address=tunnel_server_management_address,
+        tunnel_connect_address=tunnel_connect_address,
         tls_hostname=tls_hostname,
         local_address=local_address,
-        provisioning_token=provisioning_token,
-        namespace=namespace,
-        repository=repository,
+        secret_token=secret_token,
+        section_id=section_id,
         trusted_root_line=trusted_root_line,
     )
 
 
+def get_rathole_client_binary_path():
+    config_dir = os.path.dirname(get_singleton(CONFIG, "SG_CONFIG_FILE"))
+    return os.path.join(config_dir, "rathole")
+
+
 def write_rathole_client_config(
-    provisioning_token: str,
-    tunnel_server_management_host: str,
-    tunnel_server_management_port: int,
+    section_id: str,
+    secret_token: str,
+    tunnel_connect_host: str,
+    tunnel_connect_port: int,
+    local_address: str,
     tls_hostname: Optional[str],
-    repository: Repository,
-    params: Dict[str, Any],
-    config_dir: str,
 ) -> str:
     # in production, this will be None, but for dev instances, we need to
     # specify rootCA.pem
+    config_dir = os.path.dirname(get_singleton(CONFIG, "SG_CONFIG_FILE"))
     trusted_root = os.environ.get("REQUESTS_CA_BUNDLE") or os.environ.get("SSL_CERT_FILE")
     rathole_client_config = get_rathole_client_config(
         # TODO: replace these stub values with response of provisioning call
-        tunnel_server_management_address=f"{tunnel_server_management_host}:{tunnel_server_management_port}",
-        tls_hostname=tls_hostname or tunnel_server_management_host,
-        local_address=f"{params['host']}:{params['port']}",
-        provisioning_token=provisioning_token,
-        namespace=repository.namespace,
-        repository=repository.repository,
+        tunnel_connect_address=f"{tunnel_connect_host}:{tunnel_connect_port}",
+        tls_hostname=tls_hostname or tunnel_connect_host,
+        local_address=local_address,
+        secret_token=secret_token,
+        section_id=section_id,
         trusted_root=trusted_root,
     )
     config_filename = path.join(config_dir, RATHOLE_CLIENT_CONFIG_FILENAME)

splitgraph/commandline/cloud.py

@@ -9,7 +9,7 @@
 from copy import copy
 from glob import glob
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, OrderedDict, Tuple, cast
 from urllib.parse import quote, urlparse
 
 import click
@@ -18,6 +18,7 @@
 from splitgraph.cloud.models import AddExternalRepositoryRequest, IntrospectionMode
 from splitgraph.cloud.project.models import Metadata, SplitgraphYAML
 from splitgraph.cloud.tunnel_client import (
+    get_rathole_client_binary_path,
     launch_rathole_client,
     write_rathole_client_config,
 )
@@ -1239,29 +1240,65 @@ def tunnel_c(remote: str, repositories_file: List[Path], repository: "CoreReposi
             f"Repository {repository.namespace}/{repository.repository} is not tunneled"
         )
 
-    config_dir = os.path.dirname(get_singleton(CONFIG, "SG_CONFIG_FILE"))
     # TODO: Get current version of rathole client for architecture.
     # user must manually download rathole for now
-    rathole_client_binary_path = os.path.join(config_dir, "rathole")
 
     from splitgraph.cloud import GQLAPIClient
 
     client = GQLAPIClient(remote)
-    (secretToken, tunnelConnectHost, tunnelConnectPort) = client.provision_repository_tunnel(
+    (secret_token, tunnel_connect_host, tunnel_connect_port) = client.provision_repository_tunnel(
         repository.namespace, repository.repository
     )
 
+    section_id = f"{repository.namespace}/{repository.repository}"
+    local_address = f"{external.params['host']}:{external.params['port']}"
+
+    rathole_client_config_path = write_rathole_client_config(
+        section_id,
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        local_address,
+        tunnel_connect_host,
+    )
+
+    print("launching rathole client")
+    launch_rathole_client(get_rathole_client_binary_path(), rathole_client_config_path)
+
+
+@click.command("ephemeral-tunnel")
+@click.option("--remote", default="data.splitgraph.com", help="Name of the remote registry to use.")
+@click.argument("host", type=str)
+@click.argument("port", type=int)
+def ephemeral_tunnel_c(remote: str, host: str, port: int):
+    """
+    Start ephemeral tunnel.
+    """
+
+    from splitgraph.cloud import GQLAPIClient
+
+    client = GQLAPIClient(remote)
+    (
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        private_address_host,
+        private_address_port,
+    ) = client.provision_ephemeral_tunnel()
+
     rathole_client_config_path = write_rathole_client_config(
-        secretToken,
-        tunnelConnectHost,
-        tunnelConnectPort,
-        tunnelConnectHost,
-        repository,
-        external.params,
-        config_dir,
+        private_address_host,
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        f"{host}:{port}",
+        tunnel_connect_host,
+    )
+    print(
+        f"To connect to {host}:{port} from Splitgraph, use the following connection parameters:\nHost: {private_address_host}\nPort: {private_address_port}"
     )
     print("launching rathole client")
-    launch_rathole_client(rathole_client_binary_path, rathole_client_config_path)
+    launch_rathole_client(get_rathole_client_binary_path(), rathole_client_config_path)
 
 
 @click.group("cloud")
@@ -1292,3 +1329,4 @@ def cloud_c():
 cloud_c.add_command(validate_c)
 cloud_c.add_command(seed_c)
 cloud_c.add_command(tunnel_c)
+cloud_c.add_command(ephemeral_tunnel_c)