bibibi12345 commited on
Commit
952cf47
·
1 Parent(s): ed723a3

fixing project id env variable loading

Browse files
Files changed (1) hide show
  1. src/auth.py +52 -6
src/auth.py CHANGED
@@ -84,8 +84,21 @@ def authenticate_user(request: Request):
84
  def save_credentials(creds, project_id=None):
85
  global credentials_from_env
86
 
87
- # Don't save to file if credentials came from environment variable
 
88
  if credentials_from_env:
 
 
 
 
 
 
 
 
 
 
 
 
89
  return
90
 
91
  creds_data = {
@@ -180,6 +193,12 @@ def get_credentials():
180
  credentials = Credentials.from_authorized_user_info(creds_data, SCOPES)
181
  credentials_from_env = True # Mark as environment credentials
182
 
 
 
 
 
 
 
183
  # Try to refresh if expired and refresh token exists
184
  if credentials.expired and credentials.refresh_token:
185
  try:
@@ -212,6 +231,12 @@ def get_credentials():
212
  credentials = Credentials.from_authorized_user_info(minimal_creds_data, SCOPES)
213
  credentials_from_env = True # Mark as environment credentials
214
 
 
 
 
 
 
 
215
  # Force refresh since we don't have a valid token
216
  try:
217
  logging.info("Refreshing minimal environment credentials...")
@@ -486,29 +511,41 @@ def get_user_project_id(creds):
486
  if user_project_id:
487
  return user_project_id
488
 
 
489
  env_project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
490
  if env_project_id:
 
491
  user_project_id = env_project_id
492
  save_credentials(creds, user_project_id)
493
  return user_project_id
494
 
 
495
  if os.path.exists(CREDENTIAL_FILE):
496
  try:
497
  with open(CREDENTIAL_FILE, "r") as f:
498
  creds_data = json.load(f)
499
  cached_project_id = creds_data.get("project_id")
500
  if cached_project_id:
 
501
  user_project_id = cached_project_id
502
  return user_project_id
503
  except Exception as e:
504
- pass
505
 
 
 
506
  if creds.expired and creds.refresh_token:
507
  try:
 
508
  creds.refresh(GoogleAuthRequest())
509
  save_credentials(creds)
 
510
  except Exception as e:
511
- raise Exception(f"Failed to refresh credentials while getting project ID: {str(e)}")
 
 
 
 
512
 
513
  headers = {
514
  "Authorization": f"Bearer {creds.token}",
@@ -522,6 +559,7 @@ def get_user_project_id(creds):
522
 
523
  try:
524
  import requests
 
525
  resp = requests.post(
526
  f"{CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist",
527
  data=json.dumps(probe_payload),
@@ -529,12 +567,20 @@ def get_user_project_id(creds):
529
  )
530
  resp.raise_for_status()
531
  data = resp.json()
532
- user_project_id = data.get("cloudaicompanionProject")
533
- if not user_project_id:
534
  raise ValueError("Could not find 'cloudaicompanionProject' in loadCodeAssist response.")
535
 
 
 
536
  save_credentials(creds, user_project_id)
537
 
538
  return user_project_id
539
  except requests.exceptions.HTTPError as e:
540
- raise
 
 
 
 
 
 
 
84
  def save_credentials(creds, project_id=None):
85
  global credentials_from_env
86
 
87
+ # Don't save credentials to file if they came from environment variable,
88
+ # but still save project_id if provided and no file exists or file lacks project_id
89
  if credentials_from_env:
90
+ if project_id and os.path.exists(CREDENTIAL_FILE):
91
+ try:
92
+ with open(CREDENTIAL_FILE, "r") as f:
93
+ existing_data = json.load(f)
94
+ # Only update project_id if it's missing from the file
95
+ if "project_id" not in existing_data:
96
+ existing_data["project_id"] = project_id
97
+ with open(CREDENTIAL_FILE, "w") as f:
98
+ json.dump(existing_data, f, indent=2)
99
+ logging.info(f"Added project_id {project_id} to existing credential file")
100
+ except Exception as e:
101
+ logging.warning(f"Could not update project_id in credential file: {e}")
102
  return
103
 
104
  creds_data = {
 
193
  credentials = Credentials.from_authorized_user_info(creds_data, SCOPES)
194
  credentials_from_env = True # Mark as environment credentials
195
 
196
+ # Extract project_id from environment credentials if available
197
+ if "project_id" in raw_env_creds_data:
198
+ global user_project_id
199
+ user_project_id = raw_env_creds_data["project_id"]
200
+ logging.info(f"Extracted project_id from environment credentials: {user_project_id}")
201
+
202
  # Try to refresh if expired and refresh token exists
203
  if credentials.expired and credentials.refresh_token:
204
  try:
 
231
  credentials = Credentials.from_authorized_user_info(minimal_creds_data, SCOPES)
232
  credentials_from_env = True # Mark as environment credentials
233
 
234
+ # Extract project_id from environment credentials if available
235
+ if "project_id" in raw_env_creds_data:
236
+ global user_project_id
237
+ user_project_id = raw_env_creds_data["project_id"]
238
+ logging.info(f"Extracted project_id from minimal environment credentials: {user_project_id}")
239
+
240
  # Force refresh since we don't have a valid token
241
  try:
242
  logging.info("Refreshing minimal environment credentials...")
 
511
  if user_project_id:
512
  return user_project_id
513
 
514
+ # Priority 1: Check environment variable first
515
  env_project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
516
  if env_project_id:
517
+ logging.info(f"Using project ID from GOOGLE_CLOUD_PROJECT environment variable: {env_project_id}")
518
  user_project_id = env_project_id
519
  save_credentials(creds, user_project_id)
520
  return user_project_id
521
 
522
+ # Priority 2: Check cached project ID in credential file
523
  if os.path.exists(CREDENTIAL_FILE):
524
  try:
525
  with open(CREDENTIAL_FILE, "r") as f:
526
  creds_data = json.load(f)
527
  cached_project_id = creds_data.get("project_id")
528
  if cached_project_id:
529
+ logging.info(f"Using cached project ID from credential file: {cached_project_id}")
530
  user_project_id = cached_project_id
531
  return user_project_id
532
  except Exception as e:
533
+ logging.warning(f"Could not read project_id from credential file: {e}")
534
 
535
+ # Priority 3: Make API call to discover project ID
536
+ # Ensure we have valid credentials for the API call
537
  if creds.expired and creds.refresh_token:
538
  try:
539
+ logging.info("Refreshing credentials before project ID discovery...")
540
  creds.refresh(GoogleAuthRequest())
541
  save_credentials(creds)
542
+ logging.info("Credentials refreshed successfully for project ID discovery")
543
  except Exception as e:
544
+ logging.error(f"Failed to refresh credentials while getting project ID: {e}")
545
+ # Continue with existing credentials - they might still work
546
+
547
+ if not creds.token:
548
+ raise Exception("No valid access token available for project ID discovery")
549
 
550
  headers = {
551
  "Authorization": f"Bearer {creds.token}",
 
559
 
560
  try:
561
  import requests
562
+ logging.info("Attempting to discover project ID via API call...")
563
  resp = requests.post(
564
  f"{CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist",
565
  data=json.dumps(probe_payload),
 
567
  )
568
  resp.raise_for_status()
569
  data = resp.json()
570
+ discovered_project_id = data.get("cloudaicompanionProject")
571
+ if not discovered_project_id:
572
  raise ValueError("Could not find 'cloudaicompanionProject' in loadCodeAssist response.")
573
 
574
+ logging.info(f"Discovered project ID via API: {discovered_project_id}")
575
+ user_project_id = discovered_project_id
576
  save_credentials(creds, user_project_id)
577
 
578
  return user_project_id
579
  except requests.exceptions.HTTPError as e:
580
+ logging.error(f"HTTP error during project ID discovery: {e}")
581
+ if hasattr(e, 'response') and e.response:
582
+ logging.error(f"Response status: {e.response.status_code}, body: {e.response.text}")
583
+ raise Exception(f"Failed to discover project ID via API: {e}")
584
+ except Exception as e:
585
+ logging.error(f"Unexpected error during project ID discovery: {e}")
586
+ raise Exception(f"Failed to discover project ID: {e}")