diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index d3dd2664a..ef25a0534 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -285,9 +285,6 @@ class SnowflakeConnectionInfo(BaseConnectionInfo): user: SecretStr = Field( description="the username of your database", examples=["admin"] ) - password: SecretStr = Field( - description="the password of your database", examples=["password"] - ) account: SecretStr = Field( description="the account name of your database", examples=["myaccount"] ) @@ -299,6 +296,14 @@ class SnowflakeConnectionInfo(BaseConnectionInfo): description="the schema name of your database", examples=["myschema"], ) # Use `sf_schema` to avoid `schema` shadowing in BaseModel + warehouse: SecretStr = Field( + description="the warehouse name of your database", examples=["COMPUTE_WH"] + ) + private_key: SecretStr | None = Field( + description="the private key for key pair authentication", + examples=["private_key_content"], + default=None, + ) kwargs: dict[str, str] | None = Field( description="Additional arguments passed to the DBAPI connection call.", default=None, diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index a575df7fe..187acc532 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -220,10 +220,11 @@ def get_oracle_connection(info: OracleConnectionInfo) -> BaseBackend: def get_snowflake_connection(info: SnowflakeConnectionInfo) -> BaseBackend: return ibis.snowflake.connect( user=info.user.get_secret_value(), - password=info.password.get_secret_value(), account=info.account.get_secret_value(), database=info.database.get_secret_value(), schema=info.sf_schema.get_secret_value(), + warehouse=info.warehouse.get_secret_value(), + private_key=info.private_key.get_secret_value(), **info.kwargs if info.kwargs else dict(), ) diff --git a/ibis-server/tests/routers/v2/connector/test_snowflake.py b/ibis-server/tests/routers/v2/connector/test_snowflake.py index ee374d7bb..7abcab694 100644 --- a/ibis-server/tests/routers/v2/connector/test_snowflake.py +++ b/ibis-server/tests/routers/v2/connector/test_snowflake.py @@ -12,10 +12,11 @@ connection_info = { "user": os.getenv("SNOWFLAKE_USER"), - "password": os.getenv("SNOWFLAKE_PASSWORD"), "account": os.getenv("SNOWFLAKE_ACCOUNT"), "database": "SNOWFLAKE_SAMPLE_DATA", "schema": "TPCH_SF1", + "warehouse": "COMPUTE_WH", + "private_key": os.getenv("SNOWFLAKE_PRIVATE_KEY"), } manifest = { diff --git a/ibis-server/tests/routers/v3/connector/snowflake/conftest.py b/ibis-server/tests/routers/v3/connector/snowflake/conftest.py index d5a4afcad..ef11b761d 100644 --- a/ibis-server/tests/routers/v3/connector/snowflake/conftest.py +++ b/ibis-server/tests/routers/v3/connector/snowflake/conftest.py @@ -19,8 +19,9 @@ def pytest_collection_modifyitems(items): def connection_info() -> dict[str, str]: return { "user": os.getenv("SNOWFLAKE_USER"), - "password": os.getenv("SNOWFLAKE_PASSWORD"), "account": os.getenv("SNOWFLAKE_ACCOUNT"), "database": "SNOWFLAKE_SAMPLE_DATA", "schema": "TPCH_SF1", + "warehouse": "COMPUTE_WH", + "private_key": os.getenv("SNOWFLAKE_PRIVATE_KEY"), }