Merge branch 'main' into sqlalchemy-tz

Jesse Whitehouse · Jesse Whitehouse · commit 40e23d24dd43 · 2023-10-31T20:50:11.000-04:00
Signed-off-by: Jesse Whitehouse &lt;jesse.whitehouse@databricks.com&gt;
diff --git a/examples/README.md b/examples/README.md
@@ -33,6 +33,7 @@ To run all of these examples you can clone the entire repository to your disk. O
 - **`insert_data.py`** adds a tables called `squares` to your default catalog and inserts one hundred rows of example data. Then it fetches this data and prints it to the screen.
 - **`query_cancel.py`** shows how to cancel a query assuming that you can access the `Cursor` executing that query from a different thread. This is necessary because `databricks-sql-connector` does not yet implement an asynchronous API; calling `.execute()` blocks the current thread until execution completes. Therefore, the connector can't cancel queries from the same thread where they began.
 - **`interactive_oauth.py`** shows the simplest example of authenticating by OAuth (no need for a PAT generated in the DBSQL UI) while Bring Your Own IDP is in public preview. When you run the script it will open a browser window so you can authenticate. Afterward, the script fetches some sample data from Databricks and prints it to the screen. For this script, the OAuth token is not persisted which means you need to authenticate every time you run the script.
+- **`m2m_oauth.py`** shows the simplest example of authenticating by using OAuth M2M (machine-to-machine) for service principal.
 - **`persistent_oauth.py`** shows a more advanced example of authenticating by OAuth while Bring Your Own IDP is in public preview. In this case, it shows how to use a sublcass of `OAuthPersistence` to reuse an OAuth token across script executions.
 - **`set_user_agent.py`** shows how to customize the user agent header used for Thrift commands. In
 this example the string `ExamplePartnerTag` will be added to the the user agent on every request.
diff --git a/examples/m2m_oauth.py b/examples/m2m_oauth.py
@@ -0,0 +1,41 @@
+import os
+
+from databricks.sdk.core import oauth_service_principal, Config
+from databricks import sql
+
+"""
+This example shows how to use OAuth M2M (machine-to-machine) for service principal
+
+Pre-requisites:
+- Create service principal and OAuth secret in Account Console
+- Assign the service principal to the workspace
+
+See more https://docs.databricks.com/en/dev-tools/authentication-oauth.html)
+"""
+
+server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")
+
+
+def credential_provider():
+    config = Config(
+        host=f"https://{server_hostname}",
+        # Service Principal UUID
+        client_id=os.getenv("DATABRICKS_CLIENT_ID"),
+        # Service Principal Secret
+        client_secret=os.getenv("DATABRICKS_CLIENT_SECRET"))
+    return oauth_service_principal(config)
+
+
+with sql.connect(
+        server_hostname=server_hostname,
+        http_path=os.getenv("DATABRICKS_HTTP_PATH"),
+        credentials_provider=credential_provider) as connection:
+    for x in range(1, 100):
+        cursor = connection.cursor()
+        cursor.execute('SELECT 1+1')
+        result = cursor.fetchall()
+        for row in result:
+            print(row)
+        cursor.close()
+
+    connection.close()
diff --git a/src/databricks/sql/__init__.py b/src/databricks/sql/__init__.py
@@ -5,7 +5,10 @@
 # PEP 249 module globals
 apilevel = "2.0"
 threadsafety = 1  # Threads may share the module, but not connections.
-paramstyle = "pyformat"  # Python extended format codes, e.g. ...WHERE name=%(name)s
+
+# Python extended format codes, e.g. ...WHERE name=%(name)s
+# Note that when we switch to ParameterApproach.NATIVE, paramstyle will be `named`
+paramstyle = "pyformat"
 
 
 class DBAPITypeObject(object):
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
@@ -36,7 +36,7 @@
 DEFAULT_RESULT_BUFFER_SIZE_BYTES = 104857600
 DEFAULT_ARRAY_SIZE = 100000
 
-NO_NATIVE_PARAMS = []
+NO_NATIVE_PARAMS: List = []
 
 
 class Connection:
@@ -376,9 +376,12 @@ def __iter__(self):
         else:
             raise Error("There is no active result set")
 
-    def _determine_parameter_approach(self) -> ParameterApproach:
+    def _determine_parameter_approach(
+        self, params: Optional[Union[List, Dict[str, Any]]] = None
+    ) -> ParameterApproach:
         """Encapsulates the logic for choosing whether to send parameters in native vs inline mode
 
+        If params is None then ParameterApproach.NONE is returned.
         If self.use_inline_params is True then inline mode is used.
         If self.use_inline_params is False, then check if the server supports them and proceed.
             Else raise an exception.
@@ -388,6 +391,9 @@ def _determine_parameter_approach(self) -> ParameterApproach:
         If inline approach is used when the server supports native approach, a warning is logged
         """
 
