1616
1717package com .google .cloud .pubsublite .spark ;
1818
19+ import static com .google .cloud .pubsublite .internal .ExtractStatus .toCanonical ;
20+
1921import com .github .benmanes .caffeine .cache .Ticker ;
2022import com .google .auto .service .AutoService ;
2123import com .google .cloud .pubsublite .AdminClient ;
2224import com .google .cloud .pubsublite .PartitionLookupUtils ;
2325import com .google .cloud .pubsublite .SubscriptionPath ;
2426import com .google .cloud .pubsublite .TopicPath ;
25- import com .google .cloud .pubsublite .internal .CursorClient ;
2627import java .util .Objects ;
2728import java .util .Optional ;
2829import org .apache .spark .sql .sources .DataSourceRegister ;
@@ -53,12 +54,13 @@ public ContinuousReader createContinuousReader(
5354
5455 PslDataSourceOptions pslDataSourceOptions =
5556 PslDataSourceOptions .fromSparkDataSourceOptions (options );
56- CursorClient cursorClient = pslDataSourceOptions .newCursorClient ();
57- AdminClient adminClient = pslDataSourceOptions .newAdminClient ();
5857 SubscriptionPath subscriptionPath = pslDataSourceOptions .subscriptionPath ();
59- long topicPartitionCount = PartitionLookupUtils .numPartitions (subscriptionPath , adminClient );
58+ long topicPartitionCount ;
59+ try (AdminClient adminClient = pslDataSourceOptions .newAdminClient ()) {
60+ topicPartitionCount = PartitionLookupUtils .numPartitions (subscriptionPath , adminClient );
61+ }
6062 return new PslContinuousReader (
61- cursorClient ,
63+ pslDataSourceOptions . newCursorClient () ,
6264 pslDataSourceOptions .newMultiPartitionCommitter (topicPartitionCount ),
6365 pslDataSourceOptions .getSubscriberFactory (),
6466 subscriptionPath ,
@@ -76,19 +78,17 @@ public MicroBatchReader createMicroBatchReader(
7678
7779 PslDataSourceOptions pslDataSourceOptions =
7880 PslDataSourceOptions .fromSparkDataSourceOptions (options );
79- CursorClient cursorClient = pslDataSourceOptions .newCursorClient ();
80- AdminClient adminClient = pslDataSourceOptions .newAdminClient ();
8181 SubscriptionPath subscriptionPath = pslDataSourceOptions .subscriptionPath ();
8282 TopicPath topicPath ;
83- try {
83+ long topicPartitionCount ;
84+ try (AdminClient adminClient = pslDataSourceOptions .newAdminClient ()) {
8485 topicPath = TopicPath .parse (adminClient .getSubscription (subscriptionPath ).get ().getTopic ());
86+ topicPartitionCount = PartitionLookupUtils .numPartitions (topicPath , adminClient );
8587 } catch (Throwable t ) {
86- throw new IllegalStateException (
87- "Unable to get topic for subscription " + subscriptionPath , t );
88+ throw toCanonical (t ).underlying ;
8889 }
89- long topicPartitionCount = PartitionLookupUtils .numPartitions (topicPath , adminClient );
9090 return new PslMicroBatchReader (
91- cursorClient ,
91+ pslDataSourceOptions . newCursorClient () ,
9292 pslDataSourceOptions .newMultiPartitionCommitter (topicPartitionCount ),
9393 pslDataSourceOptions .getSubscriberFactory (),
9494 new LimitingHeadOffsetReader (
0 commit comments