Envvars for Snowflake adapter

splitgraph/ingestion/snowflake/__init__.py

@@ -1,3 +1,4 @@
+import json
 import urllib.parse
 from typing import Dict, Optional, cast, Mapping
 
@@ -32,6 +33,10 @@ class SnowflakeDataSource(ForeignDataWrapperDataSource):
             "schema": {"type": "string", "description": "Snowflake schema"},
             "warehouse": {"type": "string", "description": "Warehouse name"},
             "role": {"type": "string", "description": "Role"},
+            "envvars": {
+                "type": "object",
+                "description": "Environment variables to set on the engine side",
+            },
         },
         "required": ["database"],
     }
@@ -53,6 +58,7 @@ class SnowflakeDataSource(ForeignDataWrapperDataSource):
     "account": "acc-id.west-europe.azure",
     "database": "SNOWFLAKE_SAMPLE_DATA",
     "schema": "TPCH_SF100"
+    "envvars": {"HTTP_PROXY": "http://proxy.company.com"}
 }
 EOF
 
@@ -130,10 +136,14 @@ def get_server_options(self):
         if "role" in self.params:
             extra_params["role"] = self.params["role"]
 
-        db_url += urllib.parse.urlencode(extra_params)
+        if extra_params:
+            db_url += "?" + urllib.parse.urlencode(extra_params)
 
         options["db_url"] = db_url
 
+        if "envvars" in self.params:
+            options["envvars"] = json.dumps(self.params["envvars"])
+
         return options
 
     def get_remote_schema_name(self) -> str:

test/splitgraph/ingestion/test_snowflake.py

@@ -3,7 +3,7 @@
 from splitgraph.ingestion.snowflake import SnowflakeDataSource
 
 
-def test_snowflake_data_source_dburl_conversion():
+def test_snowflake_data_source_dburl_conversion_warehouse():
     source = SnowflakeDataSource(
         Mock(),
         credentials={
@@ -20,11 +20,31 @@ def test_snowflake_data_source_dburl_conversion():
     )
 
     assert source.get_server_options() == {
-        "db_url": "snowflake://username:password@abcdef.eu-west-1.aws/SOME_DB/TPCH_SF100warehouse=my_warehouse&role=role",
+        "db_url": "snowflake://username:password@abcdef.eu-west-1.aws/SOME_DB/TPCH_SF100?warehouse=my_warehouse&role=role",
         "schema": "TPCH_SF100",
         "wrapper": "multicorn.sqlalchemyfdw.SqlAlchemyFdw",
     }
 
+
+def test_snowflake_data_source_dburl_conversion_no_warehouse():
+    source = SnowflakeDataSource(
+        Mock(),
+        credentials={
+            "username": "username",
+            "password": "password",
+            "account": "abcdef.eu-west-1.aws",
+        },
+        params={"database": "SOME_DB", "schema": "TPCH_SF100",},
+    )
+
+    assert source.get_server_options() == {
+        "db_url": "snowflake://username:password@abcdef.eu-west-1.aws/SOME_DB/TPCH_SF100",
+        "schema": "TPCH_SF100",
+        "wrapper": "multicorn.sqlalchemyfdw.SqlAlchemyFdw",
+    }
+
+
+def test_snowflake_data_source_table_options():
     source = SnowflakeDataSource(
         Mock(),
         credentials={