Skip to content

Commit e6c443b

Browse files
committed
deployment changes
1 parent ddfb279 commit e6c443b

File tree

7 files changed

+166
-51
lines changed

7 files changed

+166
-51
lines changed

.github/workflows/snowpark-ci-cd.yml

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,13 +158,32 @@ jobs:
158158
echo "${{ secrets.SNOWFLAKE_PRIVATE_KEY }}" > ~/.snowflake/keys/rsa_key.p8
159159
chmod 600 ~/.snowflake/keys/rsa_key.p8
160160
161+
# Install cryptography for key handling
162+
pip install cryptography snowflake-connector-python
163+
164+
# Validate the key format
165+
python -c "
166+
from cryptography.hazmat.backends import default_backend
167+
from cryptography.hazmat.primitives import serialization
168+
try:
169+
with open('~/.snowflake/keys/rsa_key.p8', 'rb') as f:
170+
private_key = serialization.load_pem_private_key(
171+
f.read(),
172+
password=None,
173+
backend=default_backend()
174+
)
175+
print('Private key format is valid')
176+
except Exception as e:
177+
print(f'Error with private key format: {e}')
178+
# Don't exit with error yet
179+
"
180+
161181
# Create connection profile for the environment with key authentication
162182
cat > ~/.snowflake/connections.toml << EOF
163183
[dev]
164184
account = "$SNOWFLAKE_ACCOUNT"
165185
user = "$SNOWFLAKE_USER"
166186
private_key_path = "~/.snowflake/keys/rsa_key.p8"
167-
authenticator = "snowflake_jwt"
168187
warehouse = "CO2_WH_DEV"
169188
role = "CO2_ROLE_DEV"
170189
database = "CO2_DB_DEV"
@@ -174,7 +193,6 @@ jobs:
174193
account = "$SNOWFLAKE_ACCOUNT"
175194
user = "$SNOWFLAKE_USER"
176195
private_key_path = "~/.snowflake/keys/rsa_key.p8"
177-
authenticator = "snowflake_jwt"
178196
warehouse = "CO2_WH_PROD"
179197
role = "CO2_ROLE_PROD"
180198
database = "CO2_DB_PROD"

poetry.lock

Lines changed: 38 additions & 34 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pytest = "^6.2.5"
1515
jinja2 = "3.1.4"
1616
pyyaml = "^6.0.2"
1717
boto3 = "^1.37.3"
18+
cryptography = "^44.0.2"
1819

1920
[build-system]
2021
requires = ["poetry-core>=1.0.0"]

requirements.txt

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@ toml==0.10.2
99
typer==0.15.2
1010
exceptiongroup>=1.2.2
1111
iniconfig==2.0.0
12-
tomli>=2.2.1
12+
tomli>=2.2.1
13+
cryptography

rsa_key_pair_generator.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,4 @@ Note: if you dont have openssl installed in your system, try running the below c
2222
# copy the public key into the following snowflake command:
2323
ALTER USER SHUSHIL SET RSA_PUBLIC_KEY='YOUR_PUBLIC_KEY_STRING_HERE';
2424
# COPY the private key into the github secrets using the comand:
25-
(Get-Content "$HOME\.snowflake\keys\rsa_key.p8") -join "\n"
25+
(Get-Content "$HOME\.snowflake\keys\rsa_key.p8" | Select-String -NotMatch "BEGIN|END") -join ""

