26class PostgreSQLManager:
 27    """PostgreSQL接続とクエリ実行クラス
 28
 29    Azure AD認証またはユーザー名/パスワード認証でPostgreSQLに接続し、
 30    クエリを実行します。接続プールを使用して効率的に接続を管理します。
 31
 32    Features:
 33        - Azure AD認証対応（トークン自動更新）
 34        - 標準認証（ユーザー名/パスワード）対応
 35        - 接続プール管理
 36        - 非同期コンテキストマネージャー対応
 37        - SQL Injection防止（パラメータ化クエリ）
 38
 39    Example:
 40        >>> from agenticstar_platform import PostgreSQLManager, PostgreSQLConfig
 41        >>>
 42        >>> # 設定を作成
 43        >>> config = PostgreSQLConfig.from_toml("config.toml")
 44        >>>
 45        >>> # コンテキストマネージャーで使用（推奨）
 46        >>> async with PostgreSQLManager(config) as db:
 47        ...     # SELECT実行
 48        ...     result = await db.fetch_all("SELECT * FROM users WHERE status = $1", ("active",))
 49        ...     print(f"Active users: {len(result)}")
 50        ...
 51        ...     # INSERT実行
 52        ...     await db.execute(
 53        ...         "INSERT INTO logs (message) VALUES ($1)",
 54        ...         ("Operation completed",)
 55        ...     )
 56        >>>
 57        >>> # 明示的な初期化・クローズ
 58        >>> db = PostgreSQLManager(config)
 59        >>> await db.initialize()
 60        >>> try:
 61        ...     result = await db.fetch_one("SELECT COUNT(*) as cnt FROM users")
 62        ...     print(f"Total users: {result['cnt']}")
 63        ... finally:
 64        ...     await db.close()
 65    """
 66
 67    def __init__(self, config: PostgreSQLConfig):
 68        self.config = config
 69        self.pool: Optional[asyncpg.Pool] = None
 70        self._access_token: Optional[str] = None
 71        self._token_expiry: Optional[float] = None
 72        self._token_refresh_lock: Optional[asyncio.Lock] = None
 73
 74    async def __aenter__(self) -> "PostgreSQLManager":
 75        """Context Manager: 開始時に接続プールを初期化"""
 76        await self.initialize()
 77        return self
 78
 79    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
 80        """Context Manager: 終了時に接続プールをクローズ"""
 81        await self.close()
 82
 83    async def _get_azure_ad_token(self) -> str:
 84        """Azure AD認証でアクセストークンを取得"""
 85        logger.info("[IN] _get_azure_ad_token: Getting Azure AD token")
 86
 87        if not self.config.azure_ad:
 88            logger.error("[ERROR] _get_azure_ad_token: Azure AD configuration is required")
 89            raise ValueError("Azure AD configuration is required")
 90
 91        # トークンの有効期限をチェック（5分前に更新）
 92        if (
 93            self._access_token
 94            and self._token_expiry
 95            and time.time() < self._token_expiry - 300
 96        ):
 97            logger.info("[OUT] _get_azure_ad_token: Using cached token")
 98            return self._access_token
 99
