diff --git a/CHANGES.rst b/CHANGES.rst index 5e87bfa1..23483b6b 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -5,6 +5,10 @@ Changes for croud Unreleased ========== +- Added support for organization secrets of type Azure. + +- Added support for importing jobs from Azure Blob Storage. + 1.9.0 - 2023/12/04 ================== diff --git a/croud/__main__.py b/croud/__main__.py index 7ed88151..e9570ac8 100644 --- a/croud/__main__.py +++ b/croud/__main__.py @@ -51,6 +51,7 @@ export_jobs_delete, export_jobs_list, import_job_progress, + import_jobs_create_from_azure_blob_storage, import_jobs_create_from_file, import_jobs_create_from_s3, import_jobs_create_from_url, @@ -809,7 +810,8 @@ Argument( "--file-path", type=str, required=True, help="The absolute path in the S3 bucket that " - "points to the file to be imported." + "points to the file to be imported. " + "Globbing (use of *) is allowed." ), Argument( "--secret-id", type=str, required=True, @@ -824,6 +826,32 @@ ] + import_job_create_common_args, "resolver": import_jobs_create_from_s3, }, + "from-azure-blob-storage": { + "help": "Create a data import job on the specified " + "cluster from an Azure blob storage location.", + "extra_args": [ + # Type Azure Blob Storage params + Argument( + "--container-name", type=str, + required=True, + help="The name of the storage container " + "where the file to be imported is located." + ), + Argument( + "--blob-name", type=str, required=True, + help="The absolute path in the storage " + "container that points to the file to be " + "imported. Globbing (use of *) is allowed." + ), + Argument( + "--secret-id", type=str, required=True, + help="The secret that contains the access key " + "and secret key needed to access the file " + "to be imported." + ), + ] + import_job_create_common_args, + "resolver": import_jobs_create_from_azure_blob_storage, + }, }, }, "progress": { @@ -1114,18 +1142,25 @@ help="The name the secret will be known as.", ), Argument( - "--type", type=str, required=True, choices=["AWS"], - help="The type of Secret. Currently only AWS type is " - "supported.", + "--type", type=str, required=True, + choices=["AWS", "AZURE"], + help="The type of Secret. Either AWS or Azure.", ), + # AWS arguments Argument( - "--access-key", type=str, required=True, + "--access-key", type=str, required=False, help="For an AWS type secret, the access key ID.", ), Argument( - "--secret-key", type=str, required=True, + "--secret-key", type=str, required=False, help="For an AWS type secret, the secret key.", ), + # Azure arguments + Argument( + "--connection-string", type=str, required=False, + help="For an Azure type secret, the connection string " + "or URL that grants access to a resource.", + ), ], "resolver": org_secrets_create, }, diff --git a/croud/clusters/commands.py b/croud/clusters/commands.py index c2936e47..cf5c0d25 100644 --- a/croud/clusters/commands.py +++ b/croud/clusters/commands.py @@ -200,6 +200,18 @@ def import_jobs_create_from_s3(args: Namespace) -> None: import_jobs_create(args, extra_payload=extra_body) +def import_jobs_create_from_azure_blob_storage(args: Namespace) -> None: + extra_body = { + "azureblob": { + "container_name": args.container_name, + "blob_name": args.blob_name, + "secret_id": args.secret_id, + } + } + args.type = "azureblob" + import_jobs_create(args, extra_payload=extra_body) + + def _get_org_id_from_cluster_id(client, cluster_id: str) -> Optional[str]: data, errors = client.get(f"/api/v2/clusters/{cluster_id}/") if errors or not data: diff --git a/croud/organizations/commands.py b/croud/organizations/commands.py index 6c1aa760..336544f8 100644 --- a/croud/organizations/commands.py +++ b/croud/organizations/commands.py @@ -162,11 +162,28 @@ def org_secrets_create(args: Namespace) -> None: payload = { "name": args.name, "type": args.type, - "data": { + } + + if args.type == "AWS": + if not args.access_key or not args.secret_key: + print_error( + "Both access_key and secret_key are required for secret type AWS." + ) + return + payload["data"] = { "access_key": args.access_key, "secret_key": args.secret_key, - }, - } + } + elif args.type == "AZURE": + if not args.connection_string: + print_error("Argument connection-string is required for secret type Azure.") + return + payload["data"] = { + "azure_secret": { + "connection_string": args.connection_string, + } + } + data, errors = client.post( f"/api/v2/organizations/{args.org_id}/secrets/", body=payload ) diff --git a/tests/commands/test_clusters.py b/tests/commands/test_clusters.py index ab37f8bc..ac13b316 100644 --- a/tests/commands/test_clusters.py +++ b/tests/commands/test_clusters.py @@ -1644,6 +1644,58 @@ def test_import_job_create_from_file(mock_request): ) +@mock.patch.object( + Client, "request", return_value=({"id": "1", "status": "SUCCEEDED"}, None) +) +def test_import_job_create_from_azure_blob_storage(mock_request): + cluster_id = gen_uuid() + secret_id = gen_uuid() + container_name = "my-container-name" + blob_name = "path/to/my/files/*.csv.gz" + + call_command( + "croud", + "clusters", + "import-jobs", + "create", + "from-azure-blob-storage", + "--cluster-id", + cluster_id, + "--file-format", + "csv", + "--compression", + "gzip", + "--table", + "my-table", + "--create-table", + "false", + "--container-name", + container_name, + "--blob-name", + blob_name, + "--secret-id", + secret_id, + ) + body = { + "type": "azureblob", + "azureblob": { + "container_name": container_name, + "blob_name": blob_name, + "secret_id": secret_id, + }, + "format": "csv", + "destination": {"table": "my-table", "create_table": False}, + "compression": "gzip", + } + assert_rest( + mock_request, + RequestMethod.POST, + f"/api/v2/clusters/{cluster_id}/import-jobs/", + body=body, + any_times=True, + ) + + @mock.patch.object( Client, "request", diff --git a/tests/commands/test_organizations.py b/tests/commands/test_organizations.py index e127ae3b..dffd53d1 100644 --- a/tests/commands/test_organizations.py +++ b/tests/commands/test_organizations.py @@ -464,7 +464,7 @@ def test_organizations_secrets_list(mock_request): @mock.patch.object(Client, "request", return_value=({}, None)) -def test_organizations_secrets_create(mock_request): +def test_organizations_secrets_aws_create(mock_request): org_id = gen_uuid() name = "my_secret" secret_type = "AWS" @@ -499,6 +499,41 @@ def test_organizations_secrets_create(mock_request): ) +@mock.patch.object(Client, "request", return_value=({}, None)) +def test_organizations_secrets_azure_create(mock_request): + org_id = gen_uuid() + name = "my_secret" + secret_type = "AZURE" + connection_string = "https://my-storage-account.blob.core.windows.net/my-container?my-auth-params" # noqa + + call_command( + "croud", + "organizations", + "secrets", + "create", + "--org-id", + org_id, + "--name", + name, + "--type", + secret_type, + "--connection-string", + connection_string, + ) + assert_rest( + mock_request, + RequestMethod.POST, + f"/api/v2/organizations/{org_id}/secrets/", + body={ + "name": name, + "type": secret_type, + "data": { + "azure_secret": {"connection_string": connection_string}, + }, + }, + ) + + @mock.patch.object(Client, "request", return_value=({}, None)) def test_organizations_secrets_delete(mock_request): org_id = gen_uuid()