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

pocket_cli/resources/s3.py

@@ -0,0 +1,307 @@
+from __future__ import annotations
+
+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
+from pocket_cli.resources.aws.s3_utils import (
+    bucket_exists,
+    create_bucket,
+    delete_bucket_with_contents,
+)
+
+if TYPE_CHECKING:
+    from pocket.context import CloudFrontContext, S3Context
+
+
+class S3:
+    context: S3Context
+    _cloudfront_contexts: dict[str, CloudFrontContext]
+
+    def __init__(
+        self,
+        context: S3Context,
+        cloudfront_contexts: dict[str, CloudFrontContext] | None = None,
+    ) -> None:
+        self.context = context
+        self._cloudfront_contexts = cloudfront_contexts or {}
+        self.client = boto3.client("s3", region_name=context.region)
+
+    @property
+    def description(self):
+        return "Create bucket: %s" % self.context.bucket_name
+
+    def state_info(self):
+        return {"s3": {"bucket_name": self.context.bucket_name}}
+
+    def deploy_init(self):
+        pass
+
+    def create(self):
+        create_bucket(self.client, self.context.bucket_name, self.context.region)
+        self.ensure_public_access_block()
+        self._ensure_cors()
+        self._ensure_versioning()
+        self._ensure_lifecycle()
+
+    def ensure_exists(self):
+        if self.exists():
+            self.ensure_public_access_block()
+            self._ensure_cors()
+            self._ensure_versioning()
+            self._ensure_lifecycle()
+            return
+        self.create()
+
+    def delete(self):
+        delete_bucket_with_contents(self.client, self.context.bucket_name)
+
+    def update(self):
+        self.ensure_public_access_block()
+        self._ensure_cors()
+        self._ensure_versioning()
+        self._ensure_lifecycle()
+
+    def exists(self):
+        try:
+            return bucket_exists(self.client, self.context.bucket_name)
+        except ClientError as e:
+            raise Exception(
+                "Bucket might be already used by other account."
+                " Try another bucket_prefix."
+            ) from e
+
+    @property
+    def status(self) -> ResourceStatus:
+        if not self.exists():
+            return "NOEXIST"
+        if self.public_access_block_require_update:
+            return "REQUIRE_UPDATE"
+        if self.cors_require_update:
+            return "REQUIRE_UPDATE"
+        if self.versioning_require_update:
+            return "REQUIRE_UPDATE"
+        if self.lifecycle_require_update:
+            return "REQUIRE_UPDATE"
+        return "COMPLETED"
+
+    @cached_property
+    def public_access_block(self) -> dict | None:
+        try:
+            res = self.client.get_public_access_block(
+                Bucket=self.context.bucket_name,
+            )
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchPublicAccessBlockConfiguration":
+                return None
+            raise
+        return res["PublicAccessBlockConfiguration"]
+
+    @property
+    def public_access_block_should_be(self):
+        return {
+            "BlockPublicAcls": True,
+            "IgnorePublicAcls": True,
+            "BlockPublicPolicy": True,
+            "RestrictPublicBuckets": True,
+        }
+
+    @property
+    def public_access_block_require_update(self):
+        return self.public_access_block_should_be != self.public_access_block
+
+    def ensure_public_access_block(self):
+        if self.public_access_block_require_update:
+            echo.info("Update public access block configuration")
+            echo.info("Current configuration: %s" % self.public_access_block)
+            self.client.put_public_access_block(
+                Bucket=self.context.bucket_name,
+                PublicAccessBlockConfiguration=self.public_access_block_should_be,
+            )
+            del self.public_access_block
+            echo.info("Updated to: %s" % self.public_access_block_should_be)
+        else:
+            echo.info("Public access block is already configured properly.")
+
+    def _resolve_cors_origins(self) -> list[str]:
+        """CORS の AllowedOrigins を CloudFront ドメインから解決する"""
+        if not self.context.cors:
+            return []
+        origins: list[str] = []
+        for cf_name in self.context.cors.cloudfront_names:
+            cf_ctx = self._cloudfront_contexts.get(cf_name)
+            if not cf_ctx:
+                echo.warning("CORS: cloudfront '%s' が見つかりません" % cf_name)
+                continue
+            if cf_ctx.domain:
+                origins.append("https://%s" % cf_ctx.domain)
+            else:
+                origins.append("https://*.cloudfront.net")
+        return origins
+
+    def _desired_cors_rules(self) -> list[dict] | None:
+        """期待する CORS ルールを返す。CORS 未設定なら None。"""
+        if not self.context.cors:
+            return None
+        origins = self._resolve_cors_origins()
+        if not origins:
+            return None
+        return [
+            {
+                "AllowedOrigins": origins,
+                "AllowedMethods": self.context.cors.methods,
+                "AllowedHeaders": ["*"],
+                "MaxAgeSeconds": 3600,
+            }
+        ]
+
+    @cached_property
+    def current_cors_rules(self) -> list[dict] | None:
+        """現在の S3 バケット CORS ルールを返す。未設定なら None。"""
+        try:
+            res = self.client.get_bucket_cors(Bucket=self.context.bucket_name)
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchCORSConfiguration":
+                return None
+            raise
+        return res.get("CORSRules")
+
+    @property
+    def cors_require_update(self) -> bool:
+        # 宣言的: desired と current が一致しなければ drift
+        return self._desired_cors_rules() != self.current_cors_rules
+
+    def _ensure_cors(self):
+        """S3 バケットの CORS 設定を宣言的に適用する。
+
+        - cors 宣言あり: PutBucketCors で置き換え
+        - cors 宣言なし & 現状あり: DeleteBucketCors
+        - cors 宣言なし & 現状なし: 何もしない
+        """
+        desired = self._desired_cors_rules()
+        current = self.current_cors_rules
+        if desired == current:
+            return
+        if desired is None:
+            self.client.delete_bucket_cors(Bucket=self.context.bucket_name)
+            self.__dict__.pop("current_cors_rules", None)
+            echo.info("CORS 設定を削除しました")
+            return
+        self.client.put_bucket_cors(
+            Bucket=self.context.bucket_name,
+            CORSConfiguration={"CORSRules": desired},
+        )
+        self.__dict__.pop("current_cors_rules", None)
+        echo.info("CORS 設定を適用しました: %s" % desired[0]["AllowedOrigins"])
+
+    @cached_property
+    def current_versioning_status(self) -> str | None:
+        """現在の S3 バケット versioning 状態 (Enabled / Suspended / None)。"""
+        res = self.client.get_bucket_versioning(Bucket=self.context.bucket_name)
+        return res.get("Status")
+
+    @property
+    def versioning_require_update(self) -> bool:
+        """宣言的判定。
+
+        - versioning=True: 現状が Enabled でなければ drift
+        - versioning=False: 現状が Enabled なら drift
+          (None / Suspended は "未有効化" として等価扱い、no-op)
+        """
+        current = self.current_versioning_status
+        if self.context.versioning:
+            return current != "Enabled"
+        return current == "Enabled"
+
+    def _ensure_versioning(self):
+        """S3 バケットの versioning を宣言的に適用する。
+
+        - versioning=True かつ現状 != Enabled: PutBucketVersioning(Enabled)
+        - versioning=False かつ現状 == Enabled: PutBucketVersioning(Suspended)
+        - その他: no-op (None と Suspended は "未有効化" として等価)
+
+        S3 仕様上、一度 Enabled にしたバケットは Suspended にしか戻せない
+        (バージョン情報は保持される)。docs に明記する。
+        """
+        current = self.current_versioning_status
+        if self.context.versioning:
+            if current == "Enabled":
+                return
+            self.client.put_bucket_versioning(
+                Bucket=self.context.bucket_name,
+                VersioningConfiguration={"Status": "Enabled"},
+            )
+            self.__dict__.pop("current_versioning_status", None)
+            echo.info("Versioning を有効化しました")
+            return
+        if current != "Enabled":
+            return
+        self.client.put_bucket_versioning(
+            Bucket=self.context.bucket_name,
+            VersioningConfiguration={"Status": "Suspended"},
+        )
+        self.__dict__.pop("current_versioning_status", None)
+        echo.info("Versioning を Suspended にしました")
+
+    def _desired_lifecycle_rules(self) -> list[dict] | None:
+        """期待する Lifecycle ルールを返す。空のとき None (= ルールなし)。"""
+        if not self.context.lifecycle_rules:
+            return None
+        rules: list[dict] = []
+        for rule in self.context.lifecycle_rules:
+            rules.append(
+                {
+                    "ID": rule.id,
+                    "Status": "Enabled",
+                    "Filter": {"Prefix": rule.prefix},
+                    "NoncurrentVersionExpiration": {
+                        "NoncurrentDays": rule.noncurrent_version_expiration_days,
+                    },
+                }
+            )
+        return rules
+
+    @cached_property
+    def current_lifecycle_rules(self) -> list[dict] | None:
+        """現在の S3 バケット Lifecycle ルール。未設定なら None。"""
+        try:
+            res = self.client.get_bucket_lifecycle_configuration(
+                Bucket=self.context.bucket_name,
+            )
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration":
+                return None
+            raise
+        return res.get("Rules")
+
+    @property
+    def lifecycle_require_update(self) -> bool:
+        # 宣言的: desired と current が一致しなければ drift
+        return self._desired_lifecycle_rules() != self.current_lifecycle_rules
+
+    def _ensure_lifecycle(self):
+        """S3 バケットの Lifecycle 設定を宣言的に適用する。
+
+        - lifecycle_rules 宣言あり: PutBucketLifecycleConfiguration で置き換え
+        - lifecycle_rules 空 & 現状あり: DeleteBucketLifecycle
+        - lifecycle_rules 空 & 現状なし: 何もしない
+        """
+        desired = self._desired_lifecycle_rules()
+        current = self.current_lifecycle_rules
+        if desired == current:
+            return
+        if desired is None:
+            self.client.delete_bucket_lifecycle(Bucket=self.context.bucket_name)
+            self.__dict__.pop("current_lifecycle_rules", None)
+            echo.info("Lifecycle 設定を削除しました")
+            return
+        self.client.put_bucket_lifecycle_configuration(
+            Bucket=self.context.bucket_name,
+            LifecycleConfiguration={"Rules": desired},
+        )
+        self.__dict__.pop("current_lifecycle_rules", None)
+        echo.info("Lifecycle ルールを適用しました: %d 件" % len(desired))

