graphdatascience 1.16__tar.gz → 1.17a1__tar.gz

{graphdatascience-1.16/graphdatascience.egg-info → graphdatascience-1.17a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.16
+Version: 1.17a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -38,6 +38,7 @@ Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0
 Requires-Dist: requests
 Requires-Dist: tenacity>=9.0
+Requires-Dist: pydantic>=2.11
 Provides-Extra: ogb
 Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: networkx

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/graph/graph_cypher_runner.py

@@ -6,6 +6,8 @@ from typing import Any, Optional
 
 from pandas import Series
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..caller_base import CallerBase
 from ..query_runner.query_runner import QueryRunner
 from ..server_version.server_version import ServerVersion
@@ -46,7 +48,7 @@ class GraphCypherRunner(CallerBase):
         GraphCypherRunner._verify_query_ends_with_return_clause(self._namespace, query)
 
         result: Optional[dict[str, Any]] = self._query_runner.run_retryable_cypher(
-            query, params, database, custom_error=False
+            query, params, database, custom_error=False, mode=QueryMode.READ
         ).squeeze()
 
         if not result:

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -5,6 +5,8 @@ from warnings import filterwarnings
 import pandas as pd
 from pandas import DataFrame, Series
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..call_parameters import CallParameters
 from ..error.cypher_warning_handler import (
     filter_id_func_deprecation_warning,
@@ -164,7 +166,9 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
             unique_node_ids = result["nodeId"].drop_duplicates().tolist()
 
             db_properties_df = query_runner.run_retryable_cypher(
-                GraphNodePropertiesRunner._build_query(db_node_properties), params={"ids": unique_node_ids}
+                GraphNodePropertiesRunner._build_query(db_node_properties),
+                params={"ids": unique_node_ids},
+                mode=QueryMode.READ,
             )
 
             if "propertyValue" not in result.keys():

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/graph/graph_object.py

@@ -68,7 +68,7 @@ class Graph:
         """
         return self._graph_info(["database"])  # type: ignore
 
-    def configuration(self) -> "Series[Any]":
+    def configuration(self) -> Series[Any]:
         """
         Returns:
             the configuration of the graph

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/graph_data_science.py

@@ -9,6 +9,7 @@ from neo4j import Driver
 from pandas import DataFrame
 
 from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
+from graphdatascience.query_runner.query_mode import QueryMode
 
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
@@ -41,6 +42,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
         show_progress: bool = True,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ):
         """
         Construct a new GraphDataScience object.
@@ -64,14 +66,22 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
                 - True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint.
                 - False will make the client use Bolt for all operations.
         arrow_disable_server_verification : bool, default True
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
+
             A flag that overrides other TLS settings and disables server verification for TLS connections.
         arrow_tls_root_certs : Optional[bytes], default None
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
+
             PEM-encoded certificates that are used for the connection to the
             GDS Arrow Flight server.
         bookmarks : Optional[Any], default None
             The Neo4j bookmarks to require a certain state before the next query gets executed.
         show_progress : bool, default True
             A flag to indicate whether to show progress bars for running procedures.
+        arrow_client_options : Optional[dict[str, Any]], default None
+            Additional options to be passed to the Arrow Flight client.
         """
         if aura_ds:
             GraphDataScience._validate_endpoint(endpoint)
@@ -104,14 +114,19 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
                 username, password = auth
                 arrow_auth = UsernamePasswordAuthentication(username, password)
 
+            if arrow_client_options is None:
+                arrow_client_options = {}
+            if arrow_disable_server_verification:
+                arrow_client_options["disable_server_verification"] = True
+            if arrow_tls_root_certs is not None:
+                arrow_client_options["tls_root_certs"] = arrow_tls_root_certs
             self._query_runner = ArrowQueryRunner.create(
                 self._query_runner,
-                arrow_info,
-                arrow_auth,
-                self._query_runner.encrypted(),
-                arrow_disable_server_verification,
-                arrow_tls_root_certs,
-                None if arrow is True else arrow,
+                arrow_info=arrow_info,
+                arrow_authentication=arrow_auth,
+                encrypted=self._query_runner.encrypted(),
+                arrow_client_options=arrow_client_options,
+                connection_string_override=None if arrow is True else arrow,
             )
 
         self._query_runner.set_show_progress(show_progress)
@@ -199,7 +214,12 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         return self._query_runner.last_bookmarks()
 
     def run_cypher(
-        self, query: str, params: Optional[dict[str, Any]] = None, database: Optional[str] = None
+        self,
+        query: str,
+        params: Optional[dict[str, Any]] = None,
+        database: Optional[str] = None,
+        retryable: bool = False,
+        mode: QueryMode = QueryMode.WRITE,
     ) -> DataFrame:
         """
         Run a Cypher query
@@ -212,6 +232,10 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
             parameters to the query
         database: str
             the database on which to run the query
+        retryable: bool
+            whether the query can be automatically retried. Make sure the query is idempotent if set to True.
+        mode: QueryMode
+            the query mode to use (READ or WRITE). Set based on the operation performed in the query.
 
         Returns:
             The query result as a DataFrame
@@ -222,8 +246,10 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         if isinstance(self._query_runner, ArrowQueryRunner):
             qr = self._query_runner.fallback_query_runner()
 
-        # not using qr.run_retryable_cypher as we dont know if it can be retried
-        return qr.run_cypher(query, params, database, False)
+        if retryable:
+            return qr.run_retryable_cypher(query, params, database, custom_error=False, mode=mode)
+        else:
+            return qr.run_cypher(query, params, database, custom_error=False, mode=mode)
 
     def driver_config(self) -> dict[str, Any]:
         """
@@ -240,11 +266,51 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         driver: Driver,
         auth: Optional[tuple[str, str]] = None,
         database: Optional[str] = None,
-        arrow: bool = True,
+        arrow: Union[str, bool] = True,
         arrow_disable_server_verification: bool = True,
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
-    ) -> "GraphDataScience":
+        arrow_client_options: Optional[dict[str, Any]] = None,
+    ) -> GraphDataScience:
+        """
+        Construct a new GraphDataScience object from an existing Neo4j Driver.
+        This method is useful when you already have a Neo4j Driver instance and want to use it with the GDS client.
+
+        Parameters
+        ----------
+        driver: Driver
+            The Neo4j Driver instance to use.
+        auth : Optional[Tuple[str, str]], default None
+            A username, password pair for authentication.
+        database: Optional[str], default None
+            The Neo4j database to query against.
+        arrow : Union[str, bool], default True
+            Arrow connection information. This is either a string or a bool.
+
+            - If it is a string, it will be interpreted as a connection URL to a GDS Arrow Server.
+            - If it is a bool:
+                - True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint.
+                - False will make the client use Bolt for all operations.
+        arrow_disable_server_verification : bool, default True
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
+
+            A flag that overrides other TLS settings and disables server verification for TLS connections.
+        arrow_tls_root_certs : Optional[bytes], default None
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
+
+            PEM-encoded certificates that are used for the connection to the
+            GDS Arrow Flight server.
+        bookmarks : Optional[Any], default None
+            The Neo4j bookmarks to require a certain state before the next query gets executed.
+        show_progress : bool, default True
+            A flag to indicate whether to show progress bars for running procedures.
+        arrow_client_options : Optional[dict[str, Any]], default None
+            Additional options to be passed to the Arrow Flight client.
+        Returns:
+            A new GraphDataScience object. configured with the provided Neo4j Driver.
+        """
         return cls(
             driver,
             auth=auth,
@@ -253,6 +319,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
             arrow_disable_server_verification=arrow_disable_server_verification,
             arrow_tls_root_certs=arrow_tls_root_certs,
             bookmarks=bookmarks,
+            arrow_client_options=arrow_client_options,
         )
 
     @staticmethod

