check

Author: jprakash-db

src/databricks/sql/client.py

@@ -223,7 +223,8 @@ def read(self) -> Optional[OAuthToken]:
         self.port = kwargs.get("_port", 443)
         self.disable_pandas = kwargs.get("_disable_pandas", False)
         self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
-
+        self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
+        
         auth_provider = get_python_sql_connector_auth_provider(
             server_hostname, **kwargs
         )
@@ -1254,6 +1255,7 @@ def __init__(
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._next_row_index = 0
         self._use_cloud_fetch = use_cloud_fetch
+        self.is_staging_operation = execute_response.is_staging_operation
 
         if execute_response.arrow_queue:
             # In this case the server has taken the fast path and returned an initial batch of

src/databricks/sql/volume/volume_client.py

@@ -15,14 +15,23 @@ def __init__(self, conn):
         :param conn: Connection object to Databricks.
         """
         self.conn = conn
+    
+    def is_staging_operation_allowed(self, condition) -> bool:
+        if not condition:
+            raise ValueError("Staging operation is not allowed")
 
     def get_object(
         self, catalog: str, schema: str, volume: str, object_path: str, local_path: str
     ) -> bool:
         get_object_query = StringUtil.create_get_object_query(
             catalog, schema, volume, object_path, local_path
         )
-        return True
+
+        with self.conn.cursor() as cursor:
+            cursor.execute(get_object_query)
+            self.is_staging_operation_allowed(cursor.active_result_set.is_staging_operation)
+            volume_processor = VolumeProcessor(local_path=local_path)
+            return True
 
     def get_object(
         self, catalog: str, schema: str, volume: str, object_path: str

src/databricks/sql/volume/volume_processor.py

@@ -0,0 +1,12 @@
+from databricks.sql import Connection
+
+class VolumeProcessor:
+    def __init__(self, **kwargs):
+        self.connection = kwargs.get("connection")
+        self.local_path = kwargs.get("local_path")
+        self.input_stream = kwargs.get("input_stream")
+        self.content_length = kwargs.get("content_length")
+        self.abs_staging_allowed_local_paths = self.get_abs_staging_allowed_local_paths(kwargs.get("staging_allowed_local_path"))
+        
+    def process_volume(self, catalog: str, schema: str, volume: str):
+        pass