h2o-connector-service 0.1.0.dev6001__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (80) hide show
  1. h2o_connector_service-0.1.0.dev6001/PKG-INFO +28 -0
  2. h2o_connector_service-0.1.0.dev6001/pyproject.toml +42 -0
  3. h2o_connector_service-0.1.0.dev6001/setup.cfg +4 -0
  4. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/__init__.py +35 -0
  5. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/_base.py +141 -0
  6. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/browse_storage_client.py +82 -0
  7. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/connection_client.py +178 -0
  8. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/connector_client.py +85 -0
  9. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/connector_service_facade.py +160 -0
  10. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/connector_session.py +520 -0
  11. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/connector_template_client.py +89 -0
  12. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/data_source_profile_client.py +35 -0
  13. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/exceptions.py +88 -0
  14. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/__init__.py +0 -0
  15. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/__init__.pyi +1 -0
  16. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/common_pb2.py +117 -0
  17. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/common_pb2_grpc.py +2 -0
  18. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connection_pb2.py +118 -0
  19. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connection_pb2_grpc.py +2 -0
  20. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connection_service_pb2.py +162 -0
  21. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connection_service_pb2_grpc.py +200 -0
  22. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_pb2.py +57 -0
  23. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_pb2_grpc.py +2 -0
  24. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_service_pb2.py +74 -0
  25. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_service_pb2_grpc.py +102 -0
  26. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_template_pb2.py +50 -0
  27. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_template_pb2_grpc.py +2 -0
  28. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_template_service_pb2.py +72 -0
  29. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/connector_template_service_pb2_grpc.py +108 -0
  30. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_pb2.py +168 -0
  31. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_pb2_grpc.py +2 -0
  32. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_source_profile_pb2.py +67 -0
  33. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_source_profile_pb2_grpc.py +2 -0
  34. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_source_profile_service_pb2.py +38 -0
  35. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_source_profile_service_pb2_grpc.py +68 -0
  36. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_stream_service_pb2.py +50 -0
  37. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/data_stream_service_pb2_grpc.py +95 -0
  38. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/__init__.py +0 -0
  39. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/__init__.pyi +1 -0
  40. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/__init__.py +0 -0
  41. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/__init__.pyi +0 -0
  42. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/annotations_pb2.py +37 -0
  43. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/client_pb2.py +90 -0
  44. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/field_behavior_pb2.py +40 -0
  45. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/http_pb2.py +41 -0
  46. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/launch_stage_pb2.py +37 -0
  47. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/google/api/resource_pb2.py +44 -0
  48. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/storage_browse_pb2.py +33 -0
  49. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/storage_browse_pb2_grpc.py +2 -0
  50. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/storage_browse_service_pb2.py +43 -0
  51. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/storage_browse_service_pb2_grpc.py +58 -0
  52. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_image_pb2.py +32 -0
  53. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_image_pb2_grpc.py +2 -0
  54. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_image_service_pb2.py +70 -0
  55. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_image_service_pb2_grpc.py +102 -0
  56. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_pb2.py +40 -0
  57. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_pb2_grpc.py +2 -0
  58. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_reporting_service_pb2.py +49 -0
  59. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_reporting_service_pb2_grpc.py +109 -0
  60. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_service_pb2.py +74 -0
  61. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_service_pb2_grpc.py +102 -0
  62. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_template_pb2.py +50 -0
  63. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_template_pb2_grpc.py +2 -0
  64. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_template_service_pb2.py +72 -0
  65. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/gen/worker_template_service_pb2_grpc.py +108 -0
  66. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/grpc_worker_client.py +113 -0
  67. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/models.py +94 -0
  68. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/worker_client.py +85 -0
  69. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/worker_image_client.py +83 -0
  70. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service/worker_template_client.py +88 -0
  71. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service.egg-info/PKG-INFO +28 -0
  72. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service.egg-info/SOURCES.txt +78 -0
  73. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service.egg-info/dependency_links.txt +1 -0
  74. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service.egg-info/requires.txt +29 -0
  75. h2o_connector_service-0.1.0.dev6001/src/h2o_connector_service.egg-info/top_level.txt +1 -0
  76. h2o_connector_service-0.1.0.dev6001/tests/test_connection_client.py +121 -0
  77. h2o_connector_service-0.1.0.dev6001/tests/test_connector_session.py +332 -0
  78. h2o_connector_service-0.1.0.dev6001/tests/test_discovery.py +227 -0
  79. h2o_connector_service-0.1.0.dev6001/tests/test_exceptions.py +128 -0
  80. h2o_connector_service-0.1.0.dev6001/tests/test_sdk_correctness.py +238 -0
