def test_get_credentials(self, mock_credentials, mock_auth_default,
mock_service_account):
mock_service_account.return_value = mock_credentials
mock_auth_default.return_value = mock_credentials, 'mock_tuple_value'
credentials = auth_service.get_credentials()
assert credentials is not None
assert credentials.service_account_email is not None
assert credentials.signer is not None
assert credentials._token_uri is not None
assert credentials.project_id is not None
assert credentials.scopes is not None
def gcs_upload_file(project_id, bucket_name, prefix, upload_file):
"""
Uploads a dag file to a GCS bucket.
Args:
project_id (string): GCP Project Id of the Cloud Composer instance
bucket_name (string): The name of the bucket (excluding any prefixes) where the dag is to be uploaded
prefix (string): The prefix of the GCS bucket where the dag is to be uploaded
upload_file (string): Path to the dag file to be uploaded
"""
logger.log(
logging.DEBUG,
f"Upload dag to GCS: project_id {project_id}, bucket_name {bucket_name}, prefix {prefix}, upload_file {upload_file}"
)
credentials = auth_service.get_credentials()
client = storage.Client(project_id, credentials=credentials)
bucket = client.bucket(bucket_name)
upload_file_name = os.path.basename(os.path.normpath(upload_file))
blob = bucket.blob(prefix + upload_file_name)
blob.upload_from_filename(upload_file)
def list_dags(project_id, bucket_name):
"""
Lists the dag files contained in the Cloud Storage bucket location of a Cloud Composer environment.
Args:
project_id (string): GCP Project Id of the Cloud Composer instance
bucket_name (string): The bucket name of the GCS location of the Cloud Composer dag files
(without the /dags prefix)
Returns:
a list of dag files contained in the Cloud Storage bucket location of a Cloud Composer environment
"""
logger.log(logging.DEBUG, "Listing the dags")
credentials = auth_service.get_credentials()
client = storage.Client(project_id, credentials=credentials)
blobs = client.list_blobs(bucket_name, prefix="dags")
dag_list = []
for blob in blobs:
if blob.name.endswith(".py"):
dag_name = blob.name.replace("dags/", "").replace(".py", "")
dag_list.append(dag_name)
return dag_list
def gcs_download_file(project_id, bucket_name, download_file):
"""
Downloads a dag file from a GCS bucket.
Args:
project_id (string): GCP Project Id of the Cloud Composer instance
bucket_name (string): The name of the bucket (including any prefixes) where the dag is located
download_file (string): Name of the dag file to be downloaded
Returns:
the absolute file path to the downloaded file
"""
logger.log(
logging.DEBUG,
f"Downloading dag from GCS: project_id {project_id}, bucket_name {bucket_name}, download_file {download_file}"
)
credentials = auth_service.get_credentials()
client = storage.Client(project_id, credentials=credentials)
bucket = client.bucket(bucket_name)
blob = bucket.blob(download_file)
temp_dir = tempfile.gettempdir()
download_file_path = os.path.join(temp_dir, f"{os.path.basename(os.path.normpath(download_file))}")
logger.log(logging.DEBUG, f"Local download path: {download_file_path}")
blob.download_to_filename(download_file_path)
return download_file_path