Skip to content

Commit

Permalink
Add support for Snowflake
Browse files Browse the repository at this point in the history
  • Loading branch information
@betodealmeida
betodealmeida committed Aug 14, 2024
1 parent e2a8519 commit 0328c48
Show file tree
Hide file tree
Showing 4 changed files with 157 additions and 36 deletions.
6 changes: 6 additions & 0 deletions superset/db_engine_specs/snowflake.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,12 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
supports_dynamic_schema = True
supports_catalog = supports_dynamic_catalog = True

# pylint: disable=invalid-name
encrypted_extra_sensitive_fields = [
"$.auth_params.privatekey_body",
"$.auth_params.privatekey_pass",
]

_time_grain_expressions = {
None: "{col}",
TimeGrain.SECOND: "DATE_TRUNC('SECOND', {col})",
Expand Down
42 changes: 24 additions & 18 deletions tests/unit_tests/db_engine_specs/test_bigquery.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,12 +212,14 @@ def test_unmask_encrypted_extra() -> None:
}
)

assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == {
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
)


def test_unmask_encrypted_extra_field_changeed() -> None:
Expand All @@ -243,12 +245,14 @@ def test_unmask_encrypted_extra_field_changeed() -> None:
}
)

assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == {
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)


def test_unmask_encrypted_extra_when_old_is_none() -> None:
Expand All @@ -267,12 +271,14 @@ def test_unmask_encrypted_extra_when_old_is_none() -> None:
}
)

assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == {
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
)


def test_unmask_encrypted_extra_when_new_is_none() -> None:
Expand Down
42 changes: 24 additions & 18 deletions tests/unit_tests/db_engine_specs/test_gsheets.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -295,12 +295,14 @@ def test_unmask_encrypted_extra() -> None:
}
)

assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == {
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
)


def test_unmask_encrypted_extra_field_changeed() -> None:
Expand All @@ -326,12 +328,14 @@ def test_unmask_encrypted_extra_field_changeed() -> None:
}
)

assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == {
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)


def test_unmask_encrypted_extra_when_old_is_none() -> None:
Expand All @@ -350,12 +354,14 @@ def test_unmask_encrypted_extra_when_old_is_none() -> None:
}
)

assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == {
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
)


def test_unmask_encrypted_extra_when_new_is_none() -> None:
Expand Down
103 changes: 103 additions & 0 deletions tests/unit_tests/db_engine_specs/test_snowflake.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,3 +291,106 @@ def test_get_default_catalog() -> None:
sqlalchemy_uri="snowflake://user:pass@account/database_name/default",
)
assert SnowflakeEngineSpec.get_default_catalog(database) == "database_name"


def test_mask_encrypted_extra() -> None:
"""
Test that the private keys are masked when the database is edited.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec

config = json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)

assert SnowflakeEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": "XXXXXXXXXX",
"privatekey_pass": "XXXXXXXXXX",
},
}
)


def test_mask_encrypted_extra_no_fields() -> None:
"""
Test that the private key is masked when the database is edited.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec

config = json.dumps(
{
# this is a fake example and the fields are made up
"auth_method": "token",
"auth_params": {
"jwt": "SECRET",
},
}
)

assert SnowflakeEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"auth_method": "token",
"auth_params": {
"jwt": "SECRET",
},
}
)


def test_unmask_encrypted_extra() -> None:
"""
Test that the private keys can be reused from the previous `encrypted_extra`.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec

old = json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)
new = json.dumps(
{
"foo": "bar",
"auth_method": "keypair",
"auth_params": {
"privatekey_body": "XXXXXXXXXX",
"privatekey_pass": "XXXXXXXXXX",
},
}
)

assert SnowflakeEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"foo": "bar",
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)

0 comments on commit 0328c48

Please sign in to comment.