magic-pocket-cli 0.2.0__py3-none-any.whl

pocket_cli/resources/dsql.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import time
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+import boto3
+from botocore.exceptions import ClientError
+
+from pocket.resources.base import ResourceStatus
+from pocket.utils import echo
+
+if TYPE_CHECKING:
+    from pocket.context import DsqlContext
+
+
+class Dsql:
+    context: DsqlContext
+
+    def __init__(self, context: DsqlContext) -> None:
+        self.context = context
+        self._client = boto3.client("dsql", region_name=context.region)
+
+    @cached_property
+    def cluster(self) -> dict | None:
+        """Name タグで DSQL クラスターを検索"""
+        paginator = self._client.get_paginator("list_clusters")
+        for page in paginator.paginate():
+            for cluster in page["clusters"]:
+                identifier = cluster["identifier"]
+                try:
+                    detail = self._client.get_cluster(Identifier=identifier)
+                    tags = self._client.list_tags_for_resource(
+                        ResourceArn=detail["arn"]
+                    )
+                    if tags.get("tags", {}).get("Name") == self.context.tag_name:
+                        return detail
+                except ClientError:
+                    continue
+        return None
+
+    @property
+    def identifier(self) -> str | None:
+        if self.cluster:
+            return self.cluster["identifier"]
+        return None
+
+    @property
+    def endpoint(self) -> str | None:
+        if self.identifier:
+            return f"{self.identifier}.dsql.{self.context.region}.on.aws"
+        return None
+
+    @property
+    def arn(self) -> str | None:
+        if self.cluster:
+            return self.cluster["arn"]
+        return None
+
+    @property
+    def status(self) -> ResourceStatus:
+        if self.cluster is None:
+            return "NOEXIST"
+        cluster_status = self.cluster["status"]
+        if cluster_status in ("CREATING", "UPDATING", "DELETING"):
+            return "PROGRESS"
+        if cluster_status == "ACTIVE":
+            return "COMPLETED"
+        return "FAILED"
+
+    @property
+    def description(self):
+        return "Create Aurora DSQL cluster: %s" % self.context.tag_name
+
+    def state_info(self):
+        return {
+            "dsql": {
+                "tag_name": self.context.tag_name,
+                "identifier": self.identifier,
+                "endpoint": self.endpoint,
+            }
+        }
+
+    def deploy_init(self):
+        pass
+
+    def create(self):
+        echo.log("Creating DSQL cluster: %s" % self.context.tag_name)
+        res = self._client.create_cluster(
+            deletionProtectionEnabled=self.context.deletion_protection,
+            tags={"Name": self.context.tag_name},
+        )
+        identifier = res["identifier"]
+        echo.log("Cluster ID: %s" % identifier)
+        echo.log("Waiting for DSQL cluster to become active...")
+        self._wait_active(identifier, timeout=600)
+        self.clear_cache()
+        echo.success("DSQL cluster is now active.")
+        echo.success("Endpoint: %s" % self.endpoint)
+
+    def delete(self):
+        if not self.identifier:
+            return
+        echo.log("Deleting DSQL cluster: %s" % self.identifier)
+        self._client.delete_cluster(Identifier=self.identifier)
+        echo.log("Waiting for DSQL cluster deletion...")
+        self._wait_deleted(self.identifier, timeout=600)
+        echo.success("DSQL cluster was deleted.")
+
+    def _wait_active(self, identifier: str, timeout: int = 600, interval: int = 5):
+        for i in range(timeout // interval):
+            try:
+                res = self._client.get_cluster(Identifier=identifier)
+                if res["status"] == "ACTIVE":
+                    print("")
+                    return
+            except ClientError:
+                pass
+            if i == 0:
+                print("Waiting for cluster to be active", end="", flush=True)
+            print(".", end="", flush=True)
+            time.sleep(interval)
+        raise TimeoutError("Cluster did not become active within %s seconds" % timeout)
+
+    def _wait_deleted(self, identifier: str, timeout: int = 600, interval: int = 5):
+        for i in range(timeout // interval):
+            try:
+                self._client.get_cluster(Identifier=identifier)
+            except ClientError as e:
+                if e.response["Error"]["Code"] == "ResourceNotFoundException":
+                    print("")
+                    return
+                raise
+            if i == 0:
+                print("Waiting for cluster deletion", end="", flush=True)
+            print(".", end="", flush=True)
+            time.sleep(interval)
+        raise TimeoutError("Cluster not deleted within %s seconds" % timeout)
+
+    def clear_cache(self):
+        if "cluster" in self.__dict__:
+            del self.__dict__["cluster"]

pocket_cli/resources/neon.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import time
+from functools import cached_property
+from typing import TYPE_CHECKING, Literal
+
+import requests
+from pydantic import BaseModel
+
+from pocket.resources.base import ResourceStatus
+
+if TYPE_CHECKING:
+    from pocket.context import NeonContext
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(level=os.getenv("POCKET_LOGGER_LEVEL", "WARNING").upper())
+
+ResourceType = Literal["branches", "databases", "endpoints", "roles"]
+
+
+class NeonResourceIsNotReady(Exception):
+    pass
+
+
+class NeonNotFound(Exception):
+    """Neon API が 404 を返したことを示す例外。
+
+    `branch` / `role` などの個別取得で対象が存在しない場合にこの例外を投げる。
+    呼び出し側は存在しないことを明示的に判定できる。
+    """
+
+    pass
+
+
+class Project(BaseModel):
+    id: str
+    name: str
+
+
+class Branch(BaseModel):
+    id: str
+    name: str
+
+
+class Database(BaseModel):
+    name: str
+    owner_name: str
+
+
+class Endpoint(BaseModel):
+    id: str
+    host: str
+    autoscaling_limit_min_cu: float
+    autoscaling_limit_max_cu: float
+    type: Literal["read_write", "read_only"]
+
+
+class Role(BaseModel):
+    name: str
+    password: str | None = None
+
+
+class NeonApi:
+    endpoint = "https://console.neon.tech/api/v2/"
+
+    def __init__(self, key) -> None:
+        self.key = key
+
+    @property
+    def header(self):
+        return {
+            "Accept": "application/json",
+            "Authorization": "Bearer %s" % self.key,
+        }
+
+    def get(self, path):
+        logger.info("GET %s" % self.endpoint + path)
+        res = requests.get(self.endpoint + path, headers=self.header)
+        logger.debug(res.status_code)
+        logger.debug(json.dumps(res.json(), indent=2))
+        if 200 <= res.status_code < 300:
+            return res
+        if res.status_code == 404:
+            raise NeonNotFound("%s: %s" % (res.status_code, res.json()["message"]))
+        if res.status_code == 401:
+            print("Used API key: %s" % (self.key[:5] + "..." + self.key[-5:]))
+            print("API key length: %s" % len(self.key))
+        raise Exception("%s: %s" % (res.status_code, res.json()["message"]))
+
+    def post(self, path, data=None):
+        logger.warning("POST %s" % self.endpoint + path)
+        logger.debug(json.dumps(data, indent=2))
+        res = requests.post(self.endpoint + path, headers=self.header, json=data)
+        logger.debug(res.status_code)
+        logger.debug(json.dumps(res.json(), indent=2))
+        if 200 <= res.status_code < 300:
+            time.sleep(2)
+            return res
+        if res.status_code == 401:
+            print("Used API key: %s" % (self.key[:5] + "..." + self.key[-5:]))
+            print("API key length: %s" % len(self.key))
+        raise Exception("%s: %s" % (res.status_code, res.json()["message"]))
+
+    def delete(self, path, data=None):
+        logger.warning("DELETE %s" % self.endpoint + path)
+        logger.debug(json.dumps(data, indent=2))
+        res = requests.delete(self.endpoint + path, headers=self.header, json=data)
+        logger.debug(res.status_code)
+        logger.debug(json.dumps(res.json(), indent=2))
+        if 200 <= res.status_code < 300:
+            time.sleep(2)
+            return res
+        raise Exception("%s: %s" % (res.status_code, res.json()["message"]))
+
+    def projects_url(self):
+        return self.endpoint + "projects"
+
+
+class Neon:
+    context: NeonContext
+
+    def __init__(self, context: NeonContext) -> None:
+        self.context = context
+
+    def get_resource_path(self, resource_type: ResourceType) -> str:
+        requirements = {
+            "branches": ["project"],
+            "databases": ["project", "branch"],
+            "endpoints": ["project"],
+            "roles": ["project", "branch"],
+        }
+        path_templates = {
+            "branches": "projects/%(project_id)s/branches",
+            "databases": "projects/%(project_id)s/branches/%(branch_id)s/databases",
+            "endpoints": "projects/%(project_id)s/endpoints",
+            "roles": "projects/%(project_id)s/branches/%(branch_id)s/roles",
+        }
+        path_context = {}
+        for requirement in requirements[resource_type]:
+            if not getattr(self, requirement):
+                raise Exception("%s not found" % requirement)
+            path_context[requirement + "_id"] = getattr(self, requirement).id
+        return path_templates[resource_type] % path_context
+
+    def construct_path(
+        self, resource_type: ResourceType, resource_id: str | None = None
+    ):
+        path = self.get_resource_path(resource_type)
+        if resource_id:
+            path += "/" + resource_id
+        return path
+
+    def get(self, resource_type: ResourceType, resource_id: str | None = None):
+        path = self.construct_path(resource_type, resource_id)
+        return self.api.get(path)
+
+    def post(
+        self, resource_type: ResourceType, resource_id: str | None = None, data=None
+    ):
+        path = self.construct_path(resource_type, resource_id)
+        return self.api.post(path, data)
+
+    def delete(
+        self, resource_type: ResourceType, resource_id: str | None = None, data=None
+    ):
+        path = self.construct_path(resource_type, resource_id)
+        return self.api.delete(path, data)
+
+    @property
+    def api(self):
+        return NeonApi(self.context.api_key)
+
+    @cached_property
+    def role(self) -> Role | None:
+        if not self.branch:
+            return None
+        try:
+            res = self.get("roles", self.context.role_name)
+        except NeonNotFound:
+            return None
+        return Role(**res.json()["role"])
+
+    @cached_property
+    def project(self) -> Project:
+        """project_nameからNeonプロジェクトを解決する。
+
+        組織キー: GET /projects で一覧取得し name で検索。
+        プロジェクトスコープキー: エラーから project_id を取得し直接アクセス。
+        """
+        res = requests.get(self.api.endpoint + "projects", headers=self.api.header)
+        if 200 <= res.status_code < 300:
+            for p in res.json().get("projects", []):
+                if p["name"] == self.context.project_name:
+                    return Project(**p)
+            raise ValueError(f"Neon project '{self.context.project_name}' not found")
+
+        # プロジェクトスコープキー: エラーから project_id をパース
+        message = res.json().get("message", "")
+        match = re.search(r'subject_project_id:"([^"]+)"', message)
+        if not match:
+            raise ValueError(
+                f"Failed to resolve Neon project: {res.status_code}: {message}"
+            )
+        project_id = match.group(1)
+        project_res = self.api.get(f"projects/{project_id}")
+        project_data = project_res.json()["project"]
+        if project_data["name"] != self.context.project_name:
+            raise ValueError(
+                f"Neon project name mismatch: "
+                f"config='{self.context.project_name}', "
+                f"actual='{project_data['name']}'"
+            )
+        return Project(id=project_data["id"], name=project_data["name"])
+
+    @cached_property
+    def branch(self) -> Branch | None:
+        if self.project:
+            for branch in self.get("branches").json()["branches"]:
+                if branch["name"] == self.context.branch_name:
+                    return Branch(**branch)
+
+    @cached_property
+    def database(self) -> Database | None:
+        if self.branch:
+            for database in self.get("databases").json()["databases"]:
+                if database["name"] == self.context.name:
+                    return Database(**database)
+
+    @cached_property
+    def endpoint(self) -> Endpoint | None:
+        if self.branch:
+            for endpoint in self.get("endpoints").json()["endpoints"]:
+                if endpoint["branch_id"] == self.branch.id:
+                    return Endpoint(**endpoint)
+
+    @property
+    def database_url(self):
+        if self.role is None or self.endpoint is None:
+            raise NeonResourceIsNotReady("Create role and endpoint first")
+        if self.role.password is None:
+            self.set_role_password()
+        return "postgres://%s:%s@%s:5432/%s?sslmode=require" % (
+            self.context.role_name,
+            self.role.password,
+            self.endpoint.host,
+            self.context.name,
+        )
+
+    @property
+    def status(self) -> ResourceStatus:
+        if self.context.skip_check_existing:
+            # 存在確認の Neon API call を skip し COMPLETED 固定。deploy ロールに
+            # Neon credentials を渡さず deploy を完走させる用途 (settings 参照)。
+            return "COMPLETED"
+        if self.working:
+            return "COMPLETED"
+        return "NOEXIST"
+
+    @property
+    def working(self):
+        check = [self.branch, self.database, self.endpoint, self.role]
+        logger.info(str(check))
+        return all(check)
+
+    @property
+    def description(self):
+        return "Create Neon branch, database, role, and endpoint"
+
+    def create_new(self):
+        self.create()
+        self.reset_database()
+
+    def state_info(self):
+        return {
+            "neon": {
+                "project_name": self.context.project_name,
+                "branch_name": self.context.branch_name,
+            }
+        }
+
+    def deploy_init(self):
+        pass
+
+    def create(self):
+        self.create_branch()
+        self.ensure_role()
+        self.ensure_database()
+
+    def create_branch(self, base_branch: Branch | None = None):
+        if self.branch is None:
+            del self.branch
+            del self.endpoint
+        data = {
+            "branch": {
+                "name": self.context.branch_name,
+            },
+            "endpoints": [{"type": "read_write"}],
+        }
+        if base_branch:
+            data["branch"]["parent_id"] = base_branch.id
+        self.post("branches", data=data)
+
+    def delete_branch(self):
+        if not self.endpoint or not self.branch:
+            raise Exception("Branch or endpoint not found. Something is wrong.")
+        self.delete("endpoints", self.endpoint.id)
+        self.delete("branches", self.branch.id)
+
+    def ensure_database(self):
+        if self.database is None:
+            self.create_database()
+
+    def create_database(self):
+        if self.database is None:
+            del self.database
+        data = {
+            "database": {
+                "name": self.context.name,
+                "owner_name": self.context.role_name,
+            }
+        }
+        self.post("databases", data=data)
+
+    def reset_database(self):
+        if self.database:
+            self.delete("databases", self.context.name)
+        self.create_database()
+
+    def create_role(self):
+        if self.role is None:
+            del self.role
+        data = {"role": {"name": self.context.role_name}}
+        self.post("roles", data=data)
+
+    def set_role_password(self):
+        if self.role is None:
+            raise Exception("Create role first")
+        if self.role.password is None:
+            self.role.password = self.get(
+                "roles", self.role.name + "/reveal_password"
+            ).json()["password"]
+
+    def ensure_role(self):
+        if self.role is None:
+            del self.role
+            data = {"role": {"name": self.context.role_name}}
+            self.post("roles", data=data)
+            self.set_role_password()