Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
# Changelog
## v1.11.0 4/27/26
- Allow Snowflake client to connect using any parameters

## v1.10.2 4/27/26
- Prevent double logging in pytest using structlog

Expand Down
66 changes: 33 additions & 33 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@ build-backend = "hatchling.build"

[project]
name = "nypl_py_utils"
version = "1.10.2"
version = "1.11.0"
authors = [
{ name="Aaron Friedman", email="aaronfriedman@nypl.org" },
]
description = "A package containing Python utilities for use across NYPL"
readme = "README.md"
requires-python = ">=3.9"
requires-python = ">=3.10"
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: MIT License",
Expand All @@ -25,85 +25,85 @@ dependencies = []
[project.optional-dependencies]
avro-client = [
"nypl_py_utils[log-helper]",
"fastavro>=1.11.1",
"requests>=2.28.1"
"fastavro==1.12.2",
"requests==2.33.1"
]
cloudlibrary-client = [
"nypl_py_utils[log-helper]",
"requests>=2.28.1"
"requests==2.33.1"
]
kinesis-client = [
"nypl_py_utils[log-helper]",
"boto3>=1.26.5",
"botocore>=1.29.5"
"boto3==1.43.1",
"botocore==1.43.1"
]
kms-client = [
"nypl_py_utils[log-helper]",
"boto3>=1.26.5",
"botocore>=1.29.5"
"boto3==1.43.1",
"botocore==1.43.1"
]
mysql-client = [
"nypl_py_utils[log-helper]",
"mysql-connector-python>=8.0.32"
"mysql-connector-python==9.7.0"
]
oauth2-api-client = [
"nypl_py_utils[log-helper]",
"oauthlib>=3.2.2",
"requests_oauthlib>=1.3.1"
"oauthlib==3.3.1",
"requests_oauthlib==2.0.0"
]
postgresql-client = [
"nypl_py_utils[log-helper]",
"psycopg[binary]>=3.1.6"
"psycopg[binary]==3.3.3"
]
redshift-client = [
"nypl_py_utils[log-helper]",
"botocore>=1.29.5",
"redshift-connector>=2.0.909"
"botocore==1.43.1",
"redshift-connector==2.1.13"
]
s3-client = [
"nypl_py_utils[log-helper]",
"boto3>=1.26.5",
"botocore>=1.29.5"
"boto3==1.43.1",
"botocore==1.43.1"
]
secrets-manager-client = [
"nypl_py_utils[log-helper]",
"boto3>=1.26.5",
"botocore>=1.29.5"
"boto3==1.43.1",
"botocore==1.43.1"
]
sftp-client = [
"nypl_py_utils[log-helper]",
"paramiko>=3.4.1"
"paramiko==4.0.0"
]
snowflake-client = [
"nypl_py_utils[log-helper]",
"snowflake-connector-python>=4.3.0"
"snowflake-connector-python==4.3.0"
]
config-helper = [
"nypl_py_utils[kms-client,log-helper]",
"PyYAML>=6.0"
"PyYAML==6.0.3"
]
log-helper = [
"structlog>=25.5.0"
"structlog==25.5.0"
]
obfuscation-helper = [
"nypl_py_utils[log-helper]",
"bcrypt>=4.0.1"
"bcrypt==5.0.0"
]
patron-data-helper = [
"nypl_py_utils[postgresql-client,redshift-client,log-helper]>=1.1.5",
"pandas>=2.2.2"
"nypl_py_utils[postgresql-client,redshift-client,log-helper]",
"pandas==3.0.2"
]
research-catalog-identifier-helper = [
"requests>=2.28.1"
"requests==2.33.1"
]
development = [
"nypl_py_utils[avro-client,cloudlibrary-client,kinesis-client,kms-client,mysql-client,oauth2-api-client,postgresql-client,redshift-client,s3-client,secrets-manager-client,sftp-client,snowflake-client,config-helper,log-helper,obfuscation-helper,patron-data-helper,research-catalog-identifier-helper]",
"flake8>=6.0.0",
"freezegun>=1.2.2",
"mock>=4.0.3",
"pytest>=8.0.0",
"pytest-mock>=3.10.0",
"requests-mock>=1.10.0"
"flake8==7.3.0",
"freezegun==1.5.5",
"mock==5.2.0",
"pytest==9.0.3",
"pytest-mock==3.15.1",
"requests-mock==1.12.1"
]

