pip install snowflake-connector-python cryptography

export SNOWFLAKE_ACCOUNT="update"
export SNOWFLAKE_USER="update"
export SNOWFLAKE_PRIVATE_KEY_PATH="/path/to/your/private/key/rsa_key.p8"
export SNOWFLAKE_PRIVATE_KEY_PASSPHRASE="update"
export SNOWFLAKE_WAREHOUSE="update"
export SNOWFLAKE_DATABASE="update"
export SNOWFLAKE_SCHEMA="update"

import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
import os

# --- Configuration ---
# Carta's Snowflake Account
SNOWFLAKE_ACCOUNT = os.getenv("SNOWFLAKE_ACCOUNT")
# Snowflake Service Account
SNOWFLAKE_USER = os.getenv("SNOWFLAKE_USER")
# Path to your private key file (e.g., '/path/to/rsa_key.p8')
PRIVATE_KEY_PATH = os.getenv("SNOWFLAKE_PRIVATE_KEY_PATH")
# Private key passphrase sent in 1Password entry
PRIVATE_KEY_PASSPHRASE = os.getenv("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")
# Default Snowflake warehouse, database, and schema to connect to
SNOWFLAKE_WAREHOUSE = os.getenv("SNOWFLAKE_WAREHOUSE")
SNOWFLAKE_DATABASE = os.getenv("SNOWFLAKE_DATABASE")
SNOWFLAKE_SCHEMA = os.getenv("SNOWFLAKE_SCHEMA")
# Optional: Snowflake role to use (if not set, uses user's default role)
SNOWFLAKE_ROLE = os.getenv("SNOWFLAKE_ROLE")

def connect_to_snowflake_keypair():
    """
    Connects to Snowflake using key pair authentication with a passphrase.
    """
    conn = None
    try:
        # 1. Load the private key
        if not os.path.exists(PRIVATE_KEY_PATH):
            print(f"Error: Private key file not found at '{PRIVATE_KEY_PATH}'")
            return None

        with open(PRIVATE_KEY_PATH, "rb") as key:
            p_key = serialization.load_pem_private_key(
                key.read(),
                password=(
                    PRIVATE_KEY_PASSPHRASE.encode() if PRIVATE_KEY_PASSPHRASE else None
                ),
                backend=default_backend(),
            )

        # Convert the private key to DER format for the Snowflake connector
        pkb = p_key.private_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )

        # 2. Establish the Snowflake connection
        print(f"Attempting to connect to Snowflake as user '{SNOWFLAKE_USER}'...")
        connect_args = {
            "user": SNOWFLAKE_USER,
            "account": SNOWFLAKE_ACCOUNT,
            "private_key": pkb,
            "warehouse": SNOWFLAKE_WAREHOUSE,
            "database": SNOWFLAKE_DATABASE,
            "schema": SNOWFLAKE_SCHEMA,
        }
        if SNOWFLAKE_ROLE:
            connect_args["role"] = SNOWFLAKE_ROLE
        
        conn = snowflake.connector.connect(**connect_args)
        print("Successfully connected to Snowflake!")

        # 3. Execute a simple query to verify the connection
        cursor = conn.cursor()
        try:
            cursor.execute("SELECT current_version()")
            one_row = cursor.fetchone()
            print(f"Snowflake Version: {one_row[0]}")

            cursor.execute("SELECT current_user()")
            current_user = cursor.fetchone()
            print(f"Connected as user: {current_user[0]}")

            # Show all tables in the FUND_ADMIN schema
            cursor.execute(f"SHOW TABLES IN SCHEMA {SNOWFLAKE_DATABASE}.FUND_ADMIN")
            tables = cursor.fetchall()
            print(f"\nTables in {SNOWFLAKE_DATABASE}.FUND_ADMIN schema:")
            for table in tables:
                cursor.execute(f"SELECT count(*) from FUND_ADMIN.{table[1]}")
                results = cursor.fetchone()
                print(f"  - {table[1]} ({results[0]} rows)")
        finally:
            if cursor:
                cursor.close()

        return conn

    except Exception as e:
        print(f"An error occurred during Snowflake connection: {e}")
        return None

# --- Main execution ---
if __name__ == "__main__":
    snowflake_connection = connect_to_snowflake_keypair()

    if snowflake_connection:
        print("\nConnection successful! You can now perform database operations.")
        # Query distinct fund names from AGGREGATE_INVESTMENTS
        cursor = snowflake_connection.cursor()
        try:
            print("\nQuerying distinct fund names from FUND_ADMIN.AGGREGATE_INVESTMENTS...")
            cursor.execute("SELECT DISTINCT fund_name FROM FUND_ADMIN.AGGREGATE_INVESTMENTS")
            fund_names = cursor.fetchall()

            print(f"\nFound {len(fund_names)} distinct fund names:")
            for idx, (fund_name,) in enumerate(fund_names, 1):
                print(f"  {idx}. {fund_name}")
        finally:
            cursor.close()
            print("\nClosing Snowflake connection.")
            snowflake_connection.close()
    else:
        print("\nFailed to establish Snowflake connection.")

python carta_snowflake.py