Created
May 24, 2024 16:15
-
-
Save orellabac/eeadded9bcc01fa9df4a1299f10c10ab to your computer and use it in GitHub Desktop.
Using your snowpark session connection with sqlalchemy
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from sqlalchemy import create_engine | |
import sqlalchemy | |
import logging | |
def get_engine_from_session(session): | |
# Your existing Snowflake connection (replace with your actual connection) | |
existing_snowflake_connection = session._conn._conn | |
setattr(existing_snowflake_connection,"_interpolate_empty_sequences",False) | |
# sql alchemy needs pyformat binding | |
existing_snowflake_connection._paramstyle = "pyformat" | |
opts = "" | |
if session.get_current_warehouse() is not None: | |
opts += f"&warehouse={session.get_current_warehouse()}" | |
if session.get_current_role() is not None: | |
opts += f"&role={session.get_current_role()}" | |
conn_url = f"snowflake://{session.get_current_user() or ''}@{session.get_current_account()}/{session.get_current_database() or ''}/{session.get_current_schema() or ''}?{opts}" | |
# Create an engine and bind it to the existing Snowflake connection | |
engine = create_engine( | |
url=conn_url, | |
creator=lambda: existing_snowflake_connection | |
) | |
return engine |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment