def test_get_code_snippets():
ds = DataSetFactory.create(type=DataSetType.MASTER.value)
assert get_code_snippets(ds) == {}
SourceTableFactory.create(dataset=ds, schema="public", table="MY_LOVELY_TABLE")
snippets = get_code_snippets(ds)
assert """SELECT * FROM "public"."MY_LOVELY_TABLE" LIMIT 50""" in snippets['python']
assert """SELECT * FROM "public"."MY_LOVELY_TABLE" LIMIT 50""" in snippets['r']
assert snippets['sql'] == """SELECT * FROM "public"."MY_LOVELY_TABLE" LIMIT 50"""
def test_new_credentials_have_pgaudit_configuration(self):
ensure_databases_configured().handle()
user = UserFactory(email="*****@*****.**")
st = SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type=UserAccessType.REQUIRES_AUTHENTICATION))
source_tables = source_tables_for_user(user)
db_role_schema_suffix = db_role_schema_suffix_for_user(user)
user_creds_to_drop = new_private_database_credentials(
db_role_schema_suffix,
source_tables,
postgres_user(user.email),
user,
valid_for=datetime.timedelta(days=1),
)
connections[st.database.memorable_name].cursor().execute("COMMIT")
rolename = user_creds_to_drop[0]["db_user"]
query = f"SELECT rolname, rolconfig FROM pg_roles WHERE rolname = '{rolename}';"
with connections[st.database.memorable_name].cursor() as cursor:
cursor.execute(query)
results = cursor.fetchall()
assert "pgaudit.log=ALL" in results[0][1]
assert "pgaudit.log_catalog=off" in results[0][1]
def test_list_user_pagination(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
UserFactory.create(username='******')
UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.list_users.side_effect = [
{
"UserList": [{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}],
"NextToken":
"foo",
},
{
"UserList": [{
"Arn": "Arn2",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}]
},
]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
),
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
),
]
def test_missing_user_handled_gracefully(self, mock_cache,
mock_boto3_client, mock_creds):
# Arrange
user = UserFactory.create(username='******')
user2 = UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.describe_user.side_effect = [
botocore.exceptions.ClientError(
{
"Error": {
"Code": "ResourceNotFoundException",
"Message": "User not found",
}
},
'DescribeUser',
),
{
"User": {
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR"
}
},
]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
# Act
sync_quicksight_permissions(user_sso_ids_to_update=[
str(user.profile.sso_id),
str(user2.profile.sso_id)
])
# Assert
assert mock_user_client.describe_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user.profile.sso_id}',
),
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user2.profile.sso_id}',
),
]
assert len(mock_data_client.create_data_source.call_args_list) == 1
assert len(mock_data_client.update_data_source.call_args_list) == 0
def test_deletes_expired_and_unused_users(self):
ensure_databases_configured().handle()
user = UserFactory(email='*****@*****.**')
st = SourceTableFactory(
dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'
)
)
source_tables = source_tables_for_user(user)
db_role_schema_suffix = db_role_schema_suffix_for_user(user)
user_creds_to_drop = new_private_database_credentials(
db_role_schema_suffix,
source_tables,
postgres_user(user.email),
user,
valid_for=datetime.timedelta(days=31),
)
qs_creds_to_drop = new_private_database_credentials(
db_role_schema_suffix,
source_tables,
postgres_user(user.email, suffix='qs'),
user,
valid_for=datetime.timedelta(seconds=0),
)
qs_creds_to_keep = new_private_database_credentials(
db_role_schema_suffix,
source_tables,
postgres_user(user.email, suffix='qs'),
user,
valid_for=datetime.timedelta(minutes=1),
)
connections[st.database.memorable_name].cursor().execute('COMMIT')
# Make sure that `qs_creds_to_drop` has definitely expired
time.sleep(1)
with mock.patch('dataworkspace.apps.applications.utils.gevent.sleep'):
delete_unused_datasets_users()
with connections[st.database.memorable_name].cursor() as cursor:
cursor.execute(
"SELECT usename FROM pg_catalog.pg_user WHERE usename IN %s",
[
(
user_creds_to_drop[0]['db_user'],
qs_creds_to_drop[0]['db_user'],
qs_creds_to_keep[0]['db_user'],
)
],
)
assert cursor.fetchall() == [(qs_creds_to_keep[0]['db_user'],)]
def test_create_new_data_source(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.list_users.return_value = {
"UserList": [{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}]
}
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
)
]
assert mock_data_client.create_data_source.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId=mock.ANY,
Name=mock.ANY,
DataSourceParameters={
'AuroraPostgreSqlParameters': {
'Host': mock.ANY,
'Port': mock.ANY,
'Database': mock.ANY,
}
},
Credentials={
'CredentialPair': {
'Username': mock.ANY,
'Password': mock.ANY
}
},
VpcConnectionProperties={'VpcConnectionArn': mock.ANY},
Type='AURORA_POSTGRESQL',
Permissions=[{
'Principal':
'Arn',
'Actions': [
'quicksight:DescribeDataSource',
'quicksight:DescribeDataSourcePermissions',
'quicksight:PassDataSource',
],
}],
)
]
assert mock_data_client.update_data_source.call_args_list == []
assert sorted(
mock_data_client.delete_data_source.call_args_list,
key=lambda x: x.kwargs['DataSourceId'],
) == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId='data-workspace-dev-my_database-88f3887d',
),
mock.call(
AwsAccountId=mock.ANY,
DataSourceId='data-workspace-dev-test_external_db2-88f3887d',
),
]
def test_poll_until_user_created(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
user = UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.describe_user.side_effect = [
botocore.exceptions.ClientError(
{
"Error": {
"Code": "ResourceNotFoundException",
"Message": "User not found",
}
},
'DescribeUser',
),
] * 10 + [{
"User": {
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}
}]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
# Act
with mock.patch('dataworkspace.apps.applications.utils.gevent.sleep'):
sync_quicksight_permissions(
user_sso_ids_to_update=[str(user.profile.sso_id)],
poll_for_user_creation=True,
)
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
)
]
assert (mock_user_client.describe_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user.profile.sso_id}',
),
] * 11)
assert len(mock_data_client.create_data_source.call_args_list) == 1
assert len(mock_data_client.update_data_source.call_args_list) == 0
def test_create_new_data_source(self, mock_cache, mock_boto3_client, mock_creds):
# Arrange
UserFactory.create(username="******")
SourceTableFactory(
dataset=MasterDataSetFactory.create(
user_access_type=UserAccessType.REQUIRES_AUTHENTICATION
)
)
mock_user_client = mock.Mock()
mock_user_client.list_users.return_value = {
"UserList": [
{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}
]
}
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace="default",
Role="AUTHOR",
CustomPermissionsName="author-custom-permissions",
UserName="******",
Email="*****@*****.**",
)
]
assert mock_data_client.create_data_source.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId=mock.ANY,
Name=mock.ANY,
DataSourceParameters={
"AuroraPostgreSqlParameters": {
"Host": mock.ANY,
"Port": mock.ANY,
"Database": mock.ANY,
}
},
Credentials={"CredentialPair": {"Username": mock.ANY, "Password": mock.ANY}},
VpcConnectionProperties={"VpcConnectionArn": mock.ANY},
Type="AURORA_POSTGRESQL",
Permissions=[
{
"Principal": "Arn",
"Actions": [
"quicksight:DescribeDataSource",
"quicksight:DescribeDataSourcePermissions",
"quicksight:PassDataSource",
],
}
],
)
]
assert mock_data_client.update_data_source.call_args_list == []
assert sorted(
mock_data_client.delete_data_source.call_args_list,
key=lambda x: x.kwargs["DataSourceId"],
) == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId="data-workspace-dev-my_database-88f3887d",
),
mock.call(
AwsAccountId=mock.ANY,
DataSourceId="data-workspace-dev-test_external_db2-88f3887d",
),
]