graphdatascience-1.17a1/graphdatascience/procedure_surface/arrow/wcc_arrow_endpoints.py

@@ -0,0 +1,191 @@
+import json
+from typing import Any, List, Optional
+
+from pandas import DataFrame
+
+from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from ...arrow_client.v2.data_mapper_utils import deserialize_single
+from ...arrow_client.v2.job_client import JobClient
+from ...arrow_client.v2.mutation_client import MutationClient
+from ...arrow_client.v2.write_back_client import WriteBackClient
+from ...graph.graph_object import Graph
+from ..api.estimation_result import EstimationResult
+from ..api.wcc_endpoints import WccEndpoints, WccMutateResult, WccStatsResult, WccWriteResult
+from ..utils.config_converter import ConfigConverter
+
+WCC_ENDPOINT = "v2/community.wcc"
+
+
+class WccArrowEndpoints(WccEndpoints):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[WriteBackClient] = None):
+        self._arrow_client = arrow_client
+        self._write_back_client = write_back_client
+
+    def mutate(
+        self,
+        G: Graph,
+        mutate_property: str,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> WccMutateResult:
+        config = ConfigConverter.convert_to_gds_config(
+            graph_name=G.name(),
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+        )
+
+        job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
+
+        mutate_result = MutationClient.mutate_node_property(self._arrow_client, job_id, mutate_property)
+        computation_result = JobClient.get_summary(self._arrow_client, job_id)
+
+        computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = 0
+
+        return WccMutateResult(**computation_result)
+
+    def stats(
+        self,
+        G: Graph,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> WccStatsResult:
+        config = ConfigConverter.convert_to_gds_config(
+            graph_name=G.name(),
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+        )
+
+        job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
+        computation_result = JobClient.get_summary(self._arrow_client, job_id)
+
+        return WccStatsResult(**computation_result)
+
+    def stream(
+        self,
+        G: Graph,
+        min_component_size: Optional[int] = None,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> DataFrame:
+        config = ConfigConverter.convert_to_gds_config(
+            graph_name=G.name(),
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            min_component_size=min_component_size,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+        )
+
+        job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
+        return JobClient.stream_results(self._arrow_client, G.name(), job_id)
+
+    def write(
+        self,
+        G: Graph,
+        write_property: str,
+        min_component_size: Optional[int] = None,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+        write_concurrency: Optional[int] = None,
+    ) -> WccWriteResult:
+        config = ConfigConverter.convert_to_gds_config(
+            graph_name=G.name(),
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            min_component_size=min_component_size,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+        )
+
+        job_id = JobClient.run_job_and_wait(self._arrow_client, WCC_ENDPOINT, config)
+        computation_result = JobClient.get_summary(self._arrow_client, job_id)
+
+        if self._write_back_client is None:
+            raise Exception("Write back client is not initialized")
+
+        write_millis = self._write_back_client.write(
+            G.name(), job_id, write_concurrency if write_concurrency is not None else concurrency
+        )
+
+        computation_result["writeMillis"] = write_millis
+
+        return WccWriteResult(**computation_result)
+
+    def estimate(
+        self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
+    ) -> EstimationResult:
+        if G is not None:
+            payload = {"graphName": G.name()}
+        elif projection_config is not None:
+            payload = projection_config
+        else:
+            raise ValueError("Either graph_name or projection_config must be provided.")
+
+        res = self._arrow_client.do_action_with_retry("v2/community.wcc.estimate", json.dumps(payload).encode("utf-8"))
+
+        return EstimationResult(**deserialize_single(res))

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/query_runner/arrow_query_runner.py

@@ -25,8 +25,7 @@ class ArrowQueryRunner(QueryRunner):
         arrow_info: ArrowInfo,
         arrow_authentication: Optional[ArrowAuthentication] = None,
         encrypted: bool = False,
-        disable_server_verification: bool = False,
-        tls_root_certs: Optional[bytes] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
         connection_string_override: Optional[str] = None,
         retry_config: Optional[RetryConfig] = None,
     ) -> ArrowQueryRunner:
@@ -34,13 +33,12 @@ class ArrowQueryRunner(QueryRunner):
             raise ValueError("Arrow is not enabled on the server")
 
         gds_arrow_client = GdsArrowClient.create(
-            arrow_info,
-            arrow_authentication,
-            encrypted,
-            disable_server_verification,
-            tls_root_certs,
-            connection_string_override,
+            arrow_info=arrow_info,
+            auth=arrow_authentication,
+            encrypted=encrypted,
+            connection_string_override=connection_string_override,
             retry_config=retry_config,
+            arrow_client_options=arrow_client_options,
         )
 
         return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())
@@ -65,18 +63,22 @@ class ArrowQueryRunner(QueryRunner):
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._fallback_query_runner.run_cypher(query, params, database, custom_error)
+        return self._fallback_query_runner.run_cypher(query, params, database, mode, custom_error=custom_error)
 
     def run_retryable_cypher(
         self,
         query: str,
         params: Optional[dict[str, Any]] = None,
         database: Optional[str] = None,
+        mode: Optional[QueryMode] = None,
         custom_error: bool = True,
     ) -> DataFrame:
-        return self._fallback_query_runner.run_retryable_cypher(query, params, database, custom_error=custom_error)
+        return self._fallback_query_runner.run_retryable_cypher(
+            query, params, database, mode, custom_error=custom_error
+        )
 
     def call_function(self, endpoint: str, params: Optional[CallParameters] = None) -> Any:
         return self._fallback_query_runner.call_function(endpoint, params)

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/query_runner/cypher_graph_constructor.py

@@ -7,6 +7,8 @@ from uuid import uuid4
 
 from pandas import DataFrame, concat
 
+from graphdatascience.query_runner.query_mode import QueryMode
+
 from ..server_version.server_version import ServerVersion
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
@@ -105,7 +107,9 @@ class CypherGraphConstructor(GraphConstructor):
     def _should_warn_about_arrow_missing(self) -> bool:
         try:
             license: str = self._query_runner.run_retryable_cypher(
-                "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value", custom_error=False
+                "CALL gds.debug.sysInfo() YIELD key, value WHERE key = 'gdsEdition' RETURN value",
+                custom_error=False,
+                mode=QueryMode.READ,
             ).squeeze()
             should_warn = license == "Licensed"
         except Exception as e:

{graphdatascience-1.16 → graphdatascience-1.17a1}/graphdatascience/query_runner/gds_arrow_client.py

@@ -55,6 +55,7 @@ class GdsArrowClient:
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
         retry_config: Optional[RetryConfig] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ) -> GdsArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -78,14 +79,15 @@ class GdsArrowClient:
             )
 
         return GdsArrowClient(
-            host,
-            retry_config,
-            int(port),
-            auth,
-            encrypted,
-            disable_server_verification,
-            tls_root_certs,
-            arrow_endpoint_version,
+            host=host,
+            retry_config=retry_config,
+            port=int(port),
+            auth=auth,
+            encrypted=encrypted,
+            disable_server_verification=disable_server_verification,
+            tls_root_certs=tls_root_certs,
+            arrow_endpoint_version=arrow_endpoint_version,
+            arrow_client_options=arrow_client_options,
         )
 
     def __init__(
@@ -99,6 +101,7 @@ class GdsArrowClient:
         tls_root_certs: Optional[bytes] = None,
         arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.V1,
         user_agent: Optional[str] = None,
+        arrow_client_options: Optional[dict[str, Any]] = None,
     ):
         """Creates a new GdsArrowClient instance.
 
@@ -113,8 +116,12 @@ class GdsArrowClient:
         encrypted: bool
             A flag that indicates whether the connection should be encrypted (default is False)
         disable_server_verification: bool
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             A flag that disables server verification for TLS connections (default is False)
         tls_root_certs: Optional[bytes]
+            .. deprecated:: 1.16
+                Use arrow_client_options instead
             PEM-encoded certificates that are used for the connection to the GDS Arrow Flight server
         arrow_endpoint_version:
             The version of the Arrow endpoint to use (default is ArrowEndpointVersion.V1)
@@ -122,18 +129,26 @@ class GdsArrowClient:
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
         retry_config: Optional[RetryConfig]
             The retry configuration to use for the Arrow requests send by the client.
+        arrow_client_options: Optional[dict[str, Any]]
+            Additional configuration for the Arrow flight client.
+
         """
         self._arrow_endpoint_version = arrow_endpoint_version
         self._host = host
         self._port = port
         self._auth = None
         self._encrypted = encrypted
