diff --git a/dataherald/sql_database/models/types.py b/dataherald/sql_database/models/types.py index 0254711a..27e4756c 100644 --- a/dataherald/sql_database/models/types.py +++ b/dataherald/sql_database/models/types.py @@ -121,12 +121,11 @@ def set_dialect(cls, input_string): def connection_uri_format(cls, value: str, values): fernet_encrypt = FernetEncrypt() try: - decrypted_connection_uri = fernet_encrypt.decrypt(value) - dialect_prefix = cls.validate_uri(decrypted_connection_uri) + fernet_encrypt.decrypt(value) except Exception: dialect_prefix = cls.validate_uri(value) + values["dialect"] = cls.set_dialect(dialect_prefix) value = fernet_encrypt.encrypt(value) - values["dialect"] = cls.set_dialect(dialect_prefix) return value @validator("llm_api_key", pre=True, always=True) diff --git a/dataherald/tests/db/test_db.py b/dataherald/tests/db/test_db.py index 4adc1189..26a32a70 100644 --- a/dataherald/tests/db/test_db.py +++ b/dataherald/tests/db/test_db.py @@ -15,7 +15,6 @@ def __init__(self, system: System): { "_id": ObjectId("64dfa0e103f5134086f7090c"), "alias": "alias", - "dialect": "dialect", "use_ssh": False, "connection_uri": "gAAAAABkwD9Y9EpBxF1hRxhovjvedX1TeDNu-WaGqDebk_CJnpGjRlpXzDOl_puehMSbz9KDQ6OqPepl8XQpD0EchiV7he4j5tEXYE33eak87iORA7s8ko0=", # noqa: E501 "ssh_settings": None, diff --git a/dataherald/tests/test_api.py b/dataherald/tests/test_api.py index 38ca5eed..5d086bc9 100644 --- a/dataherald/tests/test_api.py +++ b/dataherald/tests/test_api.py @@ -19,7 +19,6 @@ def test_scan_all_tables(): "/api/v1/table-descriptions/sync-schemas", json={"db_connection_id": "64dfa0e103f5134086f7090c"}, ) - print(response.json()) assert response.status_code == HTTP_201_CODE