Allow connecting via private key or password

Author: aaronfriedman6

src/nypl_py_utils/classes/snowflake_client.py

@@ -6,40 +6,71 @@
 class SnowflakeClient:
     """Client for managing connections to Snowflake"""
 
-    def __init__(self, account, user, password, warehouse=None):
+    def __init__(self, account, user, private_key=None, password=None):
         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.conn = None
         self.account = account
         self.user = user
+        self.private_key = private_key
         self.password = password
-        self.warehouse = warehouse
 
-    def connect(self, **kwargs):
+    def connect(self, mfa_code=None, **kwargs):
         """
-        Connects to a Snowflake database using the given credentials. If
-        warehouse parameter is None, uses the default warehouse for the user.
+        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 timeouts) can be found here:
+            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
         """
         self.logger.info('Connecting to Snowflake')
-        try:
-            self.conn = sc.connect(
-                account=self.account,
-                user=self.user,
-                password=self.password,
-                warehouse=self.warehouse,
-                **kwargs)
-        except Exception as e:
-            raise SnowflakeClientError(
-                f'Error connecting to Snowflake: {e}') from None
+        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
 
     def execute_query(self, query, **kwargs):
         """
-        Executes an arbitrary query against the given database connection.
+        Executes an arbitrary query against the given connection.
 
         Note that:
             1) All results will be fetched by default, so this method is not
@@ -53,6 +84,8 @@ def execute_query(self, query, **kwargs):
 
         Parameters
         ----------
+        query: str
+            The SQL query to execute
         kwargs:
             All possible arguments (such as timeouts) can be found here:
             https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#execute
@@ -62,7 +95,7 @@ def execute_query(self, query, **kwargs):
         sequence
             A list of tuples
         """
-        self.logger.info('Querying database')
+        self.logger.info('Querying Snowflake')
         cursor = self.conn.cursor()
         try:
             try:
@@ -73,12 +106,12 @@ def execute_query(self, query, **kwargs):
             finally:
                 cursor.close()
         except Exception as e:
-            # If there was an error, also close the database connection
+            # If there was an error, also close the connection
             self.close_connection()
 
             short_q = str(query)
             if len(short_q) > 2500:
-                short_q = short_q[:2497] + "..."
+                short_q = short_q[:2497] + '...'
             raise SnowflakeClientError(
                 f'Error executing Snowflake query {short_q}: {e}', self.logger
             ) from None

tests/test_snowflake_client.py

@@ -13,15 +13,39 @@ def mock_snowflake_conn(self, mocker):
 
     @pytest.fixture
     def test_instance(self):
-        return SnowflakeClient('test_account', 'test_user', 'test_password')
+        return SnowflakeClient(
+            'test_account', 'test_user', private_key='test_pk')
 
-    def test_connect(self, mock_snowflake_conn, test_instance):
+    def test_init_no_pw(self):
+        with pytest.raises(SnowflakeClientError):
+            SnowflakeClient('test_account', 'test_user')
+
+    def test_init_multiple_pw(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):
         test_instance.connect()
         mock_snowflake_conn.assert_called_once_with(
             account='test_account',
             user='test_user',
-            password='test_password',
-            warehouse=None)
+            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()
 
     def test_execute_query(
             self, mock_snowflake_conn, test_instance, mocker):