Fix oauth in snowflake ibis plugin

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
 

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

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)

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": {}
-}
+}