diff --git a/CHANGELOG.md b/CHANGELOG.md index a6f3bb215..9f9d798b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Deprecate vanilla `DataType` - Remove `_Encodable` from project +- Connect to Snowflake using the incluster oauth token #### New Features & Functionality diff --git a/plugins/ibis/superduper_ibis/data_backend.py b/plugins/ibis/superduper_ibis/data_backend.py index 9a8524e39..26d9b0ffb 100644 --- a/plugins/ibis/superduper_ibis/data_backend.py +++ b/plugins/ibis/superduper_ibis/data_backend.py @@ -36,16 +36,17 @@ def _snowflake_connection_callback(): logging.info('Using env variables and OAuth to connect!') - import snowflake + import snowflake.connector conn = snowflake.connector.connect( - account=os.environ['SNOWFLAKE_ACCOUNT'], host=os.environ['SNOWFLAKE_HOST'], - schema=os.environ['SUPERDUPER_DATA_SCHEMA'], - database=os.environ['SNOWFLAKE_DATABASE'], port=int(os.environ['SNOWFLAKE_PORT']), - token=open('/snowflake/session/token').read(), + account=os.environ['SNOWFLAKE_ACCOUNT'], authenticator='oauth', + token=open('/snowflake/session/token').read(), + warehouse=os.environ['SNOWFLAKE_WAREHOUSE'], + database=os.environ['SNOWFLAKE_DATABASE'], + schema=os.environ['SUPERDUPER_DATA_SCHEMA'], ) return ibis.snowflake.from_connection(conn) @@ -68,7 +69,7 @@ def _connection_callback(uri, flavour): in_memory = True return ibis_conn, dir_name, in_memory elif uri == 'snowflake://': - return _snowflake_connection_callback(uri), 'snowflake', False + return _snowflake_connection_callback(), 'snowflake', False else: name = uri.split("//")[0] in_memory = False diff --git a/plugins/sqlalchemy/superduper_sqlalchemy/metadata.py b/plugins/sqlalchemy/superduper_sqlalchemy/metadata.py index 54f88e8f6..50172190f 100644 --- a/plugins/sqlalchemy/superduper_sqlalchemy/metadata.py +++ b/plugins/sqlalchemy/superduper_sqlalchemy/metadata.py @@ -30,19 +30,20 @@ def _connect_snowflake(): # mounted token. In this case, as a convention # we connect with `"snowflake://"` - import snowflake + import snowflake.connector def creator(): import os return snowflake.connector.connect( - account=os.environ['SNOWFLAKE_ACCOUNT'], host=os.environ['SNOWFLAKE_HOST'], - schema=os.environ['SUPERDUPER_DATA_SCHEMA'], - database=os.environ['SNOWFLAKE_DATABASE'], port=int(os.environ['SNOWFLAKE_PORT']), - token=open('/snowflake/session/token').read(), + account=os.environ['SNOWFLAKE_ACCOUNT'], authenticator='oauth', + token=open('/snowflake/session/token').read(), + warehouse=os.environ['SNOWFLAKE_WAREHOUSE'], + database=os.environ['SNOWFLAKE_DATABASE'], + schema=os.environ['SUPERDUPER_DATA_SCHEMA'], ) return create_engine("snowflake://not@used/db", creator=creator) diff --git a/templates/pdf_rag/component.json b/templates/pdf_rag/component.json index 7ef010a0a..163ad3cb6 100644 --- a/templates/pdf_rag/component.json +++ b/templates/pdf_rag/component.json @@ -335,7 +335,7 @@ "blobs": null, "files": null, "requirements": null, - "default_table": "?table:sample_pdf_rag", + "default_tables": ["?table:sample_pdf_rag"], "queries": null, "_literals": [ "template" @@ -392,4 +392,4 @@ } }, "_files": {} -} \ No newline at end of file +}