1+ # Copyright 2025 Google LLC
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
115import functions_framework
216from spanner_client import SpannerClient
317from storage_client import StorageClient
418from flask import jsonify
519import logging
620import import_utils
21+ import os
22+ import sys
23+ from absl import flags
724
825logging .getLogger ().setLevel (logging .INFO )
926
10- _PROJECT_ID = 'datcom-import-automation-prod'
11- _SPANNER_PROJECT_ID = 'datcom-store'
12- _SPANNER_INSTANCE_ID = 'dc-kg-test'
13- _SPANNER_DATABASE_ID = 'dc_graph_import'
14- _GCS_BUCKET_ID = 'datcom-prod-imports'
15- _LOCATION = 'us-central1'
27+ FLAGS = flags .FLAGS
28+
29+ if 'project_id' not in FLAGS :
30+ flags .DEFINE_string ('project_id' ,
31+ os .environ .get ('PROJECT_ID' ,
32+ 'datcom-import-automation-prod' ),
33+ 'GCP Project ID' )
34+ flags .DEFINE_string ('spanner_project_id' ,
35+ os .environ .get ('SPANNER_PROJECT_ID' , 'datcom-store' ),
36+ 'Spanner Project ID' )
37+ flags .DEFINE_string ('spanner_instance_id' ,
38+ os .environ .get ('SPANNER_INSTANCE_ID' , 'dc-kg-test' ),
39+ 'Spanner Instance ID' )
40+ flags .DEFINE_string ('spanner_database_id' ,
41+ os .environ .get ('SPANNER_DATABASE_ID' , 'dc_graph_import' ),
42+ 'Spanner Database ID' )
43+ flags .DEFINE_string ('gcs_bucket_id' ,
44+ os .environ .get ('GCS_BUCKET_ID' , 'datcom-prod-imports' ),
45+ 'GCS Bucket ID' )
46+ flags .DEFINE_string ('location' ,
47+ os .environ .get ('LOCATION' , 'us-central1' ),
48+ 'GCP Location' )
49+
50+ # Parse flags to ensure they are accessible.
51+ # Use known_only=True to avoid conflicts with framework arguments.
52+ if not FLAGS .is_parsed ():
53+ FLAGS (sys .argv , known_only = True )
1654
1755
1856def _validate_params (request_json , required_params ):
@@ -36,9 +74,9 @@ def ingestion_helper(request):
3674 return (validation_error , 400 )
3775
3876 actionType = request_json ['actionType' ]
39- spanner = SpannerClient (_SPANNER_PROJECT_ID , _SPANNER_INSTANCE_ID ,
40- _SPANNER_DATABASE_ID )
41- storage = StorageClient (_GCS_BUCKET_ID )
77+ spanner = SpannerClient (FLAGS . spanner_project_id , FLAGS . spanner_instance_id ,
78+ FLAGS . spanner_database_id )
79+ storage = StorageClient (FLAGS . gcs_bucket_id )
4280
4381 if actionType == 'get_import_list' :
4482 # Gets the list of imports that are ready for ingestion.
@@ -76,7 +114,7 @@ def ingestion_helper(request):
76114 import_list = request_json ['importList' ]
77115 workflow_id = request_json ['workflowId' ]
78116 job_id = request_json ['jobId' ]
79- metrics = import_utils .get_ingestion_metrics (_PROJECT_ID , _LOCATION ,
117+ metrics = import_utils .get_ingestion_metrics (FLAGS . project_id , FLAGS . location ,
80118 job_id )
81119 spanner .update_ingestion_status (import_list , workflow_id , job_id ,
82120 metrics )
0 commit comments