@@ -0,0 +1,28 @@
1
+ Metadata-Version: 2.4
2
+ Name: h2o-connector-service
3
+ Version: 0.1.0.dev6001
4
+ Summary: Python client SDK for the H2O Connector Service gRPC API
5
+ Requires-Python: >=3.10
6
+ Requires-Dist: httpx>=0.27
7
+ Requires-Dist: grpcio>=1.64
8
+ Requires-Dist: protobuf>=4.25
9
+ Requires-Dist: googleapis-common-protos>=1.72
10
+ Requires-Dist: h2o-cloud-discovery>=3.3.0
11
+ Requires-Dist: h2o-authn>=3.1.0
12
+ Provides-Extra: pandas
13
+ Requires-Dist: pandas>=1.5; extra == "pandas"
14
+ Provides-Extra: parquet
15
+ Requires-Dist: pyarrow>=14; extra == "parquet"
16
+ Provides-Extra: datatable
17
+ Requires-Dist: datatable>=1.0; extra == "datatable"
18
+ Provides-Extra: h2o
19
+ Requires-Dist: pandas>=1.5; extra == "h2o"
20
+ Requires-Dist: h2o>=3.44; extra == "h2o"
21
+ Requires-Dist: pyarrow>=14; extra == "h2o"
22
+ Provides-Extra: test
23
+ Requires-Dist: pytest<9,>=8.0; extra == "test"
24
+ Requires-Dist: requests>=2.31; extra == "test"
25
+ Provides-Extra: dev
26
+ Requires-Dist: pytest<9,>=8.0; extra == "dev"
27
+ Requires-Dist: requests>=2.31; extra == "dev"
28
+ Requires-Dist: grpcio-tools>=1.64; extra == "dev"
@@ -0,0 +1,42 @@
1
+ [project]
2
+ name = "h2o-connector-service"
3
+ version = "0.1.0.dev6001"
4
+ description = "Python client SDK for the H2O Connector Service gRPC API"
5
+ requires-python = ">=3.10"
6
+ dependencies = [
7
+ "httpx>=0.27",
8
+ "grpcio>=1.64",
9
+ "protobuf>=4.25",
10
+ "googleapis-common-protos>=1.72",
11
+ "h2o-cloud-discovery>=3.3.0",
12
+ "h2o-authn>=3.1.0",
13
+ ]
14
+
15
+ [project.optional-dependencies]
16
+ pandas = ["pandas>=1.5"]
17
+ parquet = ["pyarrow>=14"]
18
+ datatable = ["datatable>=1.0"]
19
+ h2o = ["pandas>=1.5", "h2o>=3.44", "pyarrow>=14"]
20
+ test = [
21
+ "pytest>=8.0,<9",
22
+ "requests>=2.31",
23
+ ]
24
+ dev = [
25
+ "pytest>=8.0,<9",
26
+ "requests>=2.31",
27
+ "grpcio-tools>=1.64",
28
+ ]
29
+
30
+ [build-system]
31
+ requires = ["setuptools>=68"]
32
+ build-backend = "setuptools.build_meta"
33
+
34
+ [tool.setuptools.packages.find]
35
+ where = ["src"]
36
+
37
+ [tool.setuptools.package-data]
38
+ h2o_connector_service = ["gen/**/*.py", "gen/**/*.pyi"]
39
+
40
+ [tool.pytest.ini_options]
41
+ testpaths = ["tests"]
42
+ addopts = "-v --tb=short"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,35 @@
1
+ """connector-service Python SDK — REST control plane + gRPC worker data plane clients."""
2
+
3
+ from h2o_connector_service._base import Client
4
+ from h2o_connector_service.browse_storage_client import BrowseStorageClient
5
+ from h2o_connector_service.connection_client import ConnectionServiceClient
6
+ from h2o_connector_service.connector_client import ConnectorServiceClient
7
+ from h2o_connector_service.connector_service_facade import ConnectorService
8
+ from h2o_connector_service.connector_session import ConnectorSession
9
+ from h2o_connector_service.connector_template_client import ConnectorTemplateServiceClient
10
+ from h2o_connector_service.data_source_profile_client import DataSourceProfileServiceClient
11
+ from h2o_connector_service.exceptions import ConnectorServiceError
12
+ from h2o_connector_service.grpc_worker_client import WorkerClient
13
+ from h2o_connector_service.models import ConnectionStatus, ConnectorMeta, StorageEntry
14
+ from h2o_connector_service.worker_client import WorkerServiceClient
15
+ from h2o_connector_service.worker_image_client import WorkerImageServiceClient
16
+ from h2o_connector_service.worker_template_client import WorkerTemplateServiceClient
17
+
18
+ __all__ = [
19
+ "Client",
20
+ "BrowseStorageClient",
21
+ "ConnectionServiceClient",
22
+ "ConnectionStatus",
23
+ "ConnectorMeta",
24
+ "ConnectorService",
25
+ "ConnectorServiceClient",
26
+ "ConnectorServiceError",
27
+ "ConnectorSession",
28
+ "ConnectorTemplateServiceClient",
29
+ "DataSourceProfileServiceClient",
30
+ "StorageEntry",
31
+ "WorkerClient",
32
+ "WorkerImageServiceClient",
33
+ "WorkerServiceClient",
34
+ "WorkerTemplateServiceClient",
35
+ ]
@@ -0,0 +1,141 @@
1
+ """Base HTTP client for the connector-service Python SDK."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import warnings
7
+ from typing import Callable, Generator
8
+
9
+ import httpx
10
+
11
+ from h2o_connector_service.exceptions import ConnectorServiceError
12
+
13
+
14
+ class _TokenAuth(httpx.Auth):
15
+ """Per-request Bearer token injection via callable token provider."""
16
+
17
+ def __init__(self, token_provider: Callable[[], str]) -> None:
18
+ self._token_provider = token_provider
19
+
20
+ def auth_flow(
21
+ self, request: httpx.Request
22
+ ) -> Generator[httpx.Request, httpx.Response, None]:
23
+ request.headers["Authorization"] = f"Bearer {self._token_provider()}"
24
+ yield request
25
+
26
+
27
+ class Client:
28
+ """Base httpx client that injects ``Authorization: Bearer <token>`` on every request.
29
+
30
+ Args:
31
+ base_url: Base URL of the api-server (e.g. ``"http://localhost:8080"``).
32
+ token: OIDC bearer token. Injected as ``Authorization: Bearer <token>``.
33
+ workspace_id: Default workspace ID for workspace-scoped operations.
34
+
35
+ .. deprecated::
36
+ Use :meth:`from_discovery` instead to construct a client with automatic
37
+ URL resolution and token refresh from h2o Cloud Discovery.
38
+ """
39
+
40
+ def __init__(self, base_url: str, token: str, workspace_id: str) -> None:
41
+ warnings.warn(
42
+ "Client(base_url, token, workspace_id) is deprecated. "
43
+ "Use Client.from_discovery() instead.",
44
+ DeprecationWarning,
45
+ stacklevel=2,
46
+ )
47
+ self.base_url = base_url.rstrip("/")
48
+ self.token = token
49
+ self.workspace_id = workspace_id
50
+ self._http = httpx.Client(
51
+ base_url=self.base_url,
52
+ headers={"Authorization": f"Bearer {token}"},
53
+ )
54
+
55
+ @classmethod
56
+ def from_discovery(
57
+ cls,
58
+ environment: str,
59
+ workspace_id: str,
60
+ *,
61
+ token_provider=None,
62
+ ) -> "Client":
63
+ """Construct a Client by resolving the connector-service URL from h2o Cloud Discovery.
64
+
65
+ Args:
66
+ environment: H2O Cloud environment URL (e.g. ``"https://cloud.h2o.ai"``).
67
+ workspace_id: Workspace ID for workspace-scoped operations.
68
+ token_provider: Optional ``h2o_authn.TokenProvider``. If ``None``,
69
+ auto-created from discovery credentials.
70
+
71
+ Raises:
72
+ LookupError: If connector-service URL cannot be determined or
73
+ token provider cannot be created from discovery.
74
+ """
75
+ import h2o_discovery
76
+
77
+ discovery = None
78
+ try:
79
+ discovery = h2o_discovery.discover(environment=environment)
80
+ except (OSError, ConnectionError, TimeoutError) as e:
81
+ base_url = os.environ.get("H2O_CONNECTOR_SERVICE_URL")
82
+ if not base_url:
83
+ raise LookupError(
84
+ "Cannot resolve connector-service from h2o Cloud Discovery and "
85
+ "H2O_CONNECTOR_SERVICE_URL is not set."
86
+ ) from e
87
+ else:
88
+ # Resolve connector-service URL with env var fallback (D-11)
89
+ svc = discovery.services.get("connector-service")
90
+ if svc is not None:
91
+ base_url = svc.uri
92
+ else:
93
+ base_url = os.environ.get("H2O_CONNECTOR_SERVICE_URL")
94
+ if not base_url:
95
+ raise LookupError(
96
+ "connector-service not found in h2o Cloud Discovery and "
97
+ "H2O_CONNECTOR_SERVICE_URL is not set."
98
+ )
99
+
100
+ # Auto-create token provider from discovery credentials (D-05)
101
+ if token_provider is None:
102
+ if discovery is None:
103
+ raise LookupError(
104
+ "Cannot create token provider without discovery. "
105
+ "Provide token_provider when using H2O_CONNECTOR_SERVICE_URL fallback."
106
+ )
107
+ import h2o_authn.discovery as _authn_discovery
108
+
109
+ try:
110
+ token_provider = _authn_discovery.create(discovery)
111
+ except (KeyError, ValueError) as e:
112
+ raise LookupError(
113
+ "Cannot create token provider from discovery. "
114
+ "Configure the H2O CLI or set H2O_CLOUD_CLIENT_PLATFORM_TOKEN."
115
+ ) from e
116
+
117
+ # Build Client bypassing deprecated __init__ (D-12: cached at construction)
118
+ instance = cls.__new__(cls)
119
+ instance.base_url = base_url.rstrip("/")
120
+ instance.token = None # Not applicable for discovery-based clients
121
+ instance.workspace_id = workspace_id
122
+ instance._http = httpx.Client(
123
+ base_url=instance.base_url,
124
+ auth=_TokenAuth(token_provider),
125
+ )
126
+ return instance
127
+
128
+ def _raise_for_status(self, response: httpx.Response) -> None:
129
+ """Raise :class:`ConnectorServiceError` if *response* is not 2xx."""
130
+ if response.is_error:
131
+ raise ConnectorServiceError(response.status_code, response.text)
132
+
133
+ def close(self) -> None:
134
+ """Close the underlying httpx client."""
135
+ self._http.close()
136
+
137
+ def __enter__(self) -> Client:
138
+ return self
139
+
140
+ def __exit__(self, *args: object) -> None:
141
+ self.close()
@@ -0,0 +1,82 @@
1
+ """REST client for StorageBrowseService (data plane proxy via api-server)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from h2o_connector_service._base import Client
6
+
7
+
8
+ class BrowseStorageClient(Client):
9
+ """Wraps StorageBrowseService.BrowseStorage via the api-server REST proxy.
10
+
11
+ The api-server proxies the gRPC ``BrowseStorage`` call to the worker pod
12
+ identified by *connection_id*. Clients call the api-server HTTP endpoint
13
+ rather than the worker pod directly, using the workspace OIDC ``authorization``
14
+ token (Bearer).
15
+
16
+ Args:
17
+ base_url: Base URL of the api-server (e.g. ``"http://localhost:8080"``).
18
+ token: OIDC bearer token injected as ``Authorization: Bearer <token>``.
19
+ """
20
+
21
+ def __init__(self, base_url: str, token: str) -> None:
22
+ # workspace_id is not used for connection-scoped browse calls;
23
+ # the api-server resolves the workspace from the connection record.
24
+ super().__init__(base_url, token, workspace_id="")
25
+
26
+ # ------------------------------------------------------------------
27
+ # GRPC-04: StorageBrowseService.BrowseStorage
28
+ # ------------------------------------------------------------------
29
+
30
+ def browse(
31
+ self,
32
+ connection_id: str,
33
+ prefix: str = "",
34
+ delimiter: str = "/",
35
+ page_size: int = 0,
36
+ page_token: str = "",
37
+ ) -> dict:
38
+ """Browse storage entries for a connection.
39
+
40
+ Wraps ``GET /v1alpha/connections/{connection_id}/browse``.
41
+
42
+ Args:
43
+ connection_id: UUID of the connection whose worker pod to browse.
44
+ prefix: Path prefix to browse (default ``""`` — root).
45
+ delimiter: Separator for folder-style listing (default ``"/"``).
46
+ Pass ``""`` for a flat listing of all objects under *prefix*.
47
+ page_size: Maximum number of entries to return per page
48
+ (``0`` uses the server default).
49
+ page_token: Opaque page token returned by a previous call
50
+ (``""`` for the first page).
51
+
52
+ Returns:
53
+ A dict with the following keys:
54
+
55
+ - ``entries``: list of ``StorageEntry`` dicts
56
+ - ``next_page_token``: empty string when no further pages remain
57
+ - ``current_prefix``: the effective prefix used for this page
58
+
59
+ Raises:
60
+ :class:`~h2o_connector_service.exceptions.ConnectorServiceError`:
61
+ If the server returns a non-2xx HTTP status.
62
+ """
63
+ params: dict = {}
64
+ if prefix:
65
+ params["prefix"] = prefix
66
+ if delimiter != "/":
67
+ # Only override when caller explicitly changes the default
68
+ params["delimiter"] = delimiter
69
+ if page_size:
70
+ params["page_size"] = page_size
71
+ if page_token:
72
+ params["page_token"] = page_token
73
+
74
+ # Always send delimiter so the server gets explicit folder-vs-flat intent
75
+ params["delimiter"] = delimiter
76
+
77
+ resp = self._http.get(
78
+ f"/v1alpha/connections/{connection_id}/browse",
79
+ params=params,
80
+ )
81
+ self._raise_for_status(resp)
82
+ return resp.json()
@@ -0,0 +1,178 @@
1
+ """REST client for ConnectionService (control plane)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import time
6
+ from typing import Union
7
+
8
+ from h2o_connector_service._base import Client
9
+ from h2o_connector_service.models import ConnectionStatus
10
+
11
+
12
+ class ConnectionServiceClient:
13
+ """Wraps ConnectionService REST RPCs.
14
+
15
+ Args:
16
+ client: Authenticated :class:`~h2o_connector_service._base.Client` instance.
17
+ """
18
+
19
+ def __init__(self, client: Client) -> None:
20
+ self._c = client
21
+
22
+ # ------------------------------------------------------------------
23
+ # CRUD RPCs
24
+ # ------------------------------------------------------------------
25
+
26
+ def create_connection(self, workspace_id: str, connection: dict) -> dict:
27
+ """Create a new connection (POST /v1alpha/workspaces/{workspace_id}/connections)."""
28
+ resp = self._c._http.post(
29
+ f"/v1alpha/workspaces/{workspace_id}/connections",
30
+ json={"connection": connection},
31
+ )
32
+ self._c._raise_for_status(resp)
33
+ return resp.json()
34
+
35
+ def get_connection(self, workspace_id: str, connection_id: str) -> dict:
36
+ """Get a connection by ID (GET /v1alpha/workspaces/{workspace_id}/connections/{connection_id})."""
37
+ resp = self._c._http.get(
38
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}",
39
+ )
40
+ self._c._raise_for_status(resp)
41
+ return resp.json()
42
+
43
+ def list_connections(self, workspace_id: str, **params: object) -> dict:
44
+ """List connections in a workspace (GET /v1alpha/workspaces/{workspace_id}/connections).
45
+
46
+ Optional keyword params are forwarded as query parameters
47
+ (e.g., ``page_size``, ``page_token``, ``filter``, ``order_by``).
48
+ """
49
+ resp = self._c._http.get(
50
+ f"/v1alpha/workspaces/{workspace_id}/connections",
51
+ params={k: v for k, v in params.items() if v is not None},
52
+ )
53
+ self._c._raise_for_status(resp)
54
+ return resp.json()
55
+
56
+ def cancel_connection(self, workspace_id: str, connection_id: str) -> dict:
57
+ """Cancel a running connection (POST .../connections/{connection_id}:cancel)."""
58
+ resp = self._c._http.post(
59
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}:cancel",
60
+ json={},
61
+ )
62
+ self._c._raise_for_status(resp)
63
+ return resp.json()
64
+
65
+ def pause_connection(
66
+ self, workspace_id: str, connection_id: str, reason: str = ""
67
+ ) -> dict:
68
+ """Pause a running connection (POST .../connections/{connection_id}:pause)."""
69
+ body: dict = {}
70
+ if reason:
71
+ body["reason"] = reason
72
+ resp = self._c._http.post(
73
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}:pause",
74
+ json=body,
75
+ )
76
+ self._c._raise_for_status(resp)
77
+ return resp.json()
78
+
79
+ def resume_connection(self, workspace_id: str, connection_id: str) -> dict:
80
+ """Resume a paused connection (POST .../connections/{connection_id}:resume)."""
81
+ resp = self._c._http.post(
82
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}:resume",
83
+ json={},
84
+ )
85
+ self._c._raise_for_status(resp)
86
+ return resp.json()
87
+
88
+ def delete_connection(self, workspace_id: str, connection_id: str) -> None:
89
+ """Delete a connection (DELETE /v1alpha/workspaces/{workspace_id}/connections/{connection_id})."""
90
+ resp = self._c._http.delete(
91
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}",
92
+ )
93
+ self._c._raise_for_status(resp)
94
+
95
+ def generate_connection_token(self, workspace_id: str, connection_id: str) -> dict:
96
+ """Generate an auth token for a connection (POST .../connections/{connection_id}:generateToken).
97
+
98
+ Returns a dict with ``auth_token`` key containing the bearer token
99
+ for authenticating with the worker gRPC endpoint.
100
+ """
101
+ resp = self._c._http.post(
102
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}:generateToken",
103
+ json={},
104
+ )
105
+ self._c._raise_for_status(resp)
106
+ return resp.json()
107
+
108
+ def revoke_connection_token(self, workspace_id: str, connection_id: str, auth_token: str) -> dict:
109
+ """Revoke an auth token for a connection (POST .../connections/{connection_id}:revokeToken).
110
+
111
+ The token is immediately invalidated; the connection itself remains active.
112
+ Returns an empty dict on success (proto Empty / RevokeConnectionTokenResponse).
113
+ """
114
+ resp = self._c._http.post(
115
+ f"/v1alpha/workspaces/{workspace_id}/connections/{connection_id}:revokeToken",
116
+ json={"auth_token": auth_token},
117
+ )
118
+ self._c._raise_for_status(resp)
119
+ return resp.json()
120
+
121
+ # ------------------------------------------------------------------
122
+ # Typed status accessor
123
+ # ------------------------------------------------------------------
124
+
125
+ def get_connection_status(self, workspace_id: str, connection_id: str) -> ConnectionStatus:
126
+ """Get a typed connection status snapshot.
127
+
128
+ Wraps :meth:`get_connection` and returns a :class:`ConnectionStatus`
129
+ dataclass with ``connection_id``, ``status``, and ``worker_endpoint`` fields.
130
+ """
131
+ raw = self.get_connection(workspace_id, connection_id)
132
+ return ConnectionStatus.from_dict(raw)
133
+
134
+ # ------------------------------------------------------------------
135
+ # Polling helper (REST-07)
136
+ # ------------------------------------------------------------------
137
+
138
+ def poll_connection_status(
139
+ self,
140
+ workspace_id: str,
141
+ connection_id: str,
142
+ until_status: Union[str, list[str]],
143
+ interval: float = 2.0,
144
+ timeout: float = 300.0,
145
+ ) -> dict:
146
+ """Poll :meth:`get_connection` until the connection reaches a target status.
147
+
148
+ Replaces ``WatchConnection`` (server-streaming SSE) which is not
149
+ supported by the gRPC-Gateway configuration in this project.
150
+
151
+ Args:
152
+ workspace_id: Workspace that owns the connection.
153
+ connection_id: UUID of the connection to watch.
154
+ until_status: Target status string or list of status strings
155
+ (e.g. ``"CONNECTION_STATUS_WORKER_READY"`` or
156
+ ``["CONNECTION_STATUS_WORKER_READY", "CONNECTION_STATUS_WORKER_ACTIVE"]``).
157
+ interval: Polling interval in seconds (default 2.0).
158
+ timeout: Maximum total wait time in seconds (default 300.0).
159
+
160
+ Returns:
161
+ The connection dict when a target status is reached.
162
+
163
+ Raises:
164
+ TimeoutError: If the connection does not reach *until_status* within *timeout* seconds.
165
+ """
166
+ if isinstance(until_status, str):
167
+ until_status = [until_status]
168
+ deadline = time.monotonic() + timeout
169
+ while True:
170
+ conn = self.get_connection(workspace_id, connection_id)
171
+ if conn.get("status") in until_status:
172
+ return conn
173
+ if time.monotonic() >= deadline:
174
+ raise TimeoutError(
175
+ f"Connection {connection_id!r} did not reach {until_status!r} "
176
+ f"within {timeout}s"
177
+ )
178
+ time.sleep(interval)
@@ -0,0 +1,85 @@
1
+ """REST client for ConnectorService (control plane)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ from h2o_connector_service._base import Client
8
+
9
+
10
+ class ConnectorServiceClient:
11
+ """Wraps ConnectorService REST RPCs (workspace-scoped connectors).
12
+
13
+ Args:
14
+ client: Authenticated :class:`~h2o_connector_service._base.Client` instance.
15
+ """
16
+
17
+ def __init__(self, client: Client) -> None:
18
+ self._c = client
19
+
20
+ def create_connector(self, workspace_id: str, connector: dict) -> dict:
21
+ """Create a new connector (POST /v1alpha/workspaces/{workspace_id}/connectors).
22
+
23
+ Args:
24
+ workspace_id: Workspace to create the connector in.
25
+ connector: Connector resource dict (sent as ``{"connector": connector}`` body).
26
+ """
27
+ resp = self._c._http.post(
28
+ f"/v1alpha/workspaces/{workspace_id}/connectors",
29
+ json={"connector": connector},
30
+ )
31
+ self._c._raise_for_status(resp)
32
+ return resp.json()
33
+
34
+ def get_connector(self, name: str) -> dict:
35
+ """Get a connector by resource name (GET /v1alpha/{name}).
36
+
37
+ Args:
38
+ name: Full resource name, e.g. ``"workspaces/my-ws/connectors/my-connector"``.
39
+ """
40
+ resp = self._c._http.get(f"/v1alpha/{name}")
41
+ self._c._raise_for_status(resp)
42
+ return resp.json()
43
+
44
+ def list_connectors(self, workspace_id: str, **params: object) -> dict:
45
+ """List connectors in a workspace (GET /v1alpha/workspaces/{workspace_id}/connectors).
46
+
47
+ Optional keyword params are forwarded as query parameters
48
+ (e.g., ``page_size``, ``page_token``, ``filter``, ``order_by``).
49
+ """
50
+ resp = self._c._http.get(
51
+ f"/v1alpha/workspaces/{workspace_id}/connectors",
52
+ params={k: v for k, v in params.items() if v is not None},
53
+ )
54
+ self._c._raise_for_status(resp)
55
+ return resp.json()
56
+
57
+ def update_connector(
58
+ self, connector: dict, update_mask: Optional[str] = None
59
+ ) -> dict:
60
+ """Update a connector (PATCH /v1alpha/{connector.metadata.name}).
61
+
62
+ Args:
63
+ connector: Connector dict with ``metadata.name`` set to the full resource name.
64
+ update_mask: Comma-separated field mask (e.g. ``"enabled,flow_control"``).
65
+ """
66
+ name = connector["metadata"]["name"]
67
+ params: dict = {}
68
+ if update_mask:
69
+ params["update_mask"] = update_mask
70
+ resp = self._c._http.patch(
71
+ f"/v1alpha/{name}",
72
+ json={"connector": connector},
73
+ params=params,
74
+ )
75
+ self._c._raise_for_status(resp)
76
+ return resp.json()
77
+
78
+ def delete_connector(self, name: str) -> None:
79
+ """Delete a connector (DELETE /v1alpha/{name}).
80
+
81
+ Args:
82
+ name: Full resource name, e.g. ``"workspaces/my-ws/connectors/my-connector"``.
83
+ """
84
+ resp = self._c._http.delete(f"/v1alpha/{name}")
85
+ self._c._raise_for_status(resp)