pocket_cli/resources/tidb.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import secrets
+import ssl
+import string
+import time
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+import requests
+from pydantic import BaseModel
+from requests.auth import HTTPDigestAuth
+
+from pocket.resources.base import ResourceStatus
+
+if TYPE_CHECKING:
+    from pocket.context import TiDbContext
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(level=os.getenv("POCKET_LOGGER_LEVEL", "WARNING").upper())
+
+
+class TiDbResourceIsNotReady(Exception):
+    pass
+
+
+class Project(BaseModel):
+    id: str
+    name: str
+
+
+class Cluster(BaseModel):
+    id: str
+    name: str
+    status: str
+    host: str
+    port: int
+    user: str
+
+
+class TiDbApi:
+    serverless_endpoint = "https://serverless.tidbapi.com/v1beta1/"
+    iam_endpoint = "https://iam.tidbapi.com/v1beta1/"
+
+    def __init__(self, public_key: str, private_key: str) -> None:
+        self.auth = HTTPDigestAuth(public_key, private_key)
+
+    def _request(self, method: str, url: str, data=None):
+        logger.info("%s %s" % (method, url))
+        if data:
+            logger.debug(json.dumps(data, indent=2))
+        res = requests.request(method, url, auth=self.auth, json=data)
+        logger.debug(res.status_code)
+        logger.debug(json.dumps(res.json(), indent=2))
+        if 200 <= res.status_code < 300:
+            if method != "GET":
+                time.sleep(2)
+            return res
+        raise RuntimeError("%s: %s" % (res.status_code, res.text))
+
+    def iam_get(self, path: str):
+        return self._request("GET", self.iam_endpoint + path)
+
+    def serverless_get(self, path: str):
+        return self._request("GET", self.serverless_endpoint + path)
+
+    def serverless_post(self, path: str, data=None):
+        return self._request("POST", self.serverless_endpoint + path, data)
+
+    def serverless_put(self, path: str, data=None):
+        return self._request("PUT", self.serverless_endpoint + path, data)
+
+    def serverless_delete(self, path: str):
+        return self._request("DELETE", self.serverless_endpoint + path)
+
+    def list_projects(self) -> list[dict]:
+        return self.iam_get("projects").json().get("projects", [])
+
+    def list_clusters(self, project_id: str) -> list[dict]:
+        return (
+            self.serverless_get("clusters?filter=projectId=%s" % project_id)
+            .json()
+            .get("clusters", [])
+        )
+
+    def create_cluster(self, data: dict) -> dict:
+        return self.serverless_post("clusters", data).json()
+
+    def get_cluster(self, cluster_id: str) -> dict:
+        return self.serverless_get("clusters/%s" % cluster_id).json()
+
+    def delete_cluster(self, cluster_id: str):
+        return self.serverless_delete("clusters/%s" % cluster_id)
+
+    def change_password(self, cluster_id: str, password: str):
+        return self.serverless_put(
+            "clusters/%s/password" % cluster_id,
+            {"password": password},
+        )
+
+
+class TiDb:
+    context: TiDbContext
+    _root_password: str | None
+
+    def __init__(self, context: TiDbContext) -> None:
+        self.context = context
+        self._root_password = None
+
+    @cached_property
+    def api(self) -> TiDbApi:
+        if not self.context.public_key or not self.context.private_key:
+            raise TiDbResourceIsNotReady(
+                "TiDB API keys not configured. "
+                "Set tidb_public_key and tidb_private_key in .env"
+            )
+        return TiDbApi(self.context.public_key, self.context.private_key)
+
+    @cached_property
+    def project(self) -> Project | None:
+        projects = self.api.list_projects()
+        if self.context.tidb_project:
+            for p in projects:
+                if p["name"] == self.context.tidb_project:
+                    return Project(id=str(p["id"]), name=p["name"])
+            return None
+        if len(projects) == 1:
+            p = projects[0]
+            return Project(id=str(p["id"]), name=p["name"])
+        if len(projects) > 1:
+            names = [p["name"] for p in projects]
+            raise TiDbResourceIsNotReady(
+                "複数の TiDB Cloud プロジェクトが見つかりました: %s\n"
+                "pocket.toml の [tidb] で project を指定してください。\n"
+                '例: project = "%s"' % (names, names[0])
+            )
+        return None
+
+    @cached_property
+    def cluster(self) -> Cluster | None:
+        if not self.project:
+            return None
+        for c in self.api.list_clusters(self.project.id):
+            if c.get("displayName") == self.context.cluster_name:
+                endpoints = c.get("endpoints", {}).get("public", {})
+                user_prefix = c.get("userPrefix", "")
+                return Cluster(
+                    id=str(c["clusterId"]),
+                    name=c["displayName"],
+                    status=c.get("state", "UNKNOWN"),
+                    host=endpoints.get("host", ""),
+                    port=int(endpoints.get("port", 4000)),
+                    user="%s.root" % user_prefix,
+                )
+        return None
+
+    @property
+    def status(self) -> ResourceStatus:
+        if self.context.skip_check_existing:
+            # 存在確認の TiDB API call を skip し COMPLETED 固定 (settings 参照)。
+            return "COMPLETED"
+        if not self.context.public_key or not self.context.private_key:
+            return "NOEXIST"
+        if self.cluster and self.cluster.status == "ACTIVE":
+            return "COMPLETED"
+        return "NOEXIST"
+
+    @property
+    def description(self) -> str:
+        return "Create TiDB cluster and database"
+
+    @property
+    def database_url(self) -> str:
+        if not self.cluster:
+            raise TiDbResourceIsNotReady("Cluster not found")
+        if self.cluster.status != "ACTIVE":
+            raise TiDbResourceIsNotReady("Cluster is not available")
+        password = self._root_password
+        if not password:
+            password = self._reset_password()
+        return "mysql://%s:%s@%s:%d/%s" % (
+            self.cluster.user,
+            password,
+            self.cluster.host,
+            self.cluster.port,
+            self.context.database_name,
+        )
+
+    def _generate_password(self) -> str:
+        chars = string.ascii_letters + string.digits
+        return "".join(secrets.choice(chars) for _ in range(24))
+
+    def _reset_password(self) -> str:
+        if not self.cluster:
+            raise TiDbResourceIsNotReady("Cluster not found")
+        new_password = self._generate_password()
+        self.api.change_password(self.cluster.id, new_password)
+        self._root_password = new_password
+        return new_password
+
+    def deploy_init(self):
+        pass
+
+    def create(self):
+        self._ensure_cluster()
+        self._ensure_database()
+
+    def _ensure_cluster(self):
+        if self.cluster:
+            return
+        if not self.project:
+            raise TiDbResourceIsNotReady(
+                "No TiDB Cloud project found. Create one at https://tidbcloud.com/"
+            )
+        password = self._generate_password()
+        data = {
+            "displayName": self.context.cluster_name,
+            "region": {"name": "regions/aws-%s" % self.context.region},
+            "rootPassword": password,
+            "labels": {"tidb.cloud/project": self.project.id},
+            "endpoints": {
+                "public": {
+                    "authorizedNetworks": [
+                        {
+                            "startIpAddress": "0.0.0.0",
+                            "endIpAddress": "255.255.255.255",
+                            "displayName": "Allow all",
+                        }
+                    ]
+                }
+            },
+        }
+        self.api.create_cluster(data)
+        self._root_password = password
+        self._wait_for_cluster()
+
+    def _wait_for_cluster(self):
+        logger.info("Waiting for cluster to become available...")
+        del self.cluster
+        for _ in range(60):
+            if self.cluster and self.cluster.status == "ACTIVE":
+                return
+            time.sleep(5)
+            del self.cluster
+        raise TiDbResourceIsNotReady("Cluster did not become available in time")
+
+    def _get_sql_connection(self):
+        try:
+            import pymysql  # noqa: PLC0415
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "pymysql is required for TiDB database management.\n"
+                "Install it with: uv add pymysql"
+            ) from None
+        if not self.cluster:
+            raise TiDbResourceIsNotReady("Cluster not found")
+        password = self._root_password
+        if not password:
+            password = self._reset_password()
+        ssl_ctx = ssl.create_default_context()
+        return pymysql.connect(
+            host=self.cluster.host,
+            port=self.cluster.port,
+            user=self.cluster.user,
+            password=password,
+            ssl=ssl_ctx,
+        )
+
+    def _ensure_database(self):
+        conn = self._get_sql_connection()
+        try:
+            with conn.cursor() as cursor:
+                cursor.execute(
+                    "CREATE DATABASE IF NOT EXISTS `%s`" % self.context.database_name
+                )
+        finally:
+            conn.close()
+
+    def delete_cluster(self):
+        if not self.cluster:
+            raise TiDbResourceIsNotReady("Cluster not found")
+        self.api.delete_cluster(self.cluster.id)
+        del self.cluster
+
+    def reset_database(self):
+        conn = self._get_sql_connection()
+        try:
+            with conn.cursor() as cursor:
+                cursor.execute(
+                    "DROP DATABASE IF EXISTS `%s`" % self.context.database_name
+                )
+                cursor.execute("CREATE DATABASE `%s`" % self.context.database_name)
+        finally:
+            conn.close()