"""Account management module""" import csv import json from pathlib import Path from datetime import datetime, timedelta from typing import List, Dict, Optional, Set from ..common.cache import CacheManager from ..data.bigquery import BigQueryClient from ..data.models import AccountInfo class AccountManager: """账户管理器""" def __init__(self, account_details_file: str = None, cache_dir: str = None, config_manager = None): """初始化账户管理器 Args: account_details_file: 账户详情CSV文件路径 cache_dir: 缓存目录 config_manager: 配置管理器实例 """ # 如果没有指定文件路径，使用默认路径 if account_details_file is None: # 查找 aws_account_details.csv 文件 current_file = Path(__file__) aws_dir = current_file.parent.parent.parent # 向上三级到 aws/ 目录 account_details_file = aws_dir / 'aws_account_details.csv' self.account_details_file = account_details_file self.account_details = {} self.engineering_accounts = set() self.config_manager = config_manager # 初始化缓存管理器 if cache_dir: self.cache = CacheManager(cache_dir, default_expiry_hours=48) else: self.cache = CacheManager(default_expiry_hours=48) # 初始化BigQuery客户端 self.bq_client = BigQueryClient() # 加载账户详情 self.load_account_details() def get_account_details(self) -> Dict: """获取所有账户详情 Returns: 账户详情字典 """ return self.account_details def load_account_details(self): """从CSV文件加载账户详情""" csv_path = Path(self.account_details_file) if not csv_path.exists(): print(f"⚠️ 账户详情文件不存在: {csv_path}") return try: with open(csv_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: account_id = row.get('account_id', '').strip() if account_id: # 创建账户信息对象 account_info = AccountInfo.from_csv_row(row) self.account_details[account_id] = account_info # 所有在 aws_account_details.csv 中的账户都是 Engineering Group 账户 self.engineering_accounts.add(account_id) print(f"✓ 加载了 {len(self.account_details)} 个账户详情") print(f"✓ 共 {len(self.engineering_accounts)} 个 Engineering Group 账户") except Exception as e: print(f"✗ 加载账户详情失败: {e}") def is_engineering_account(self, account_id: str) -> bool: """判断是否为Engineering Group账户 所有在 aws_account_details.csv 中的账户都是 Engineering Group 账户 """ return account_id in self.account_details def get_account_info(self, account_id: str) -> Optional[AccountInfo]: """获取账户信息 Args: account_id: 账户ID Returns: 账户信息对象 """ return self.account_details.get(account_id) def classify_account(self, account_id: str) -> str: """分类账户环境 Args: account_id: 账户ID Returns: 环境名称 """ account_info = self.get_account_info(account_id) if not account_info: return 'unknown' env_value = (account_info.environment or '').strip().lower() # 生产环境标签 prod_labels = {'prod', 'production'} if env_value in prod_labels: return 'prod' # 公共测试账户单独归类 if account_info.is_shared: return 'shared' # 默认视为非生产环境 nonprod_labels = { 'nonprod', 'non-prod', 'dev', 'development', 'staging', 'test', 'testing', 'qa' } if env_value in nonprod_labels or not env_value: return 'nonprod' return 'unknown' def get_nonprod_accounts(self, refresh_cache: bool = False, validate_tables: bool = True) -> List[Dict]: """获取非生产环境账户列表 Args: refresh_cache: 是否刷新缓存 Returns: 账户列表 """ print("正在查询非生产账户列表...") env_config = self.config_manager.get_environment_config('nonprod') if self.config_manager else {} default_dataset = env_config.get('dataset_name', 'gcp-digital-bi.aws_nonprod_billing_native') table_pattern = env_config.get('table_pattern', '{account_id}_aws_billing_data') accounts = [] for account_id, account_info in self.account_details.items(): if self.classify_account(account_id) != 'nonprod': continue dataset = account_info.dataset_name or default_dataset table_name = account_info.table_name or table_pattern.format(account_id=account_id) if not dataset: print(f"⚠️ 非生产账户缺少dataset配置，跳过: {account_id}") continue if validate_tables: if not self._table_exists_cached(dataset, table_name, refresh_cache): print(f"⚠️ 表不存在，跳过账户 {account_id}: {dataset}.{table_name}") continue if not self._table_has_recent_data_cached( dataset, table_name, datetime.now() - timedelta(days=90), datetime.now(), refresh_cache, ): print(f"⚠️ 表无最近90天数据，跳过账户 {account_id}: {dataset}.{table_name}") continue accounts.append({ 'account_id': account_id, 'account_name': account_info.account_name, 'environment': 'nonprod', 'is_shared': account_info.is_shared, 'dataset_name': dataset, 'table_name': table_name }) print(f"✓ 发现非生产账户: {account_id}") print(f"✓ 发现 {len(accounts)} 个非生产账户") return accounts def get_prod_accounts(self, refresh_cache: bool = False, validate_tables: bool = True) -> List[Dict]: """获取生产环境账户列表 Args: refresh_cache: 是否刷新缓存 Returns: 账户列表 """ env_config = self.config_manager.get_environment_config('prod') if self.config_manager else {} default_dataset = env_config.get('dataset_name', 'gcp-digital-bi.aws_prod_billing') table_pattern = env_config.get('table_pattern', '{account_id}_aws_billing_data') accounts = [] if validate_tables: print("正在验证生产账户...") else: print("正在加载生产账户（跳过表存在性校验）...") for account_id, account_info in self.account_details.items(): if self.classify_account(account_id) != 'prod': continue dataset = account_info.dataset_name or default_dataset table_name = account_info.table_name or table_pattern.format(account_id=account_id) if not dataset: print(f" ⚠️ 生产账户缺少dataset配置，跳过: {account_id}") continue if validate_tables: if not self._table_exists_cached(dataset, table_name, refresh_cache): print(f" ✗ {account_id}: 表不存在 {dataset}.{table_name}") continue print(f" ✓ {account_id}: {table_name}") accounts.append({ 'account_id': account_id, 'account_name': account_info.account_name, 'environment': 'prod', 'table_name': table_name, 'dataset_name': dataset }) print(f"✓ 发现 {len(accounts)} 个生产账户") return accounts def get_shared_accounts(self) -> List[Dict]: """获取共享测试账户列表 Returns: 共享账户列表 """ shared_accounts = [] for account_id, account_info in self.account_details.items(): if account_info.is_shared: dataset = account_info.dataset_name table_name = account_info.table_name or f"{account_id}_aws_billing_data" shared_accounts.append({ 'account_id': account_id, 'account_name': account_info.account_name, 'environment': 'shared', 'is_shared': True, 'dataset_name': dataset, 'table_name': table_name }) return shared_accounts def validate_table_has_data(self, dataset: str, table_name: str, start_date: datetime, end_date: datetime) -> bool: """验证表在指定日期范围内是否有数据 Args: dataset: 数据集名称 table_name: 表名 start_date: 开始日期 end_date: 结束日期 Returns: 是否有数据 """ query = f""" SELECT COUNT(*) as row_count FROM `{dataset}.{table_name}` WHERE DATE(line_item_usage_start_date) >= '{start_date.strftime('%Y-%m-%d')}' AND DATE(line_item_usage_start_date) <= '{end_date.strftime('%Y-%m-%d')}' LIMIT 1 """ try: results = list(self.bq_client.execute_query(query)) if results: return results[0].get('row_count', 0) > 0 except Exception: pass return False def _table_exists_cached(self, dataset: str, table_name: str, refresh_cache: bool) -> bool: cache_key = f"table_exists:{dataset}.{table_name}" if not refresh_cache: cached = self.cache.get(cache_key) if isinstance(cached, bool): return cached exists = self.bq_client.table_exists(dataset, table_name) self.cache.set(cache_key, exists, expiry_hours=24) return exists def _table_has_recent_data_cached(self, dataset: str, table_name: str, start_date: datetime, end_date: datetime, refresh_cache: bool) -> bool: cache_key = f"table_recent:{dataset}.{table_name}" if not refresh_cache: cached = self.cache.get(cache_key) if isinstance(cached, bool): return cached has_data = self.validate_table_has_data(dataset, table_name, start_date, end_date) self.cache.set(cache_key, has_data, expiry_hours=24) return has_data def get_accounts_by_environment(self, environment: str, refresh_cache: bool = False, validate_tables: bool = True) -> List[Dict]: """根据环境获取账户列表 Args: environment: 环境名称 refresh_cache: 是否刷新缓存 Returns: 账户列表 """ if environment == 'nonprod': return self.get_nonprod_accounts(refresh_cache, validate_tables) elif environment == 'prod': return self.get_prod_accounts(refresh_cache, validate_tables) elif environment == 'shared': return self.get_shared_accounts() else: return [] def get_all_accounts(self, refresh_cache: bool = False, validate_tables: bool = True) -> Dict[str, List[Dict]]: """获取所有环境的账户 Args: refresh_cache: 是否刷新缓存 Returns: 按环境分组的账户字典 """ return { 'nonprod': self.get_nonprod_accounts(refresh_cache, validate_tables), 'prod': self.get_prod_accounts(refresh_cache, validate_tables), 'shared': self.get_shared_accounts(), }