// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/CachedPartitionCountReader.java
// @ThreadSafe
// public class CachedPartitionCountReader implements PartitionCountReader {
//   private final AdminClient adminClient;
//   private final Supplier<Integer> supplier;
// 
//   public CachedPartitionCountReader(AdminClient adminClient, TopicPath topicPath) {
//     this.adminClient = adminClient;
//     this.supplier =
//         Suppliers.memoizeWithExpiration(
//             () -> PartitionLookupUtils.numPartitions(topicPath, adminClient), 1, TimeUnit.MINUTES);
//   }
// 
//   @Override
//   public void close() {
//     adminClient.close();
//   }
// 
//   public int getPartitionCount() {
//     return supplier.get();
//   }
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/LimitingHeadOffsetReader.java
// public class LimitingHeadOffsetReader implements PerTopicHeadOffsetReader {
//   private static final GoogleLogger log = GoogleLogger.forEnclosingClass();
// 
//   private final TopicStatsClient topicStatsClient;
//   private final TopicPath topic;
//   private final PartitionCountReader partitionCountReader;
//   private final AsyncLoadingCache<Partition, Offset> cachedHeadOffsets;
// 
//   @VisibleForTesting
//   public LimitingHeadOffsetReader(
//       TopicStatsClient topicStatsClient,
//       TopicPath topic,
//       PartitionCountReader partitionCountReader,
//       Ticker ticker) {
//     this.topicStatsClient = topicStatsClient;
//     this.topic = topic;
//     this.partitionCountReader = partitionCountReader;
//     this.cachedHeadOffsets =
//         Caffeine.newBuilder()
//             .ticker(ticker)
//             .expireAfterWrite(1, TimeUnit.MINUTES)
//             .buildAsync(this::loadHeadOffset);
//   }
// 
//   private CompletableFuture<Offset> loadHeadOffset(Partition partition, Executor executor) {
// 
//     CompletableFuture<Offset> result = new CompletableFuture<>();
//     ApiFutures.addCallback(
//         topicStatsClient.computeHeadCursor(topic, partition),
//         new ApiFutureCallback<Cursor>() {
//           @Override
//           public void onFailure(Throwable t) {
//             result.completeExceptionally(t);
//           }
// 
//           @Override
//           public void onSuccess(Cursor c) {
//             result.complete(Offset.of(c.getOffset()));
//           }
//         },
//         MoreExecutors.directExecutor());
//     return result;
//   }
// 
//   @Override
//   public PslSourceOffset getHeadOffset() {
//     Set<Partition> keySet = new HashSet<>();
//     for (int i = 0; i < partitionCountReader.getPartitionCount(); i++) {
//       keySet.add(Partition.of(i));
//     }
//     CompletableFuture<Map<Partition, Offset>> future = cachedHeadOffsets.getAll(keySet);
//     try {
//       return PslSourceOffset.builder().partitionOffsetMap(future.get()).build();
//     } catch (Throwable t) {
//       throw new IllegalStateException("Unable to compute head offset for topic: " + topic, t);
//     }
//   }
// 
//   @Override
//   public void close() {
//     try (AutoCloseable a = topicStatsClient;
//         Closeable b = partitionCountReader) {
//     } catch (Exception e) {
//       log.atWarning().withCause(e).log("Unable to close LimitingHeadOffsetReader.");
//     }
//   }
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PartitionCountReader.java
// public interface PartitionCountReader extends Closeable {
//   int getPartitionCount();
// 
//   @Override
//   void close();
// }


