From 79bc7db2c963d881b2f7bf6e39fca6d915a22006 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Sun, 16 Feb 2025 21:48:32 +0530 Subject: [PATCH 1/2] Added the enforce schema option --- src/databricks/sql/client.py | 4 ++++ src/databricks/sql/thrift_backend.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index aefed1ef0..9d24c0b42 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -737,6 +737,7 @@ def execute( self, operation: str, parameters: Optional[TParameterCollection] = None, + enforceEmbeddedSchema=False, ) -> "Cursor": """ Execute a query and wait for execution to complete. @@ -801,6 +802,7 @@ def execute( use_cloud_fetch=self.connection.use_cloud_fetch, parameters=prepared_params, async_op=False, + enforceEmbeddedSchema=enforceEmbeddedSchema, ) self.active_result_set = ResultSet( self.connection, @@ -822,6 +824,7 @@ def execute_async( self, operation: str, parameters: Optional[TParameterCollection] = None, + enforceEmbeddedSchema=False, ) -> "Cursor": """ @@ -862,6 +865,7 @@ def execute_async( use_cloud_fetch=self.connection.use_cloud_fetch, parameters=prepared_params, async_op=True, + enforceEmbeddedSchema=enforceEmbeddedSchema, ) return self diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 9beab0371..dbacee378 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -883,6 +883,7 @@ def execute_command( use_cloud_fetch=True, parameters=[], async_op=False, + enforceEmbeddedSchema=False, ): assert session_handle is not None @@ -910,6 +911,7 @@ def execute_command( }, useArrowNativeTypes=spark_arrow_types, parameters=parameters, + enforceEmbeddedSchemaCorrectness=enforceEmbeddedSchema, ) resp = self.make_request(self._client.ExecuteStatement, req) From e94c3fc3e67afa3fa12e85cc3df31ea32bb15614 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Tue, 18 Feb 2025 23:03:23 +0530 Subject: [PATCH 2/2] Renamed the variable --- src/databricks/sql/client.py | 8 ++++---- src/databricks/sql/thrift_backend.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 9d24c0b42..7e3aa02a9 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -737,7 +737,7 @@ def execute( self, operation: str, parameters: Optional[TParameterCollection] = None, - enforceEmbeddedSchema=False, + enforce_embedded_schema_correctness=False, ) -> "Cursor": """ Execute a query and wait for execution to complete. @@ -802,7 +802,7 @@ def execute( use_cloud_fetch=self.connection.use_cloud_fetch, parameters=prepared_params, async_op=False, - enforceEmbeddedSchema=enforceEmbeddedSchema, + enforce_embedded_schema_correctness=enforce_embedded_schema_correctness, ) self.active_result_set = ResultSet( self.connection, @@ -824,7 +824,7 @@ def execute_async( self, operation: str, parameters: Optional[TParameterCollection] = None, - enforceEmbeddedSchema=False, + enforce_embedded_schema_correctness=False, ) -> "Cursor": """ @@ -865,7 +865,7 @@ def execute_async( use_cloud_fetch=self.connection.use_cloud_fetch, parameters=prepared_params, async_op=True, - enforceEmbeddedSchema=enforceEmbeddedSchema, + enforce_embedded_schema_correctness=enforce_embedded_schema_correctness, ) return self diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index dbacee378..aa0f2ede4 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -883,7 +883,7 @@ def execute_command( use_cloud_fetch=True, parameters=[], async_op=False, - enforceEmbeddedSchema=False, + enforce_embedded_schema_correctness=False, ): assert session_handle is not None @@ -911,7 +911,7 @@ def execute_command( }, useArrowNativeTypes=spark_arrow_types, parameters=parameters, - enforceEmbeddedSchemaCorrectness=enforceEmbeddedSchema, + enforceEmbeddedSchemaCorrectness=enforce_embedded_schema_correctness, ) resp = self.make_request(self._client.ExecuteStatement, req)