From 1d2b9c538847065e66867b922bc3ddd714a3986b Mon Sep 17 00:00:00 2001 From: claudevdm <33973061+claudevdm@users.noreply.github.com> Date: Tue, 25 Feb 2025 10:48:20 -0500 Subject: [PATCH] Pass extra write args from AlloyDB writer to WriteToJdbc. (#34067) --- sdks/python/apache_beam/ml/rag/ingestion/alloydb.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/rag/ingestion/alloydb.py b/sdks/python/apache_beam/ml/rag/ingestion/alloydb.py index 0a03706eb38..943ad89eeb4 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/alloydb.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/alloydb.py @@ -706,4 +706,11 @@ def expand(self, pcoll: beam.PCollection[Chunk]): jdbc_url=self.config.connection_config.jdbc_url, username=self.config.connection_config.username, password=self.config.connection_config.password, - statement=self.config.query_builder.build_insert())) + statement=self.config.query_builder.build_insert(), + connection_properties=self.config.connection_config. + connection_properties, + connection_init_sqls=self.config.connection_config. + connection_init_sqls, + autosharding=self.config.connection_config.autosharding, + max_connections=self.config.connection_config.max_connections, + write_batch_size=self.config.connection_config.write_batch_size))