-        self._disable_server_verification = disable_server_verification
-        self._tls_root_certs = tls_root_certs
         self._user_agent = user_agent
         self._retry_config = retry_config
         self._logger = logging.getLogger("gds_arrow_client")
 
+        self._arrow_client_options = arrow_client_options if arrow_client_options is not None else {}
+
+        if disable_server_verification:
+            self._arrow_client_options["disable_server_verification"] = True
+        if tls_root_certs is not None:
+            self._arrow_client_options["tls_root_certs"] = tls_root_certs
+
         if auth:
             if not isinstance(auth, ArrowAuthentication):
                 username, password = auth
@@ -149,18 +164,27 @@ class GdsArrowClient:
             if self._encrypted
             else flight.Location.for_grpc_tcp(self._host, self._port)
         )
-        client_options: dict[str, Any] = {"disable_server_verification": self._disable_server_verification}
+
+        client_options = self._arrow_client_options.copy()
+
         if self._auth:
             user_agent = f"neo4j-graphdatascience-v{__version__} pyarrow-v{arrow_version}"
             if self._user_agent:
                 user_agent = self._user_agent
 
-            client_options["middleware"] = [
-                AuthFactory(self._auth_middleware),
-                UserAgentFactory(useragent=user_agent),
-            ]
-        if self._tls_root_certs:
-            client_options["tls_root_certs"] = self._tls_root_certs
+            if "middleware" in client_options:
+                if not isinstance(client_options["middleware"], list):
+                    raise TypeError("client_options['middleware'] must be a list")
+            else:
+                client_options["middleware"] = []
+
+            client_options["middleware"].extend(
+                [
+                    AuthFactory(self._auth_middleware),
+                    UserAgentFactory(useragent=user_agent),
+                ]
+            )
+
         return flight.FlightClient(location, **client_options)
 
     def connection_info(self) -> tuple[str, int]: