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