diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 5a8354fd78dc7..d80625dc0c385 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1618,7 +1618,7 @@ def get_impersonation_key(cls, user: Optional[User]) -> Any: return user.username if user else None @classmethod - def mask_encrypted_extra(cls, encrypted_extra: str) -> str: + def mask_encrypted_extra(cls, encrypted_extra: Optional[str]) -> Optional[str]: """ Mask ``encrypted_extra``. @@ -1631,7 +1631,9 @@ def mask_encrypted_extra(cls, encrypted_extra: str) -> str: # pylint: disable=unused-argument @classmethod - def unmask_encrypted_extra(cls, old: str, new: str) -> str: + def unmask_encrypted_extra( + cls, old: Optional[str], new: Optional[str] + ) -> Optional[str]: """ Remove masks from ``encrypted_extra``. diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 4016c8f08f524..38b1f92023a10 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -396,10 +396,13 @@ def get_parameters_from_uri( raise ValidationError("Invalid service credentials") @classmethod - def mask_encrypted_extra(cls, encrypted_extra: str) -> str: + def mask_encrypted_extra(cls, encrypted_extra: Optional[str]) -> Optional[str]: + if encrypted_extra is None: + return encrypted_extra + try: config = json.loads(encrypted_extra) - except json.JSONDecodeError: + except (json.JSONDecodeError, TypeError): return encrypted_extra try: @@ -410,14 +413,19 @@ def mask_encrypted_extra(cls, encrypted_extra: str) -> str: return json.dumps(config) @classmethod - def unmask_encrypted_extra(cls, old: str, new: str) -> str: + def unmask_encrypted_extra( + cls, old: Optional[str], new: Optional[str] + ) -> Optional[str]: """ Reuse ``private_key`` if available and unchanged. """ + if old is None or new is None: + return new + try: old_config = json.loads(old) new_config = json.loads(new) - except json.JSONDecodeError: + except (TypeError, json.JSONDecodeError): return new if "credentials_info" not in new_config: diff --git a/superset/db_engine_specs/gsheets.py b/superset/db_engine_specs/gsheets.py index 3ee284788c245..3562e0d0b1384 100644 --- a/superset/db_engine_specs/gsheets.py +++ b/superset/db_engine_specs/gsheets.py @@ -140,10 +140,13 @@ def get_parameters_from_uri( raise ValidationError("Invalid service credentials") @classmethod - def mask_encrypted_extra(cls, encrypted_extra: str) -> str: + def mask_encrypted_extra(cls, encrypted_extra: Optional[str]) -> Optional[str]: + if encrypted_extra is None: + return encrypted_extra + try: config = json.loads(encrypted_extra) - except json.JSONDecodeError: + except (TypeError, json.JSONDecodeError): return encrypted_extra try: @@ -154,14 +157,19 @@ def mask_encrypted_extra(cls, encrypted_extra: str) -> str: return json.dumps(config) @classmethod - def unmask_encrypted_extra(cls, old: str, new: str) -> str: + def unmask_encrypted_extra( + cls, old: Optional[str], new: Optional[str] + ) -> Optional[str]: """ Reuse ``private_key`` if available and unchanged. """ + if old is None or new is None: + return new + try: old_config = json.loads(old) new_config = json.loads(new) - except json.JSONDecodeError: + except (TypeError, json.JSONDecodeError): return new if "service_account_info" not in new_config: diff --git a/superset/models/core.py b/superset/models/core.py index 8c1b58171d470..ca15137935ef9 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -248,7 +248,7 @@ def backend(self) -> str: return sqlalchemy_url.get_backend_name() @property - def masked_encrypted_extra(self) -> str: + def masked_encrypted_extra(self) -> Optional[str]: return self.db_engine_spec.mask_encrypted_extra(self.encrypted_extra) @property @@ -261,10 +261,12 @@ def parameters(self) -> Dict[str, Any]: masked_encrypted_extra = db_engine_spec.mask_encrypted_extra( self.encrypted_extra ) - try: - encrypted_config = json.loads(masked_encrypted_extra) - except (TypeError, json.JSONDecodeError): - encrypted_config = {} + encrypted_config = {} + if masked_encrypted_extra is not None: + try: + encrypted_config = json.loads(masked_encrypted_extra) + except (TypeError, json.JSONDecodeError): + pass try: # pylint: disable=useless-suppression diff --git a/tests/unit_tests/db_engine_specs/test_bigquery.py b/tests/unit_tests/db_engine_specs/test_bigquery.py index 13f7fd3150fc4..8c33489faf598 100644 --- a/tests/unit_tests/db_engine_specs/test_bigquery.py +++ b/tests/unit_tests/db_engine_specs/test_bigquery.py @@ -186,9 +186,61 @@ def test_unmask_encrypted_extra() -> None: } ) - assert json.loads(BigQueryEngineSpec.unmask_encrypted_extra(old, new)) == { + assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == { "credentials_info": { "project_id": "yellow-unicorn-314419", "private_key": "SECRET", }, } + + +def test_unmask_encrypted_extra_when_empty() -> None: + """ + Test that a None value works for ``encrypted_extra``. + """ + from superset.db_engine_specs.bigquery import BigQueryEngineSpec + + old = None + new = json.dumps( + { + "credentials_info": { + "project_id": "yellow-unicorn-314419", + "private_key": "XXXXXXXXXX", + }, + } + ) + + assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == { + "credentials_info": { + "project_id": "yellow-unicorn-314419", + "private_key": "XXXXXXXXXX", + }, + } + + +def test_unmask_encrypted_extra_when_new_is_empty() -> None: + """ + Test that a None value works for ``encrypted_extra``. + """ + from superset.db_engine_specs.bigquery import BigQueryEngineSpec + + old = json.dumps( + { + "credentials_info": { + "project_id": "black-sanctum-314419", + "private_key": "SECRET", + }, + } + ) + new = None + + assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) is None + + +def test_mask_encrypted_extra_when_empty() -> None: + """ + Test that the encrypted extra will return a none value if the field is empty. + """ + from superset.db_engine_specs.bigquery import BigQueryEngineSpec + + assert BigQueryEngineSpec.mask_encrypted_extra(None) is None diff --git a/tests/unit_tests/db_engine_specs/test_gsheets.py b/tests/unit_tests/db_engine_specs/test_gsheets.py index 219f259c1f39f..a226653ed57e0 100644 --- a/tests/unit_tests/db_engine_specs/test_gsheets.py +++ b/tests/unit_tests/db_engine_specs/test_gsheets.py @@ -228,9 +228,52 @@ def test_unmask_encrypted_extra() -> None: } ) - assert json.loads(GSheetsEngineSpec.unmask_encrypted_extra(old, new)) == { + assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == { "service_account_info": { "project_id": "yellow-unicorn-314419", "private_key": "SECRET", }, } + + +def test_unmask_encrypted_extra_when_old_is_none() -> None: + """ + Test that a None value works for ``encrypted_extra``. + """ + from superset.db_engine_specs.gsheets import GSheetsEngineSpec + + old = None + new = json.dumps( + { + "service_account_info": { + "project_id": "yellow-unicorn-314419", + "private_key": "XXXXXXXXXX", + }, + } + ) + + assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == { + "service_account_info": { + "project_id": "yellow-unicorn-314419", + "private_key": "XXXXXXXXXX", + }, + } + + +def test_unmask_encrypted_extra_when_new_is_none() -> None: + """ + Test that a None value works for ``encrypted_extra``. + """ + from superset.db_engine_specs.gsheets import GSheetsEngineSpec + + old = json.dumps( + { + "service_account_info": { + "project_id": "yellow-unicorn-314419", + "private_key": "XXXXXXXXXX", + }, + } + ) + new = None + + assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) is None