import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
import com.github.benmanes.caffeine.cache.Ticker;
import com.google.auto.service.AutoService;
import com.google.cloud.pubsublite.AdminClient;
import com.google.cloud.pubsublite.SubscriptionPath;
import com.google.cloud.pubsublite.TopicPath;
import com.google.cloud.pubsublite.spark.internal.CachedPartitionCountReader;
import com.google.cloud.pubsublite.spark.internal.LimitingHeadOffsetReader;
import com.google.cloud.pubsublite.spark.internal.PartitionCountReader;
import java.util.Objects;
import java.util.Optional;
import org.apache.spark.sql.sources.DataSourceRegister;
import org.apache.spark.sql.sources.v2.ContinuousReadSupport;
import org.apache.spark.sql.sources.v2.DataSourceOptions;
import org.apache.spark.sql.sources.v2.DataSourceV2;
import org.apache.spark.sql.sources.v2.MicroBatchReadSupport;
import org.apache.spark.sql.sources.v2.StreamWriteSupport;
import org.apache.spark.sql.sources.v2.reader.streaming.ContinuousReader;
import org.apache.spark.sql.sources.v2.reader.streaming.MicroBatchReader;
import org.apache.spark.sql.sources.v2.writer.streaming.StreamWriter;
import org.apache.spark.sql.streaming.OutputMode;
import org.apache.spark.sql.types.StructType;
    }
    PartitionCountReader partitionCountReader =
        new CachedPartitionCountReader(pslReadDataSourceOptions.newAdminClient(), topicPath);
    return new PslContinuousReader(
        pslReadDataSourceOptions.newCursorClient(),
        pslReadDataSourceOptions.newMultiPartitionCommitter(
            partitionCountReader.getPartitionCount()),
        pslReadDataSourceOptions.getSubscriberFactory(),
        subscriptionPath,
        Objects.requireNonNull(pslReadDataSourceOptions.flowControlSettings()),
        partitionCountReader);
  }

  @Override
  public MicroBatchReader createMicroBatchReader(
      Optional<StructType> schema, String checkpointLocation, DataSourceOptions options) {
    if (schema.isPresent()) {
      throw new IllegalArgumentException(
          "PubSub Lite uses fixed schema and custom schema is not allowed");
    }

    PslReadDataSourceOptions pslReadDataSourceOptions =
        PslReadDataSourceOptions.fromSparkDataSourceOptions(options);
    SubscriptionPath subscriptionPath = pslReadDataSourceOptions.subscriptionPath();
    TopicPath topicPath;
    try (AdminClient adminClient = pslReadDataSourceOptions.newAdminClient()) {
      topicPath = TopicPath.parse(adminClient.getSubscription(subscriptionPath).get().getTopic());
    } catch (Throwable t) {
      throw toCanonical(t).underlying;
    }
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/CachedPartitionCountReader.java
// @ThreadSafe
// public class CachedPartitionCountReader implements PartitionCountReader {
//   private final AdminClient adminClient;
//   private final Supplier<Integer> supplier;
// 
//   public CachedPartitionCountReader(AdminClient adminClient, TopicPath topicPath) {
//     this.adminClient = adminClient;
//     this.supplier =
//         Suppliers.memoizeWithExpiration(
//             () -> PartitionLookupUtils.numPartitions(topicPath, adminClient), 1, TimeUnit.MINUTES);
//   }
// 
//   @Override
//   public void close() {
//     adminClient.close();
//   }
// 
//   public int getPartitionCount() {
//     return supplier.get();
//   }
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/LimitingHeadOffsetReader.java
// public class LimitingHeadOffsetReader implements PerTopicHeadOffsetReader {
//   private static final GoogleLogger log = GoogleLogger.forEnclosingClass();
// 
//   private final TopicStatsClient topicStatsClient;
//   private final TopicPath topic;
//   private final PartitionCountReader partitionCountReader;
//   private final AsyncLoadingCache<Partition, Offset> cachedHeadOffsets;
// 
//   @VisibleForTesting
//   public LimitingHeadOffsetReader(
//       TopicStatsClient topicStatsClient,
//       TopicPath topic,
//       PartitionCountReader partitionCountReader,
//       Ticker ticker) {
//     this.topicStatsClient = topicStatsClient;
//     this.topic = topic;
//     this.partitionCountReader = partitionCountReader;
//     this.cachedHeadOffsets =
//         Caffeine.newBuilder()
//             .ticker(ticker)
//             .expireAfterWrite(1, TimeUnit.MINUTES)
//             .buildAsync(this::loadHeadOffset);
//   }
// 
//   private CompletableFuture<Offset> loadHeadOffset(Partition partition, Executor executor) {
// 
//     CompletableFuture<Offset> result = new CompletableFuture<>();
//     ApiFutures.addCallback(
//         topicStatsClient.computeHeadCursor(topic, partition),
//         new ApiFutureCallback<Cursor>() {
//           @Override
//           public void onFailure(Throwable t) {
//             result.completeExceptionally(t);
//           }
// 
//           @Override
//           public void onSuccess(Cursor c) {
//             result.complete(Offset.of(c.getOffset()));
//           }
//         },
//         MoreExecutors.directExecutor());
//     return result;
//   }
// 
//   @Override
//   public PslSourceOffset getHeadOffset() {
//     Set<Partition> keySet = new HashSet<>();
//     for (int i = 0; i < partitionCountReader.getPartitionCount(); i++) {
//       keySet.add(Partition.of(i));
//     }
//     CompletableFuture<Map<Partition, Offset>> future = cachedHeadOffsets.getAll(keySet);
//     try {
//       return PslSourceOffset.builder().partitionOffsetMap(future.get()).build();
//     } catch (Throwable t) {
//       throw new IllegalStateException("Unable to compute head offset for topic: " + topic, t);
//     }
//   }
// 
//   @Override
//   public void close() {
//     try (AutoCloseable a = topicStatsClient;
//         Closeable b = partitionCountReader) {
//     } catch (Exception e) {
//       log.atWarning().withCause(e).log("Unable to close LimitingHeadOffsetReader.");
//     }
//   }
// }
// 
// Path: src/main/java/com/google/cloud/pubsublite/spark/internal/PartitionCountReader.java
// public interface PartitionCountReader extends Closeable {
//   int getPartitionCount();
// 
//   @Override
//   void close();
// }


// Path: src/main/java/com/google/cloud/pubsublite/spark/PslDataSource.java
import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
import com.github.benmanes.caffeine.cache.Ticker;
import com.google.auto.service.AutoService;
import com.google.cloud.pubsublite.AdminClient;
import com.google.cloud.pubsublite.SubscriptionPath;
import com.google.cloud.pubsublite.TopicPath;
import com.google.cloud.pubsublite.spark.internal.CachedPartitionCountReader;
import com.google.cloud.pubsublite.spark.internal.LimitingHeadOffsetReader;
import com.google.cloud.pubsublite.spark.internal.PartitionCountReader;
import java.util.Objects;
import java.util.Optional;
import org.apache.spark.sql.sources.DataSourceRegister;
import org.apache.spark.sql.sources.v2.ContinuousReadSupport;
import org.apache.spark.sql.sources.v2.DataSourceOptions;
import org.apache.spark.sql.sources.v2.DataSourceV2;
import org.apache.spark.sql.sources.v2.MicroBatchReadSupport;
import org.apache.spark.sql.sources.v2.StreamWriteSupport;
import org.apache.spark.sql.sources.v2.reader.streaming.ContinuousReader;
import org.apache.spark.sql.sources.v2.reader.streaming.MicroBatchReader;
import org.apache.spark.sql.sources.v2.writer.streaming.StreamWriter;
import org.apache.spark.sql.streaming.OutputMode;
import org.apache.spark.sql.types.StructType;

    }
    PartitionCountReader partitionCountReader =
        new CachedPartitionCountReader(pslReadDataSourceOptions.newAdminClient(), topicPath);
    return new PslContinuousReader(
        pslReadDataSourceOptions.newCursorClient(),
        pslReadDataSourceOptions.newMultiPartitionCommitter(
            partitionCountReader.getPartitionCount()),
        pslReadDataSourceOptions.getSubscriberFactory(),
        subscriptionPath,
        Objects.requireNonNull(pslReadDataSourceOptions.flowControlSettings()),
        partitionCountReader);
  }

  @Override
  public MicroBatchReader createMicroBatchReader(
      Optional<StructType> schema, String checkpointLocation, DataSourceOptions options) {
    if (schema.isPresent()) {
      throw new IllegalArgumentException(
          "PubSub Lite uses fixed schema and custom schema is not allowed");
    }

    PslReadDataSourceOptions pslReadDataSourceOptions =
        PslReadDataSourceOptions.fromSparkDataSourceOptions(options);
    SubscriptionPath subscriptionPath = pslReadDataSourceOptions.subscriptionPath();
    TopicPath topicPath;
    try (AdminClient adminClient = pslReadDataSourceOptions.newAdminClient()) {
      topicPath = TopicPath.parse(adminClient.getSubscription(subscriptionPath).get().getTopic());
    } catch (Throwable t) {
      throw toCanonical(t).underlying;
    }
    PartitionCountReader partitionCountReader =