Pass extra write args from AlloyDB writer to WriteToJdbc. (apache#34067)

lostluck · Feb 25, 2025 · 1d2b9c5 · 1d2b9c5
1 parent cbe91f0
commit 1d2b9c5
Showing 1 changed file with 8 additions and 1 deletion.
diff --git 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))