Pass engine_args to create_engine in get_embedding_column_definition (#…

…27)

tidb_vector/integrations/utils.py

@@ -44,19 +44,25 @@ def check_table_existence(
         engine.dispose()
 
 
-def get_embedding_column_definition(connection_string, table_name, column_name):
+def get_embedding_column_definition(
+    connection_string: str,
+    table_name: str,
+    column_name: str,
+    engine_args: Optional[Dict[str, Any]] = None,
+):
     """
     Retrieves the column definition of an embedding column from a database table.
 
     Args:
         connection_string (str): The connection string to the database.
         table_name (str): The name of the table.
         column_name (str): The name of the column.
+        engine_args (Optional[Dict[str, Any]]): Additional arguments for the engine.
 
     Returns:
         tuple: A tuple containing the dimension (int or None) and distance metric (str or None).
     """
-    engine = sqlalchemy.create_engine(connection_string)
+    engine = sqlalchemy.create_engine(connection_string, **(engine_args or {}))
     try:
         with engine.connect() as connection:
             query = f"""SELECT COLUMN_TYPE, COLUMN_COMMENT

tidb_vector/integrations/vector_client.py

@@ -126,7 +126,10 @@ def _check_table_compatibility(self) -> None:
             return
 
         actual_dim, actual_distance_strategy = get_embedding_column_definition(
-            self.connection_string, self._table_name, "embedding"
+            connection_string=self.connection_string,
+            table_name=self._table_name,
+            column_name="embedding",
+            engine_args=self._engine_args,
         )
         if actual_dim is not None:
             # If the vector dimension is not set, set it to the actual dimension