Skip to content

Commit

Permalink
Support for Azure secrets and blob Storage import jobs (#476)
Browse files Browse the repository at this point in the history 
* Added support for organization secrets of type Azure

* Added support for importing jobs from Azure Blob Storage
  • Loading branch information
@juanpardo
juanpardo authored Dec 15, 2023
1 parent 110298b commit f997c4b
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 10 deletions.
4 changes: 4 additions & 0 deletions CHANGES.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
==================

Expand Down
47 changes: 41 additions & 6 deletions croud/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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": {
Expand Down Expand Up @@ -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,
},
Expand Down
12 changes: 12 additions & 0 deletions croud/clusters/commands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
23 changes: 20 additions & 3 deletions croud/organizations/commands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
52 changes: 52 additions & 0 deletions tests/commands/test_clusters.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
37 changes: 36 additions & 1 deletion tests/commands/test_organizations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit f997c4b

Please sign in to comment.