100        # トークン取得中の同期処理（複数API同時実行時の重複取得防止）
101        # Note: Lockはinitialize()で作成済み（イベントループ紐付けを保証）
102        async with self._token_refresh_lock:
103            # ロック取得後に再チェック（他のAPIコールが既に更新した可能性）
104            if (
105                self._access_token
106                and self._token_expiry
107                and time.time() < self._token_expiry - 300
108            ):
109                return self._access_token
110
111            try:
112                credential = ClientSecretCredential(
113                    tenant_id=self.config.azure_ad.tenant_id,
114                    client_id=self.config.azure_ad.client_id,
115                    client_secret=self.config.azure_ad.client_secret,
116                )
117
118                # Try different scopes for Azure PostgreSQL Flexible Server
119                scopes_to_try = [
120                    "https://ossrdbms-aad.database.windows.net/.default",
121                    "https://database.windows.net/.default",
122                    "https://management.azure.com/.default",
123                ]
124
125                token = None
126                last_error = None
127
128                for scope in scopes_to_try:
129                    try:
130                        token = await asyncio.to_thread(credential.get_token, scope)
131                        break
132                    except Exception as e:
133                        last_error = e
134                        continue
135
136                if not token:
137                    raise last_error or Exception("All token scopes failed")
138
139                self._access_token = token.token
140                self._token_expiry = token.expires_on
141
142                logger.info("[OUT] _get_azure_ad_token: Azure AD token acquired successfully")
143                return self._access_token
144
145            except Exception as e:
146                logger.error(
147                    "[ERROR] _get_azure_ad_token: Failed to acquire Azure AD token - error_type=%s",
148                    type(e).__name__,
149                )
150                raise
151
152    async def _create_connection_params(self) -> Dict[str, Any]:
153        """接続パラメータを作成"""
154        logger.info("[IN] _create_connection_params: Creating database connection parameters")
155        params = {
156            "host": self.config.host,
157            "port": self.config.port,
158            "database": self.config.database,
159            "command_timeout": self.config.command_timeout,
160        }
161
162        # SSL設定（disable以外は設定する）
163        ssl_mode = getattr(self.config, 'ssl_mode', 'require')
164        if ssl_mode != "disable":
165            params["ssl"] = ssl_mode
166
167        if self.config.use_azure_ad:
168            token = await self._get_azure_ad_token()
169
170            # Use configured username if available, otherwise fallback to client_id
171            if self.config.azure_ad and self.config.azure_ad.username:
172                username = self.config.azure_ad.username
173            elif self.config.azure_ad:
174                username = self.config.azure_ad.client_id
175            else:
176                raise ValueError("Azure AD configuration is required when use_azure_ad is True")
177
178            params.update(
179                {
180                    "user": username,
181                    "password": token,
182                }
183            )
184        else:
185            # 標準認証の場合（ユーザー名/パスワード）
186            if not self.config.username or not self.config.password:
187                raise ValueError("Username and password are required when use_azure_ad is False")
188
189            params.update(
190                {
191                    "user": self.config.username,
192                    "password": self.config.password,
193                }
194            )
195
196        logger.info("[OUT] _create_connection_params: Connection parameters created successfully")
197        return params
198
199    async def _ensure_initialized(self) -> None:
200        """接続プールが初期化されていることを保証"""
201        if not self.pool:
202            await self.initialize()
203
204    async def initialize(self) -> None:
205        """データベース接続プールを初期化"""
206        logger.info("[IN] initialize: Initializing PostgreSQL connection pool")
207        try:
208            # トークンリフレッシュ用のLockを作成（イベントループ内で作成することで紐付けを保証）
209            if self._token_refresh_lock is None:
210                self._token_refresh_lock = asyncio.Lock()
211
212            logger.info("[IN] initialize.create_connection_params: Creating connection parameters")
213            connection_params = await self._create_connection_params()
214            logger.info("[OUT] initialize.create_connection_params: Connection parameters created")
215
216            logger.info(
217                "[IN] initialize.create_pool: Creating connection pool - min_size=%s, max_size=%s",
218                self.config.pool_min_size,
219                self.config.pool_max_size,
220            )
221            self.pool = await asyncpg.create_pool(
222                min_size=self.config.pool_min_size,
223                max_size=self.config.pool_max_size,
224                **connection_params,
225            )
226            logger.info("[OUT] initialize.create_pool: Connection pool created successfully")
227            logger.info("[OUT] initialize: PostgreSQL connection pool initialized successfully")
228
229        except Exception as e:
230            logger.error(
231                "[ERROR] initialize: PostgreSQL initialization failed - error_type=%s, error=%s",
232                type(e).__name__,
233                str(e),
234            )
235            raise
236
237    async def close(self) -> None:
238        """接続プールを閉じる
239
240        Note:
241            例外が発生してもリソースを確実にクリーンアップします。
242        """
243        logger.info("[IN] close: Closing PostgreSQL connection pool")
244        try:
245            if self.pool:
246                await self.pool.close()
247                logger.info("[OUT] close: PostgreSQL connection pool closed successfully")
248            else:
249                logger.info("[OUT] close: No connection pool to close")
250        except Exception as e:
251            logger.warning(f"Error closing PostgreSQL connection pool: {e}")
252        finally:
253            self.pool = None
254            self._access_token = None
255            self._token_expiry = None
256
257    async def execute_query(
258        self, query: str, params: tuple = ()
259    ) -> Dict[str, Any]:
260        """
261        汎用クエリ実行メソッド
262
263        SELECT/INSERT/UPDATE/DELETE等のSQLクエリを実行し、
264        結果を統一フォーマットで返します。
265
266        Args:
267            query: 実行するSQLクエリ。プレースホルダーは $1, $2 形式
268            params: クエリパラメータ（タプル）
269
270        Returns:
271            Dict[str, Any]: 実行結果
272                - success: bool - 実行成功/失敗
273                - data: List[Dict] (SELECT時) or Dict (INSERT/UPDATE/DELETE時)
274                - error: Optional[str] - エラーメッセージ（失敗時のみ）
275                - error_code: Optional[str] - エラーコード（失敗時のみ）
276
277        Example:
278            >>> # SELECTクエリ
279            >>> result = await db.execute_query(
280            ...     "SELECT * FROM users WHERE created_at > $1",
281            ...     (datetime(2024, 1, 1),)
282            ... )
283            >>> if result["success"]:
284            ...     for user in result["data"]:
285            ...         print(user["name"])
286            >>>
287            >>> # INSERT with RETURNING
288            >>> result = await db.execute_query(
289            ...     "INSERT INTO users (name) VALUES ($1) RETURNING id",
290            ...     ("John",)
291            ... )
292            >>> if result["success"]:
293            ...     new_id = result["data"][0]["id"]
294        """
295        logger.info(
296            "[IN] execute_query: Executing SQL query - query_type=%s, params_count=%s",
297            (query.strip().upper().split()[0] if query.strip() else "EMPTY"),
298            len(params),
299        )
300        try:
301            await self._ensure_initialized()
302
303            async with self.pool.acquire() as conn:
304                # クエリの種類を判定
305                normalized_query = query.strip().upper()
306                is_select = normalized_query.startswith("SELECT")
307                has_returning = "RETURNING" in normalized_query
308
309                if is_select or has_returning:
310                    # SELECT文またはRETURNING句がある場合はfetchを使用
311                    rows = await conn.fetch(query, *params)
312                    # Rowオブジェクトを辞書に変換
313                    result_data = [dict(row) for row in rows]
314                    logger.info(
315                        "[OUT] execute_query: Query executed successfully - result_count=%s",
316                        len(result_data),
317                    )
318                    return {"success": True, "data": result_data}
319                else:
320                    # INSERT/UPDATE/DELETE等（RETURNING句なし）の場合はexecuteを使用
321                    result = await conn.execute(query, *params)
322                    logger.info("[OUT] execute_query: Query executed successfully")
323                    return {"success": True, "data": {"result": result}}
324
325        except Exception as e:
326            logger.error(
327                "[ERROR] execute_query: Query execution failed - error_type=%s, error=%s",
328                type(e).__name__,
329                str(e),
330            )
331            return {
332                "success": False,
333                "data": None,
334                "error": f"Database query error: {str(e)}",
335                "error_code": "DB_QUERY_ERROR",
336            }
337
338    async def fetch_one(
339        self, query: str, params: tuple = ()
340    ) -> Optional[Dict[str, Any]]:
341        """
342        単一行を取得
343
344        1行だけ返すことが期待されるクエリに使用します。
345        結果が複数行ある場合でも最初の1行のみ返します。
346
347        Args:
348            query: 実行するSQLクエリ。プレースホルダーは $1, $2 形式
349            params: クエリパラメータ（タプル）
350
351        Returns:
352            Optional[Dict[str, Any]]: 結果行（辞書形式）、結果がない場合はNone
353
354        Example:
355            >>> # ユーザーを1件取得
356            >>> user = await db.fetch_one(
357            ...     "SELECT * FROM users WHERE id = $1",
358            ...     (user_id,)
359            ... )
360            >>> if user:
361            ...     print(f"Found: {user['name']}")
362            ... else:
363            ...     print("User not found")
364            >>>
365            >>> # 集計結果を取得
366            >>> count = await db.fetch_one("SELECT COUNT(*) as total FROM users")
367            >>> print(f"Total: {count['total']}")
368        """
369        await self._ensure_initialized()
370        async with self.pool.acquire() as conn:
371            row = await conn.fetchrow(query, *params)
372            return dict(row) if row else None
373
374    async def fetch_all(
375        self, query: str, params: tuple = ()
376    ) -> List[Dict[str, Any]]:
377        """
378        全行を取得
379
380        複数行を返すSELECTクエリに使用します。
381        結果がない場合は空のリストを返します。
382
383        Args:
384            query: 実行するSQLクエリ。プレースホルダーは $1, $2 形式
385            params: クエリパラメータ（タプル）
386
387        Returns:
388            List[Dict[str, Any]]: 結果行のリスト（各行は辞書形式）
389
390        Example:
391            >>> # 全ユーザーを取得
392            >>> users = await db.fetch_all("SELECT * FROM users ORDER BY created_at DESC")
393            >>> for user in users:
394            ...     print(f"{user['id']}: {user['name']}")
395            >>>
396            >>> # 条件付きで取得
397            >>> active_users = await db.fetch_all(
398            ...     "SELECT * FROM users WHERE status = $1 AND role = $2",
399            ...     ("active", "admin")
400            ... )
401            >>> print(f"Found {len(active_users)} admin users")
402        """
403        await self._ensure_initialized()
404        async with self.pool.acquire() as conn:
405            rows = await conn.fetch(query, *params)
406            return [dict(row) for row in rows]
407
408    async def execute(self, query: str, params: tuple = ()) -> str:
409        """
410        実行のみ（結果を返さないクエリ用）
411
412        INSERT/UPDATE/DELETE等の結果データが不要なクエリに使用します。
413        DDL文（CREATE TABLE等）の実行にも使用できます。
414
415        Args:
416            query: 実行するSQLクエリ。プレースホルダーは $1, $2 形式
417            params: クエリパラメータ（タプル）
418
419        Returns:
420            str: 実行結果のステータス（例: "INSERT 0 1", "UPDATE 5", "DELETE 3"）
421
422        Example:
423            >>> # データ挿入
424            >>> result = await db.execute(
425            ...     "INSERT INTO logs (level, message) VALUES ($1, $2)",
426            ...     ("INFO", "Operation completed")
427            ... )
428            >>> print(result)  # "INSERT 0 1"
429            >>>
430            >>> # データ更新
431            >>> result = await db.execute(
432            ...     "UPDATE users SET status = $1 WHERE last_login < $2",
433            ...     ("inactive", datetime(2024, 1, 1))
434            ... )
435            >>> print(result)  # "UPDATE 15"
436            >>>
437            >>> # データ削除
438            >>> result = await db.execute(
439            ...     "DELETE FROM sessions WHERE expired_at < $1",
440            ...     (datetime.now(),)
441            ... )
442            >>> print(result)  # "DELETE 100"
443        """
444        await self._ensure_initialized()
445        async with self.pool.acquire() as conn:
446            return await conn.execute(query, *params)

