File size: 3,940 Bytes
1e23d14 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import logging
import time
from typing import Any, Optional
logger = logging.getLogger(__name__)
try:
from google.cloud import compute_v1
from google.api_core.exceptions import NotFound
except ImportError:
compute_v1 = None # type: ignore
NotFound = Exception
class ComputeClient:
def __init__(self, project_id: str, zone: str,
credentials_file: Optional[str] = None) -> None:
self.project_id = project_id
self.zone = zone
if compute_v1 is None:
raise ImportError("google-cloud-compute is not installed")
self._instances = compute_v1.InstancesClient()
self._operations = compute_v1.ZoneOperationsClient()
logger.debug("Compute client for %s/%s", project_id, zone)
def _wait_for_operation(self, operation_name: str,
timeout: int = 300) -> None:
deadline = time.time() + timeout
while time.time() < deadline:
op = self._operations.get(project=self.project_id,
zone=self.zone,
operation=operation_name)
if op.status == compute_v1.Operation.Status.DONE:
if op.error:
raise RuntimeError(f"Operation failed: {op.error}")
return
time.sleep(5)
raise TimeoutError(f"Operation {operation_name} timed out")
def list_instances(self, filter_expr: str = "") -> list[dict[str, Any]]:
request = compute_v1.ListInstancesRequest(
project=self.project_id, zone=self.zone, filter=filter_expr
)
instances = []
for inst in self._instances.list(request=request):
instances.append({
"name": inst.name,
"status": inst.status,
"machine_type": inst.machine_type.split("/")[-1],
"zone": self.zone,
})
return instances
def get_instance(self, name: str) -> dict[str, Any]:
inst = self._instances.get(project=self.project_id,
zone=self.zone, instance=name)
return {
"name": inst.name,
"status": inst.status,
"machine_type": inst.machine_type.split("/")[-1],
"network_ips": [
iface.network_i_p
for iface in inst.network_interfaces
],
}
def start_instance(self, name: str) -> None:
op = self._instances.start(project=self.project_id,
zone=self.zone, instance=name)
self._wait_for_operation(op.name)
logger.info("Started instance %s", name)
def stop_instance(self, name: str) -> None:
op = self._instances.stop(project=self.project_id,
zone=self.zone, instance=name)
self._wait_for_operation(op.name)
logger.info("Stopped instance %s", name)
def instance_status(self, name: str) -> str:
try:
inst = self._instances.get(project=self.project_id,
zone=self.zone, instance=name)
return inst.status
except NotFound:
return "NOT_FOUND"
def set_labels(self, name: str, labels: dict[str, str]) -> None:
inst = self._instances.get(project=self.project_id,
zone=self.zone, instance=name)
req = compute_v1.SetLabelsInstanceRequest(
project=self.project_id, zone=self.zone, instance=name,
instances_set_labels_request_resource=compute_v1.InstancesSetLabelsRequest(
label_fingerprint=inst.label_fingerprint,
labels=labels,
),
)
op = self._instances.set_labels(request=req)
self._wait_for_operation(op.name)
logger.info("Labels updated for instance %s", name)
|