[tool.pytest.ini_options]
Expand Down
77 changes: 22 additions & 55 deletions src/nypl_py_utils/classes/snowflake_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,67 +6,34 @@
class SnowflakeClient:
"""Client for managing connections to Snowflake"""

def __init__(self, account, user, private_key=None, password=None):
def __init__(self, connection_params):
"""
See the `connect` method below for what `connection_params` should
look like
"""
self.logger = create_log('snowflake_client')
if (password is None) == (private_key is None):
raise SnowflakeClientError(
'Either password or private key must be set (but not both)',
self.logger
) from None

self.connection_params = connection_params
self.conn = None
self.account = account
self.user = user
self.private_key = private_key
self.password = password

def connect(self, mfa_code=None, **kwargs):
def connect(self):
"""
Connects to Snowflake using the given credentials. If you're connecting
locally, you should be using the password and mfa_code. If the
connection is for production code, a private_key should be set up.

Parameters
----------
mfa_code: str, optional
The six-digit MFA code. Only necessary for connecting as a human
user.
kwargs:
All possible arguments (such as which warehouse to use or how
long to wait before timing out) can be found here:
https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#connect
Connects to Snowflake using the given connection parameters. In
practice, there are likely two sets of parameters that will be used:
1. Local development: `connection_name` and `private_key_file_pwd`.
This method requires a `connections.toml` file with a matching
`connection_name` connection.
2. Production code: `account`, `user`, and `private_key`.

All possible parameters can be found here:
https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#connect
"""
self.logger.info('Connecting to Snowflake')
if self.private_key is not None:
try:
self.conn = sc.connect(
account=self.account,
user=self.user,
private_key=self.private_key,
**kwargs)
except Exception as e:
raise SnowflakeClientError(
f'Error connecting to Snowflake: {e}', self.logger
) from None
else:
if mfa_code is None:
raise SnowflakeClientError(
'When using a password, an MFA code must also be provided',
self.logger
) from None

pw = self.password + mfa_code
try:
self.conn = sc.connect(
account=self.account,
user=self.user,
password=pw,
passcode_in_password=True,
**kwargs)
except Exception as e:
raise SnowflakeClientError(
f'Error connecting to Snowflake: {e}', self.logger
) from None
try:
self.conn = sc.connect(**self.connection_params)
except Exception as e:
raise SnowflakeClientError(
f'Error connecting to Snowflake: {e}', self.logger
) from None

def execute_query(self, query, **kwargs):
"""
Expand Down
33 changes: 4 additions & 29 deletions tests/test_snowflake_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,38 +14,13 @@ def mock_snowflake_conn(self, mocker):
@pytest.fixture
def test_instance(self):
return SnowflakeClient(
'test_account', 'test_user', private_key='test_pk')
{'account': 'test_account', 'user': 'test_user'}
)

def test_init_no_pw(self):
with pytest.raises(SnowflakeClientError):
SnowflakeClient('test_account', 'test_user')

def test_init_multiple_auth(self):
with pytest.raises(SnowflakeClientError):
SnowflakeClient('test_account', 'test_user', 'test_pk', 'test_pw')

def test_connect_with_pk(self, mock_snowflake_conn, test_instance):
def test_connect(self, mock_snowflake_conn, test_instance):
test_instance.connect()
mock_snowflake_conn.assert_called_once_with(
account='test_account',
user='test_user',
private_key='test_pk')

def test_connect_with_pw(self, mock_snowflake_conn):
test_instance = SnowflakeClient(
'test_account', 'test_user', password='test_pw')
test_instance.connect('123456')
mock_snowflake_conn.assert_called_once_with(
account='test_account',
user='test_user',
password='test_pw123456',
passcode_in_password=True)

def test_connect_no_mfa(self, mock_snowflake_conn):
test_instance = SnowflakeClient(
'test_account', 'test_user', password='test_pw')
with pytest.raises(SnowflakeClientError):
test_instance.connect()
account='test_account', user='test_user')

def test_execute_query(
self, mock_snowflake_conn, test_instance, mocker):
Expand Down
Loading