scripts/generate_snowflake_keys.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
import os
2+
from cryptography.hazmat.backends import default_backend
3+
from cryptography.hazmat.primitives.asymmetric import rsa
4+
from cryptography.hazmat.primitives import serialization
5+
import argparse
6+
7+
def generate_key_pair(output_dir, key_size=2048):
8+
"""Generate RSA key pair for Snowflake authentication."""
9+
# Create output directory if it doesn't exist
10+
os.makedirs(output_dir, exist_ok=True)
11+
12+
# Generate private key
13+
private_key = rsa.generate_private_key(
14+
public_exponent=65537,
15+
key_size=key_size,
16+
backend=default_backend()
17+
)
18+
19+
# Get public key
20+
public_key = private_key.public_key()
21+
22+
# Write private key in PEM PKCS8 format (no encryption)
23+
private_key_path = os.path.join(output_dir, "rsa_key.p8")
24+
with open(private_key_path, "wb") as f:
25+
f.write(
26+
private_key.private_bytes(
27+
encoding=serialization.Encoding.PEM,
28+
format=serialization.PrivateFormat.PKCS8,
29+
encryption_algorithm=serialization.NoEncryption()
30+
)
31+
)
32+
33+
# Write public key
34+
public_key_path = os.path.join(output_dir, "rsa_key.pub")
35+
with open(public_key_path, "wb") as f:
36+
f.write(
37+
public_key.public_bytes(
38+
encoding=serialization.Encoding.PEM,
39+
format=serialization.PublicFormat.SubjectPublicKeyInfo
40+
)
41+
)
42+
43+
# Create a SQL script to register the public key with Snowflake
44+
public_key_text = public_key.public_bytes(
45+
encoding=serialization.Encoding.PEM,
46+
format=serialization.PublicFormat.SubjectPublicKeyInfo
47+
).decode('utf-8').strip()
48+
49+
sql_path = os.path.join(output_dir, "register_key.sql")
50+
with open(sql_path, "w") as f:
51+
f.write(f"""-- Execute this SQL in Snowflake to register your public key
52+
ALTER USER YOUR_USERNAME SET RSA_PUBLIC_KEY='{public_key_text}';
53+
54+
-- Verify key registration
55+
DESC USER YOUR_USERNAME;
56+
""")
57+
58+
print(f"Keys generated successfully:")
59+
print(f" Private key: {private_key_path}")
60+
print(f" Public key: {public_key_path}")
61+
print(f" SQL script: {sql_path}")
62+
print("\nIMPORTANT: Update the SQL script with your username and execute it in Snowflake.")
63+
64+
return private_key_path, public_key_path, sql_path
65+
66+
if __name__ == "__main__":
67+
parser = argparse.ArgumentParser(description="Generate RSA key pair for Snowflake authentication")
68+
parser.add_argument("--output", default="~/.snowflake/keys", help="Output directory for keys")
69+
parser.add_argument("--key-size", type=int, default=2048, help="Key size in bits")
70+
args = parser.parse_args()
71+
72+
output_dir = os.path.expanduser(args.output)
73+
generate_key_pair(output_dir, args.key_size)

scripts/snowflake_deployer.py

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -66,21 +66,39 @@ def create_snowflake_connection(conn_config):
6666
if 'private_key_path' in conn_config:
6767
# Expand the path (e.g., ~ to home directory)
6868
key_path = os.path.expanduser(conn_config['private_key_path'])
69+
logger.info(f"Using private key from: {key_path}")
6970

70-
# Read the private key
71-
with open(key_path, 'rb') as key_file:
72-
private_key = key_file.read()
71+
# Read the private key properly
72+
try:
73+
from cryptography.hazmat.backends import default_backend
74+
from cryptography.hazmat.primitives import serialization
7375

74-
# Connect with private key
75-
return snowflake.connector.connect(
76-
account=conn_config.get('account'),
77-
user=conn_config.get('user'),
78-
private_key=private_key,
79-
warehouse=conn_config.get('warehouse'),
80-
database=conn_config.get('database'),
81-
schema=conn_config.get('schema'),
82-
role=conn_config.get('role')
83-
)
76+
with open(key_path, "rb") as key_file:
77+
p_key = serialization.load_pem_private_key(
78+
key_file.read(),
79+
password=None,
80+
backend=default_backend()
81+
)
82+
83+
pkb = p_key.private_bytes(
84+
encoding=serialization.Encoding.DER,
85+
format=serialization.PrivateFormat.PKCS8,
86+
encryption_algorithm=serialization.NoEncryption()
87+
)
88+
89+
# Connect with private key
90+
return snowflake.connector.connect(
91+
account=conn_config.get('account'),
92+
user=conn_config.get('user'),
93+
private_key=pkb,
94+
warehouse=conn_config.get('warehouse'),
95+
database=conn_config.get('database'),
96+
schema=conn_config.get('schema'),
97+
role=conn_config.get('role')
98+
)
99+
except Exception as e:
100+
logger.error(f"Error processing private key: {str(e)}")
101+
raise
84102
else:
85103
# Connect with password
86104
return snowflake.connector.connect(

0 commit comments

Comments
 (0)