Refactor to allow uploads_base_path to be either a single string object

Jesse Whitehouse · Jesse Whitehouse · commit bdb948ad4e5f · 2022-12-28T11:33:05.000-06:00
or a list of strings.

Signed-off-by: Jesse Whitehouse &lt;jesse.whitehouse@databricks.com&gt;
diff --git a/examples/staging_ingestion.py b/examples/staging_ingestion.py
@@ -19,7 +19,7 @@
     REMOVE 'stage://tmp/some.user@databricks.com/salesdata/september.csv'
 
 Ingestion queries are passed to cursor.execute() like any other query. For GET and PUT commands, a local file
-will be read or written. For security, this local file must be contained within, or descended from, the
+will be read or written. For security, this local file must be contained within, or descended from, an
 uploads_base_path of the connection.
 
 Additionally, the connection can only manipulate files within the cloud storage location of the authenticated user.
@@ -29,6 +29,8 @@
 1. Set the INGESTION_USER constant to the account email address of the authenticated user
 2. Set the FILEPATH constant to the path of a file that will be uploaded (this example assumes its a CSV file)
 3. Run this file
+
+Note: uploads_base_path can be either a Pathlike object or a list of Pathlike objects.
 """
 
 INGESTION_USER = "some.user@example.com"
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
@@ -300,29 +300,47 @@ def _check_not_closed(self):
         if not self.open:
             raise Error("Attempting operation on closed cursor")
 
-    def _handle_staging_operation(self, uploads_base_path: str):
+    def _handle_staging_operation(self, uploads_base_path: Union[List, str]):
         """Fetch the HTTP request instruction from a staging ingestion command
         and call the designated handler.
-        
+
         Raise an exception if localFile is specified by the server but the localFile
         is not descended from uploads_base_path.
         """
 
-        if uploads_base_path is None:
+        if isinstance(uploads_base_path, type(str())):
+            _uploads_base_paths = [uploads_base_path]
+        elif isinstance(uploads_base_path, type(list())):
+            _uploads_base_paths = uploads_base_path
+        else:
             raise Error(
-                "You must provide an uploads_base_path when initialising a connection to perform ingestion commands"
+                "You must provide at least one uploads_base_path when initialising a connection to perform ingestion commands"
             )
 
+        abs_uploads_base_paths = [os.path.abspath(i) for i in _uploads_base_paths]
+
         row = self.active_result_set.fetchone()
 
         # Must set to None in cases where server response does not include localFile
         abs_localFile = None
 
+        # Default to not allow staging operations
+        allow_operation = False
         if getattr(row, "localFile", None):
             abs_localFile = os.path.abspath(row.localFile)
-            abs_uploads_base_path = os.path.abspath(uploads_base_path)
-            if os.path.commonpath([abs_localFile, abs_uploads_base_path]) != abs_uploads_base_path:
-                raise Error("Local file operations are restricted to paths within the configured uploads_base_path")
+            for abs_uploads_base_path in abs_uploads_base_paths:
+                # If the indicated local file matches at least one allowed base path, allow the operation
+                if (
+                    os.path.commonpath([abs_localFile, abs_uploads_base_path])
+                    == abs_uploads_base_path
+                ):
+                    allow_operation = True
+                else:
+                    continue
+            if not allow_operation:
+                raise Error(
+                    "Local file operations are restricted to paths within the configured uploads_base_path"
+                )
 
         # TODO: Experiment with DBR sending real headers.
         # The specification says headers will be in JSON format but the current null value is actually an empty list []
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
@@ -6,6 +6,7 @@
 import threading
 import lz4.frame
 from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
+from typing import List, Union
 
 import pyarrow
 import thrift.transport.THttpClient
@@ -61,7 +62,7 @@ def __init__(
         http_path: str,
         http_headers,
         auth_provider: AuthProvider,
-        uploads_base_path: str = None,
+        uploads_base_path: Union[str, List] = None,
         **kwargs,
     ):
         # Internal arguments in **kwargs:
diff --git a/tests/e2e/driver_tests.py b/tests/e2e/driver_tests.py
@@ -864,6 +864,45 @@ def test_staging_ingestion_invalid_staging_path_fails_at_server(self):
                 query = f"PUT '{target_file}' INTO 'stageRANDOMSTRINGOFCHARACTERS://tmp/{self.staging_ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
                 cursor.execute(query)
 
+    def test_staging_ingestion_supports_multiple_uploadsbasepath_values(self):
+        """uploads_base_path may be either a path-like object or a list of path-like objects.
+
+        This test confirms that two configured base paths:
+        1 - doesn't raise an exception
+        2 - allows uploads from both paths
+        3 - doesn't allow uploads from a third path
+        """
+
+        def generate_file_and_path_and_queries():
+            """
+                1. Makes a temp file with some contents.
+                2. Write a query to PUT it into a staging location
+                3. Write a query to REMOVE it from that location (for cleanup)
+            """
+            fh, temp_path = tempfile.mkstemp()
+            with open(fh, "wb") as fp:
+                original_text = "hello world!".encode("utf-8")
+                fp.write(original_text)
+            put_query = f"PUT '{temp_path}' INTO 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/{id(temp_path)}.csv' OVERWRITE"
+            remove_query = f"REMOVE 'stage://tmp/{self.staging_ingestion_user}/tmp/11/15/{id(temp_path)}.csv'"
+            return fh, temp_path, put_query, remove_query
+
+        fh1, temp_path1, put_query1, remove_query1 = generate_file_and_path_and_queries()
+        fh2, temp_path2, put_query2, remove_query2 = generate_file_and_path_and_queries()
+        fh3, temp_path3, put_query3, remove_query3 = generate_file_and_path_and_queries()
+
+        with self.connection(extra_params={"uploads_base_path": [temp_path1, temp_path2]}) as conn:
+            cursor = conn.cursor()
+
+            cursor.execute(put_query1)
+            cursor.execute(put_query2)
+            
+            with pytest.raises(Error, match="Local file operations are restricted to paths within the configured uploads_base_path"):
+                cursor.execute(put_query3)
+
+            # Then clean up the files we made
+            cursor.execute(remove_query1)
+            cursor.execute(remove_query2)
 
 
 def main(cli_args):