101@dataclass
102class PostgreSQLConfig:
103    """PostgreSQL設定
104
105    Note:
106        passwordはrepr=Falseでログ出力から除外されます。
107
108    Example:
109        >>> # 辞書から作成
110        >>> config = PostgreSQLConfig.from_dict({
111        ...     "host": "localhost",
112        ...     "port": 5432,
113        ...     "database": "mydb",
114        ...     "username": "user",
115        ...     "password": "pass"
116        ... })
117
118        >>> # TOMLファイルから作成
119        >>> config = PostgreSQLConfig.from_toml("config.toml", section="database")
120    """
121
122    provider: str = "postgresql"
123    host: str = "localhost"
124    port: int = 5432
125    database: str = "autonomousagent_db"
126    use_azure_ad: bool = False
127
128    # 標準認証用（use_azure_ad=Falseの場合）
129    username: Optional[str] = None
130    password: Optional[str] = field(default=None, repr=False)  # シークレット: reprから除外
131
132    # 接続プール設定
133    pool_min_size: int = 5
134    pool_max_size: int = 20
135    command_timeout: int = 60
136    pool_timeout: Optional[int] = 30
137    max_overflow: Optional[int] = 10
138
139    # Azure AD認証設定
140    azure_ad: Optional[AzureADConfig] = None
141
142    # API Proxy設定（CLIモード用）
143    api_proxy_url: Optional[str] = None
144
145    # SSL設定（デフォルト: require, ローカル開発用: disable）
146    ssl_mode: str = "require"
147
148    def __str__(self) -> str:
149        return (
150            f"PostgreSQLConfig(host={self.host!r}, port={self.port}, "
151            f"database={self.database!r}, use_azure_ad={self.use_azure_ad}, password=***)"
152        )
153
154    @classmethod
155    def from_dict(cls, data: Dict[str, Any]) -> "PostgreSQLConfig":
156        """辞書からPostgreSQLConfigを作成
157
158        Args:
159            data: 設定辞書
160
161        Returns:
162            PostgreSQLConfig instance
163
164        Example:
165            >>> config = PostgreSQLConfig.from_dict({
166            ...     "host": "db.example.com",
167            ...     "port": 5432,
168            ...     "database": "mydb",
169            ...     "username": "user",
170            ...     "password": "secret",
171            ...     "use_azure_ad": False,
172            ...     "pool_max_size": 30
173            ... })
174        """
175        # Azure AD設定がある場合はネストして処理
176        azure_ad = None
177        if "azure_ad" in data and data.get("use_azure_ad"):
178            azure_ad = AzureADConfig.from_dict(data["azure_ad"])
179
180        return cls(
181            provider=data.get("provider", "postgresql"),
182            host=data.get("host", "localhost"),
183            port=data.get("port", 5432),
184            database=data.get("database", "autonomousagent_db"),
185            use_azure_ad=data.get("use_azure_ad", False),
186            username=data.get("username"),
187            password=data.get("password"),
188            pool_min_size=data.get("pool_min_size", 5),
189            pool_max_size=data.get("pool_max_size", 20),
190            command_timeout=data.get("command_timeout", 60),
191            pool_timeout=data.get("pool_timeout", 30),
192            max_overflow=data.get("max_overflow", 10),
193            azure_ad=azure_ad,
194            api_proxy_url=data.get("api_proxy_url"),
195        )
196
197    @classmethod
198    def from_toml(cls, toml_path: str, section: str = "database") -> "PostgreSQLConfig":
199        """TOMLファイルからPostgreSQLConfigを作成
200
201        Args:
202            toml_path: TOMLファイルパス
203            section: セクション名（ドット区切りでネスト対応）
204
205        Returns:
206            PostgreSQLConfig instance
207
208        Example:
209            >>> # config.toml の [database] セクションを読み込み
210            >>> config = PostgreSQLConfig.from_toml("config.toml")
211
212            >>> # ネストされたセクション
213            >>> config = PostgreSQLConfig.from_toml("config.toml", section="services.database")
214        """
215        if not TOML_AVAILABLE:
216            raise ImportError("toml library is required: pip install toml")
217
218        with open(toml_path, "r") as f:
219            toml_data = toml.load(f)
220
221        # ネストされたセクションをドット区切りで辿る
222        data = toml_data
223        for key in section.split("."):
224            data = data[key]
225
226        return cls.from_dict(data)
227
228    @classmethod
229    def from_env(cls, prefix: str = "DB_") -> "PostgreSQLConfig":
230        """環境変数からPostgreSQLConfigを作成
231
232        環境変数名は {prefix}{FIELD_NAME} の形式で読み込みます。
233
234        Args:
235            prefix: 環境変数名のプレフィックス（デフォルト: "DB_"）
236
237        Returns:
238            PostgreSQLConfig instance
239
240        Example:
241            >>> # 環境変数から作成（DB_HOST, DB_PORT, DB_DATABASE, etc.）
242            >>> config = PostgreSQLConfig.from_env()
243
244            >>> # カスタムプレフィックス（POSTGRES_HOST, POSTGRES_PORT, etc.）
245            >>> config = PostgreSQLConfig.from_env(prefix="POSTGRES_")
246
247        Environment Variables:
248            {prefix}HOST: データベースホスト（デフォルト: localhost）
249            {prefix}PORT: ポート番号（デフォルト: 5432）
250            {prefix}DATABASE: データベース名（デフォルト: autonomousagent_db）
251            {prefix}USER / {prefix}USERNAME: ユーザー名
252            {prefix}PASSWORD: パスワード
253            {prefix}USE_AZURE_AD: Azure AD認証を使用するか（true/false）
254            {prefix}POOL_MIN_SIZE: 最小プールサイズ（デフォルト: 5）
255            {prefix}POOL_MAX_SIZE: 最大プールサイズ（デフォルト: 20）
256            {prefix}API_PROXY_URL: API Proxy URL（CLIモード用）
257        """
258        def get_env(key: str, default: str = "") -> str:
259            return os.environ.get(f"{prefix}{key}", default)
260
261        def get_env_int(key: str, default: int) -> int:
262            val = os.environ.get(f"{prefix}{key}")
263            return int(val) if val else default
264
265        def get_env_bool(key: str, default: bool = False) -> bool:
266            val = os.environ.get(f"{prefix}{key}", "").lower()
267            return val in ("true", "1", "yes") if val else default
268
269        # ユーザー名は USER または USERNAME どちらでも可
270        username = get_env("USER") or get_env("USERNAME")
271
272        return cls(
273            host=get_env("HOST", "localhost"),
274            port=get_env_int("PORT", 5432),
275            database=get_env("DATABASE", "autonomousagent_db"),
276            username=username or None,
277            password=get_env("PASSWORD") or None,
278            use_azure_ad=get_env_bool("USE_AZURE_AD"),
279            pool_min_size=get_env_int("POOL_MIN_SIZE", 5),
280            pool_max_size=get_env_int("POOL_MAX_SIZE", 20),
281            command_timeout=get_env_int("COMMAND_TIMEOUT", 60),
282            pool_timeout=get_env_int("POOL_TIMEOUT", 30),
283            max_overflow=get_env_int("MAX_OVERFLOW", 10),
284            api_proxy_url=get_env("API_PROXY_URL") or None,
285        )

