mirror of
https://github.com/not-lucky/GeminiKeyManagement.git
synced 2025-12-06 08:44:01 +05:30
refactor: enhance code structure and add Terms of Service handling for API interactions
This commit is contained in:
@@ -1,14 +1,16 @@
|
||||
"""This module contains functions for interacting with various Google Cloud Platform APIs."""
|
||||
"""
|
||||
Functions for interacting with Google Cloud Platform APIs.
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
import concurrent.futures
|
||||
from datetime import datetime, timezone
|
||||
from google.cloud import resourcemanager_v3, service_usage_v1, api_keys_v2
|
||||
from google.api_core import exceptions as google_exceptions
|
||||
from . import config, utils
|
||||
from . import config, utils, exceptions
|
||||
|
||||
def enable_api(project_id, credentials, dry_run=False):
|
||||
"""Enables the Generative Language API for a given project."""
|
||||
"""Enables the Generative Language API."""
|
||||
service_name = config.GENERATIVE_LANGUAGE_API
|
||||
service_path = f"projects/{project_id}/services/{service_name}"
|
||||
service_usage_client = service_usage_v1.ServiceUsageClient(credentials=credentials)
|
||||
@@ -21,7 +23,7 @@ def enable_api(project_id, credentials, dry_run=False):
|
||||
|
||||
enable_request = service_usage_v1.EnableServiceRequest(name=service_path)
|
||||
operation = service_usage_client.enable_service(request=enable_request)
|
||||
# This is a long-running operation, so we wait for it to complete.
|
||||
# Wait for the operation to complete.
|
||||
operation.result()
|
||||
logging.info(f" Successfully enabled Generative Language API for project {project_id}")
|
||||
return True
|
||||
@@ -30,17 +32,20 @@ def enable_api(project_id, credentials, dry_run=False):
|
||||
logging.warning(f" Permission denied to enable API for project {project_id}. Skipping.")
|
||||
return False
|
||||
except google_exceptions.GoogleAPICallError as err:
|
||||
if 'UREQ_TOS_NOT_ACCEPTED' in str(err):
|
||||
tos_url = "https://console.developers.google.com/terms/generative-language-api"
|
||||
raise exceptions.TermsOfServiceNotAcceptedError(
|
||||
f"Terms of Service for the Generative Language API have not been accepted for project {project_id}.",
|
||||
url=tos_url
|
||||
)
|
||||
logging.error(f" Error enabling API for project {project_id}: {err}")
|
||||
return False
|
||||
|
||||
def create_api_key(project_id, credentials, dry_run=False):
|
||||
"""
|
||||
Creates a new API key in the specified project.
|
||||
The key is restricted to only allow access to the Generative Language API.
|
||||
"""
|
||||
"""Creates a new, restricted API key."""
|
||||
if dry_run:
|
||||
logging.info(f" [DRY RUN] Would create API key for project {project_id}")
|
||||
# In a dry run, return a mock key object to allow the rest of the logic to proceed.
|
||||
# Return a mock key object for dry run
|
||||
return api_keys_v2.Key(
|
||||
name=f"projects/{project_id}/locations/global/keys/mock-key-id",
|
||||
uid="mock-key-id",
|
||||
@@ -79,7 +84,7 @@ def create_api_key(project_id, credentials, dry_run=False):
|
||||
return None
|
||||
|
||||
def delete_api_keys(project_id, credentials, dry_run=False):
|
||||
"""Deletes all API keys with the configured display name from a project."""
|
||||
"""Deletes all API keys with the display name 'Gemini API Key' and returns their UIDs."""
|
||||
deleted_keys_uids = []
|
||||
try:
|
||||
api_keys_client = api_keys_v2.ApiKeysClient(credentials=credentials)
|
||||
@@ -113,80 +118,3 @@ def delete_api_keys(project_id, credentials, dry_run=False):
|
||||
logging.error(f" An API error occurred while deleting keys for project {project_id}: {err}")
|
||||
return []
|
||||
|
||||
|
||||
|
||||
def _create_single_project(project_number, creds, dry_run, timeout_seconds=300, initial_delay=5):
|
||||
"""
|
||||
Creates a new GCP project and waits for it to be ready.
|
||||
Readiness is determined by successfully enabling the Generative Language API.
|
||||
"""
|
||||
random_string = utils.generate_random_string()
|
||||
project_id = f"project{project_number}-{random_string}"
|
||||
display_name = f"Project{project_number}"
|
||||
|
||||
logging.info(f"Attempting to create project: ID='{project_id}', Name='{display_name}'")
|
||||
|
||||
if dry_run:
|
||||
logging.info(f"[DRY RUN] Would create project '{display_name}' with ID '{project_id}'.")
|
||||
return None
|
||||
|
||||
try:
|
||||
resource_manager = resourcemanager_v3.ProjectsClient(credentials=creds)
|
||||
project_to_create = resourcemanager_v3.Project(
|
||||
project_id=project_id,
|
||||
display_name=display_name
|
||||
)
|
||||
operation = resource_manager.create_project(project=project_to_create)
|
||||
logging.info(f"Waiting for project creation operation for '{display_name}' to complete...")
|
||||
created_project = operation.result()
|
||||
logging.info(f"Successfully initiated creation for project '{display_name}'.")
|
||||
|
||||
# After creation, there can be a delay before the project is fully available
|
||||
# for API enablement. This loop polls until the API can be enabled.
|
||||
start_time = time.time()
|
||||
delay = initial_delay
|
||||
while time.time() - start_time < timeout_seconds:
|
||||
if enable_api(project_id, creds):
|
||||
logging.info(f"Generative AI API enabled for project '{display_name}' ({project_id}). Project is ready.")
|
||||
return created_project
|
||||
else:
|
||||
logging.info(f"Waiting for project '{display_name}' ({project_id}) to become ready... Retrying in {delay} seconds.")
|
||||
time.sleep(delay)
|
||||
delay = min(delay * 2, 30)
|
||||
|
||||
logging.error(f"Timed out waiting for project '{display_name}' ({project_id}) to become ready after {timeout_seconds} seconds.")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to create project '{display_name}': {e}")
|
||||
return None
|
||||
|
||||
def create_projects_if_needed(projects, creds, dry_run=False, max_workers=5):
|
||||
"""Creates new GCP projects in parallel until the account has at least 12 projects."""
|
||||
existing_project_count = len(projects)
|
||||
logging.info(f"Found {existing_project_count} existing projects.")
|
||||
newly_created_projects = []
|
||||
|
||||
if existing_project_count >= 12:
|
||||
logging.info("Account already has 12 or more projects. No new projects will be created.")
|
||||
return newly_created_projects
|
||||
|
||||
projects_to_create_count = 12 - existing_project_count
|
||||
logging.info(f"Need to create {projects_to_create_count} more projects.")
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_project_number = {
|
||||
executor.submit(_create_single_project, str(i + 1).zfill(2), creds, dry_run): i
|
||||
for i in range(existing_project_count, 12)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_project_number):
|
||||
try:
|
||||
created_project = future.result()
|
||||
if created_project:
|
||||
newly_created_projects.append(created_project)
|
||||
except Exception as exc:
|
||||
project_number = future_to_project_number[future]
|
||||
logging.error(f"Project number {project_number} generated an exception: {exc}", exc_info=True)
|
||||
|
||||
return newly_created_projects
|
||||
|
||||
Reference in New Issue
Block a user