Merge pull request #61 from NYPL/snowflake-connect

Allow connecting to snowflake with any parameters

Author: aaronfriedman6

CHANGELOG.md

@@ -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
 

pyproject.toml

@@ -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",
@@ -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]

src/nypl_py_utils/classes/snowflake_client.py

@@ -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):
         """

tests/test_snowflake_client.py

@@ -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):