s3-querier 1.0.0 → 1.0.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "s3-querier",
-  "version": "1.0.0",
+  "version": "1.0.2",
   "description": "Query S3-compatible storage with DuckDB and SQL",
   "type": "module",
   "main": "src/s3-querier.js",

package/src/s3/s3.js

@@ -3,7 +3,7 @@ import { dirname } from 'node:path';
 import { S3Client, ListObjectsV2Command, GetObjectCommand } from '@aws-sdk/client-s3';
 
 import { logger } from '../utils/logger.js';
-import { datesInRange, hoursInRange, buildPath } from '../utils/file-path-builder/file-path-builder.js';
+import { datesInRange, hoursInRange, monthsInRange, buildPath } from '../utils/file-path-builder/file-path-builder.js';
 import { regexFromPattern } from '../utils/date-regex/date-regex.js';
 import { buildIbmIamClient } from './auth/ibm-iam-client.js';
 
@@ -139,12 +139,14 @@ export default class S3 {
    * @returns {(from:Date, to:Date, filePattern:string) => string[]} A function that creates a list of prefixes
    */
   prefixStrategy(from, to, filePattern) {
-    const hasDateToken = filePattern.match(/\{(yyyy|MM|dd|hh|mm)\}/g);
-    const hasGlob = filePattern.match(/\*/g);
+    const hasDayToken = /\{(dd|hh|mm)\}/.test(filePattern);
+    const hasMonthToken = /\{(yyyy|MM)\}/.test(filePattern);
+    const hasGlob = filePattern.includes('*');
     const hourDiff = (new Date(to) - new Date(from)) / 1000 / 60 / 60;
 
-    if (hasDateToken && hourDiff < 24) return this.prefixHours;
-    if (hasDateToken) return this.prefixDays;
+    if (hasDayToken && hourDiff < 24) return this.prefixHours;
+    if (hasDayToken) return this.prefixDays;
+    if (hasMonthToken) return this.prefixMonths;
     if (hasGlob) return this.prefixGlob;
     return (_from, _to, pattern) => [pattern];
   }
@@ -182,6 +184,22 @@ export default class S3 {
     });
   }
 
+  /**
+   * Returns a list of prefixes based on a range of months, one per calendar month.
+   * Used for Hive-style paths with {yyyy}/{MM} tokens but no {dd}.
+   *
+   * @param {Date} from From date
+   * @param {Date} to To date
+   * @param {string} filePattern The file pattern to use
+   * @returns {string[]} The list of prefixes for filtering
+   */
+  prefixMonths(from, to, filePattern) {
+    const monthRange = monthsInRange(new Date(from), new Date(to));
+    const splitToken = filePattern.includes('{MM}') ? '{MM}' : '{yyyy}';
+    const [trimmed] = filePattern.split(splitToken);
+    return monthRange.map((date) => buildPath(`${trimmed}${splitToken}`, date));
+  }
+
   /**
    * Returns a single entry array with a file pattern trimmed to the first glob
    *
@@ -368,17 +386,20 @@ export default class S3 {
    */
   async objectToFile(key) {
     const file = `${this.mount}/${key}`;
+    const tmp = `${file}.${process.pid}.${Date.now()}.tmp`;
     try {
       const response = await this.s3.send(new GetObjectCommand({ Bucket: this.bucket, Key: key }));
       const chunks = [];
       for await (const chunk of response.Body) {
         chunks.push(chunk);
       }
-      await fsPromise.writeFile(file, Buffer.concat(chunks));
+      await fsPromise.writeFile(tmp, Buffer.concat(chunks));
+      await fsPromise.rename(tmp, file);
       await this.processFile(file);
       return file;
     } catch (error) {
       logger.error(`${error.$metadata?.httpStatusCode ?? error.statusCode} - ${file}`);
+      await fsPromise.unlink(tmp).catch(() => {});
       throw error;
     }
   }
@@ -419,7 +440,7 @@ export default class S3 {
   }
 }
 
-function buildS3Client({ apiKey, accessKeyId, secretAccessKey, endpoint, region }) {
+export function buildS3Client({ apiKey, accessKeyId, secretAccessKey, endpoint, region = 'us-east-1' }) {
   const config = { ...(endpoint && { endpoint }), region, forcePathStyle: true };
 
   if (apiKey) return buildIbmIamClient(config, apiKey);

package/src/utils/file-path-builder/file-path-builder.js

@@ -28,6 +28,22 @@ export function datesInRange(from, to) {
   });
 }
 
+/**
+ * Given a date range returns one date object per calendar month.
+ * Returns noon-UTC dates so that getMonth()/getFullYear() produce the correct
+ * UTC month in any timezone.
+ *
+ * @param {Date} from The from Date object
+ * @param {Date} to The to Date object
+ * @returns {Date[]} An array of dates, one per month within the range
+ */
+export function monthsInRange(from, to) {
+  const startYear = from.getUTCFullYear();
+  const startMonth = from.getUTCMonth();
+  const count = (to.getUTCFullYear() - startYear) * 12 + (to.getUTCMonth() - startMonth) + 1;
+  return Array.from({ length: count }, (_, index) => noonUtcForMonthOffset(startYear, startMonth, index));
+}
+
 /**
  * Given a date range returns an array of date objects
  *
@@ -53,3 +69,9 @@ export function zeroDateMins(date) {
   zeroedDate.setMinutes(0, 0);
   return zeroedDate;
 }
+
+function noonUtcForMonthOffset(startYear, startMonth, offset) {
+  const year = startYear + Math.floor((startMonth + offset) / 12);
+  const month = (startMonth + offset) % 12;
+  return new Date(Date.UTC(year, month, 1, 12, 0, 0));
+}