21@dataclass
22class AzureADConfig:
23    """Azure AD認証設定
24
25    Note:
26        client_secretはrepr=Falseでログ出力から除外されます。
27
28    Example:
29        >>> # 辞書から作成
30        >>> config = AzureADConfig.from_dict({
31        ...     "tenant_id": "xxx",
32        ...     "client_id": "yyy",
33        ...     "client_secret": "zzz"
34        ... })
35
36        >>> # TOMLファイルから作成
37        >>> config = AzureADConfig.from_toml("config.toml", section="database.azure_ad")
38    """
39
40    tenant_id: str
41    client_id: str
42    client_secret: str = field(repr=False)  # シークレット: reprから除外
43    username: str = ""  # PostgreSQL Azure AD authentication username
44
45    def __str__(self) -> str:
46        return f"AzureADConfig(tenant_id={self.tenant_id!r}, client_id={self.client_id!r}, client_secret=***)"
47
48    @classmethod
49    def from_dict(cls, data: Dict[str, Any]) -> "AzureADConfig":
50        """辞書からAzureADConfigを作成
51
52        Args:
53            data: 設定辞書
54
55        Returns:
56            AzureADConfig instance
57
58        Example:
59            >>> config = AzureADConfig.from_dict({
60            ...     "tenant_id": "your-tenant-id",
61            ...     "client_id": "your-client-id",
62            ...     "client_secret": "your-secret"
63            ... })
64        """
65        return cls(
66            tenant_id=data["tenant_id"],
67            client_id=data["client_id"],
68            client_secret=data["client_secret"],
69            username=data.get("username", ""),
70        )
71
72    @classmethod
73    def from_toml(cls, toml_path: str, section: str = "azure_ad") -> "AzureADConfig":
74        """TOMLファイルからAzureADConfigを作成
75
76        Args:
77            toml_path: TOMLファイルパス
78            section: セクション名（ドット区切りでネスト対応）
79
80        Returns:
81            AzureADConfig instance
82
83        Example:
84            >>> # config.toml の [database.azure_ad] セクションを読み込み
85            >>> config = AzureADConfig.from_toml("config.toml", section="database.azure_ad")
86        """
87        if not TOML_AVAILABLE:
88            raise ImportError("toml library is required: pip install toml")
89
90        with open(toml_path, "r") as f:
91            toml_data = toml.load(f)
92
93        # ネストされたセクションをドット区切りで辿る
94        data = toml_data
95        for key in section.split("."):
96            data = data[key]
97
98        return cls.from_dict(data)