+        if params is None:
+            return ParameterApproach.NONE
+
         server_supports_native_approach = (
             self.connection.server_parameterized_queries_enabled(
                 self.connection.protocol_version
@@ -415,7 +421,7 @@ def _determine_parameter_approach(self) -> ParameterApproach:
             )
 
     def _prepare_inline_parameters(
-        self, stmt: str, params: Union[List, Dict[str, Any]]
+        self, stmt: str, params: Optional[Union[List, Dict[str, Any]]]
     ) -> Tuple[str, List]:
         """Return a statement and list of native parameters to be passed to thrift_backend for execution
 
@@ -440,7 +446,7 @@ def _prepare_inline_parameters(
         return rendered_statement, NO_NATIVE_PARAMS
 
     def _prepare_native_parameters(
-        self, stmt: str, params: Union[List[Any], Dict[str, Any]]
+        self, stmt: str, params: Optional[Union[List[Any], Dict[str, Any]]]
     ) -> Tuple[str, List[TSparkParameter]]:
         """Return a statement and a list of native parameters to be passed to thrift_backend for execution
 
@@ -455,12 +461,12 @@ def _prepare_native_parameters(
             can be wrapped in a DbsqlParameter class.
 
         Returns a tuple of:
-            stmt: the passed statement with the param markers replaced by literal rendered values
+            stmt: the passed statement` with the param markers replaced by literal rendered values
             params: a list of TSparkParameters that will be passed in native mode
         """
 
         stmt = stmt
-        params = named_parameters_to_tsparkparams(params)
+        params = named_parameters_to_tsparkparams(params)  # type: ignore
 
         return stmt, params
 
@@ -621,7 +627,7 @@ def _handle_staging_remove(self, presigned_url: str, headers: dict = None):
     def execute(
         self,
         operation: str,
-        parameters: Optional[Union[List[Any], Dict[str, str]]] = None,
+        parameters: Optional[Union[List[Any], Dict[str, Any]]] = None,
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -652,18 +658,16 @@ def execute(
         :returns self
         """
 
-        if parameters:
-            param_approach = self._determine_parameter_approach()
-        else:
-            param_approach = ParameterApproach.NONE
+        param_approach = self._determine_parameter_approach(parameters)
+        if param_approach == ParameterApproach.NONE:
             prepared_params = NO_NATIVE_PARAMS
             prepared_operation = operation
 
-        if param_approach == ParameterApproach.INLINE:
+        elif param_approach == ParameterApproach.INLINE:
             prepared_operation, prepared_params = self._prepare_inline_parameters(
                 operation, parameters
             )
-        if param_approach == ParameterApproach.NATIVE:
+        elif param_approach == ParameterApproach.NATIVE:
             prepared_operation, prepared_params = self._prepare_native_parameters(
                 operation, parameters
             )
diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
@@ -633,7 +633,9 @@ def calculate_decimal_cast_string(input: Decimal) -> str:
     return f"DECIMAL({overall},{after})"
 
 
-def named_parameters_to_tsparkparams(parameters: Union[List[Any], Dict[str, str]]):
+def named_parameters_to_tsparkparams(
+    parameters: Union[List[Any], Dict[str, str]]
+) -> List[TSparkParameter]:
     tspark_params = []
     if isinstance(parameters, dict):
         dbsql_params = named_parameters_to_dbsqlparams_v1(parameters)
diff --git a/src/databricks/sqlalchemy/test_local/e2e/test_basic.py b/src/databricks/sqlalchemy/test_local/e2e/test_basic.py
@@ -22,6 +22,10 @@
 except ImportError:
     from sqlalchemy.ext.declarative import declarative_base
 
+from databricks.sqlalchemy.test.test_suite import start_protocol_patch
+
+start_protocol_patch()
+
 
 USER_AGENT_TOKEN = "PySQL e2e Tests"
 
diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py
@@ -130,6 +130,13 @@ def conditional_protocol_patch(self, bypass=False):
                 pass
 
     def _inline_roundtrip(self, params: dict):
+        """This INSERT, SELECT, DELETE dance is necessary because simply selecting
+        ```
+        "SELECT %(param)s"
+        ```
+        in INLINE mode would always return a str and the nature of the test is to
+        confirm that types are maintained.
+        """
         target_column = self.inline_type_map[type(params.get("p"))]
         INSERT_QUERY = f"INSERT INTO pysql_e2e_inline_param_test_table (`{target_column}`) VALUES (%(p)s)"
         SELECT_QUERY = f"SELECT {target_column} `col` FROM pysql_e2e_inline_param_test_table LIMIT 1"
