From aef5f66ff61bd411a28890ba68629e11ec077907 Mon Sep 17 00:00:00 2001
From: Ivan Katkov <44121163+Pseudolukian@users.noreply.github.com>
Date: Fri, 29 Nov 2024 08:54:26 +0100
Subject: [PATCH] Add root_certificates option for ydb.DriverConfig (#525)

* Add root_certificates option for ydb.DriverConfig in example

---------

Co-authored-by: Oleg Ovcharuk <vgvoleg@gmail.com>
---
 examples/static-credentials/example.py | 27 +++++++++++++++++++++++++-
 ydb/auth_helpers.py                    |  5 +++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/examples/static-credentials/example.py b/examples/static-credentials/example.py
index 71409f5c..7a31e07f 100644
--- a/examples/static-credentials/example.py
+++ b/examples/static-credentials/example.py
@@ -2,17 +2,42 @@
 
 
 def test_driver_works(driver: ydb.Driver):
+    """Tests the functionality of the YDB driver.
+
+    Waits for the driver to become ready and executes a simple SQL query to verify that the driver works as expected.
+
+    Args:
+        driver (ydb.Driver): The YDB driver instance to test.
+
+    Raises:
+        AssertionError: If the SQL query does not return the expected result.
+    """
     driver.wait(5)
     pool = ydb.QuerySessionPool(driver)
     result = pool.execute_with_retries("SELECT 1 as cnt")
     assert result[0].rows[0].cnt == 1
 
 
-def auth_with_static_credentials(endpoint: str, database: str, user: str, password: str):
+def auth_with_static_credentials(endpoint: str, database: str, user: str, password: str, ca_path: str):
+    """Authenticate using static credentials.
+
+    Args:
+        endpoint (str): Accepts a string in the format `grpcs://<node-fqdn>:2136` or `grpcs://<node-ip>:2136`.
+        database (str): Accepts a string, the database name in the format `/Root/<database-name>`.
+        user (str): Username.
+        password (str): User password.
+        ca_path (str): Path to CA cert
+
+    Notes:
+        The argument `root_certificates` of the function `ydb.DriverConfig` takes the content of the cluster's root certificate
+        for connecting to cluster nodes via TLS.
+        Note that the VM from which you are connecting must be in the cluster's domain for which the CA certificate is issued.
+    """
     driver_config = ydb.DriverConfig(
         endpoint=endpoint,
         database=database,
         credentials=ydb.StaticCredentials.from_user_password(user, password),
+        root_certificates=ydb.load_ydb_root_certificate(ca_path),
     )
 
     with ydb.Driver(driver_config=driver_config) as driver:
diff --git a/ydb/auth_helpers.py b/ydb/auth_helpers.py
index 6399c3cf..abf7331a 100644
--- a/ydb/auth_helpers.py
+++ b/ydb/auth_helpers.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import os
+from typing import Optional
 
 
 def read_bytes(f):
@@ -7,8 +8,8 @@ def read_bytes(f):
         return fr.read()
 
 
-def load_ydb_root_certificate():
-    path = os.getenv("YDB_SSL_ROOT_CERTIFICATES_FILE", None)
+def load_ydb_root_certificate(path: Optional[str] = None):
+    path = path if path is not None else os.getenv("YDB_SSL_ROOT_CERTIFICATES_FILE", None)
     if path is not None and os.path.exists(path):
         return read_bytes(path)
     return None