From 05a8f146fa5d3c751dfade4ad371e94aa4f39489 Mon Sep 17 00:00:00 2001 From: Gabriel Bressan Date: Thu, 28 May 2026 12:11:16 -0400 Subject: [PATCH 1/2] feat: improve s3 handler file listing --- .../handlers/s3_handler/README.md | 34 ++- .../handlers/s3_handler/connection_args.py | 21 ++ .../handlers/s3_handler/s3_handler.py | 215 +++++++++++++++--- 3 files changed, 238 insertions(+), 32 deletions(-) diff --git a/mindsdb/integrations/handlers/s3_handler/README.md b/mindsdb/integrations/handlers/s3_handler/README.md index bac14fe76e2..af4f6ca80a6 100644 --- a/mindsdb/integrations/handlers/s3_handler/README.md +++ b/mindsdb/integrations/handlers/s3_handler/README.md @@ -20,7 +20,10 @@ WITH parameters = { "aws_access_key_id": "AQAXEQK89OX07YS34OP", "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - "bucket": "my-bucket" + "bucket": "my-bucket", + "path_prefix": "rules/country=US/", + "list_cache_ttl_seconds": 300, + "include_metadata": false }; ``` @@ -37,6 +40,9 @@ Optional connection parameters include the following: * `aws_session_token`: The AWS session token that identifies the user or IAM role. This becomes necessary when using temporary security credentials. * `bucket`: The name of the Amazon S3 bucket. If not provided, all available buckets can be queried, however, this can affect performance, especially when listing all of the available objects. +* `path_prefix`: An optional S3 key prefix used to pre-filter objects returned by the `files` table. +* `list_cache_ttl_seconds`: Optional cache TTL, in seconds, for S3 object listings used by the `files` table. Defaults to 300 seconds. +* `include_metadata`: Whether to include custom S3 object metadata in the `files` table. Defaults to false. When enabled, the handler calls `HeadObject` for each listed file. ## Usage @@ -67,6 +73,28 @@ SELECT * FROM s3_datasource.files LIMIT 10 ``` +The `files` table includes standard object details returned by S3 listing operations: `path`, `name`, `extension`, `bucket`, `size`, `last_modified`, `etag`, and `storage_class`. These fields do not require per-object metadata requests. + +Use path filters to list a subset of objects. Prefix-like filters can be pushed down to S3 listing, and the cached listing can then be filtered by DuckDB: + +```sql +SELECT path, size, last_modified +FROM s3_datasource.files +WHERE path LIKE 'rules/country=US/state=CA/%' + AND path LIKE '%dimension=pay_transparency/%'; +``` + +Use exact path filters when you already know which documents to read. This avoids listing the bucket and directly checks the requested objects: + +```sql +SELECT path, content +FROM s3_datasource.files +WHERE path IN ( + 'rules/country=US/scope=federal/dimension=eeo_aap/rule_id=eeoc-eeo/version=2026-01-01/source.md', + 'rules/country=US/state=CA/scope=state/dimension=pay_transparency/rule_id=ca-pay-transparency/version=2026-01-01/source.md' +); +``` + The content of files can also be retrieved by explicitly requesting the `content` column. This column is empty by default to avoid unnecessary data transfer: ```sql @@ -74,6 +102,8 @@ SELECT path, content FROM s3_datasource.files LIMIT 10 ``` +Custom S3 object metadata can be returned in the `metadata` column when `include_metadata` is set to true. This is intended for inspection and enrichment workflows, not high-volume filtering, because S3 requires one `HeadObject` request per file to retrieve custom metadata. + This table will return all objects regardless of the file format, however, only the supported file formats mentioned above can be queried. @@ -100,4 +130,4 @@ This table will return all objects regardless of the file format, however, only * Incorrect: SELECT * FROM integration.travel/travel_data.csv * Incorrect: SELECT * FROM integration.'travel/travel_data.csv' * Correct: SELECT * FROM integration.\`travel/travel_data.csv\` - \ No newline at end of file + diff --git a/mindsdb/integrations/handlers/s3_handler/connection_args.py b/mindsdb/integrations/handlers/s3_handler/connection_args.py index b7ed9b69e19..37f3fdb4304 100644 --- a/mindsdb/integrations/handlers/s3_handler/connection_args.py +++ b/mindsdb/integrations/handlers/s3_handler/connection_args.py @@ -23,6 +23,24 @@ 'required': True, 'label': 'Amazon S3 Bucket' }, + path_prefix={ + 'type': ARG_TYPE.STR, + 'description': 'Optional S3 key prefix used to pre-filter objects returned by the files table.', + 'required': False, + 'label': 'S3 Path Prefix' + }, + include_metadata={ + 'type': ARG_TYPE.BOOL, + 'description': 'Whether to include custom S3 object metadata in the files table. Defaults to false and requires one HeadObject request per file when enabled.', + 'required': False, + 'label': 'Include Object Metadata' + }, + list_cache_ttl_seconds={ + 'type': ARG_TYPE.INT, + 'description': 'Time-to-live in seconds for cached S3 object listings used by the files table. Defaults to 300 seconds.', + 'required': False, + 'label': 'List Cache TTL Seconds' + }, region_name={ 'type': ARG_TYPE.STR, 'description': 'The AWS region to connect to. Default is `us-east-1`.', @@ -44,4 +62,7 @@ aws_session_token='FQoGZXIvYXdzEHcaDmJjJj...', region_name='us-east-2', bucket='my-bucket', + path_prefix='rules/country=US/', + include_metadata=False, + list_cache_ttl_seconds=300, ) diff --git a/mindsdb/integrations/handlers/s3_handler/s3_handler.py b/mindsdb/integrations/handlers/s3_handler/s3_handler.py index c643b50837b..8b696a9512a 100644 --- a/mindsdb/integrations/handlers/s3_handler/s3_handler.py +++ b/mindsdb/integrations/handlers/s3_handler/s3_handler.py @@ -1,4 +1,6 @@ +import json import os +import time from typing import List from contextlib import contextmanager from io import BytesIO @@ -34,7 +36,10 @@ class ListFilesTable(APIResource): def list( self, targets: List[str] = None, conditions: List[FilterCondition] = None, limit: int = None, *args, **kwargs ) -> pd.DataFrame: + conditions = conditions or [] buckets = None + exact_paths = None + prefix = self.handler.path_prefix for condition in conditions: if condition.column == "bucket": if condition.op == FilterOperator.IN: @@ -42,27 +47,43 @@ def list( elif condition.op == FilterOperator.EQUAL: buckets = [condition.value] condition.applied = True + elif condition.column == "path": + if condition.op == FilterOperator.IN: + exact_paths = condition.value + condition.applied = True + elif condition.op == FilterOperator.EQUAL: + exact_paths = [condition.value] + condition.applied = True + elif condition.op == FilterOperator.LIKE and isinstance(condition.value, str): + prefix = self.handler.combine_prefixes(prefix, self.handler.get_prefix_from_like(condition.value)) data = [] - for obj in self.handler.get_objects(limit=limit, buckets=buckets): - path = obj["Key"] - path = path.replace("`", "") - item = { - "path": path, - "bucket": obj["Bucket"], - "name": path[path.rfind("/") + 1 :], - "extension": path[path.rfind(".") + 1 :], - } - - if targets and "public_url" in targets: - item["public_url"] = self.handler.generate_sas_url(path, obj["Bucket"]) + if exact_paths is not None: + objects = self.handler.get_objects_by_paths(exact_paths, buckets=buckets) + else: + has_unapplied_conditions = any(not condition.applied for condition in conditions) + list_limit = None if has_unapplied_conditions else limit + objects = self.handler.get_objects(limit=list_limit, buckets=buckets, prefix=prefix) - data.append(item) + for obj in objects: + data.append(self.handler.object_to_file_row(obj, targets=targets)) return pd.DataFrame(data=data, columns=self.get_columns()) def get_columns(self) -> List[str]: - return ["path", "name", "extension", "bucket", "content", "public_url"] + return [ + "path", + "name", + "extension", + "bucket", + "size", + "last_modified", + "etag", + "storage_class", + "content", + "public_url", + "metadata", + ] class FileTable(APIResource): @@ -102,6 +123,10 @@ def __init__(self, name: Text, connection_data: Optional[Dict], **kwargs): self.is_connected = False self.cache_thread_safe = True self.bucket = self.connection_data.get("bucket") + self.path_prefix = self._normalize_prefix(self.connection_data.get("path_prefix")) + self.include_metadata = self._coerce_bool(self.connection_data.get("include_metadata", False)) + self.list_cache_ttl_seconds = int(self.connection_data.get("list_cache_ttl_seconds") or 300) + self._objects_cache = {} self._regions = {} self._files_table = ListFilesTable(self) @@ -276,6 +301,81 @@ def _get_bucket(self, key): ar = key.split("/") return ar[0], "/".join(ar[1:]) + @staticmethod + def _coerce_bool(value) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + return bool(value) + + @staticmethod + def _normalize_prefix(prefix): + if prefix is None: + return None + prefix = str(prefix).strip() + if prefix == "": + return None + return prefix.lstrip("/") + + @staticmethod + def get_prefix_from_like(value: str): + """Extract the S3-prefix-safe part of a SQL LIKE pattern.""" + wildcard_positions = [pos for pos in (value.find("%"), value.find("_")) if pos >= 0] + if not wildcard_positions: + return value + wildcard_pos = min(wildcard_positions) + if wildcard_pos == 0: + return None + return value[:wildcard_pos] + + @classmethod + def combine_prefixes(cls, base_prefix, filter_prefix): + base_prefix = cls._normalize_prefix(base_prefix) + filter_prefix = cls._normalize_prefix(filter_prefix) + if not base_prefix: + return filter_prefix + if not filter_prefix: + return base_prefix + if filter_prefix.startswith(base_prefix): + return filter_prefix + if base_prefix.startswith(filter_prefix): + return base_prefix + return base_prefix + + def _display_path(self, bucket: str, key: str) -> str: + if self.bucket is None: + return f"{bucket}/{key}" + return key + + def object_to_file_row(self, obj: dict, targets: List[str] = None) -> dict: + path = obj["Key"].replace("`", "") + name = path[path.rfind("/") + 1 :] + extension = name[name.rfind(".") + 1 :] if "." in name else "" + raw_key = obj.get("_S3Key", obj["Key"]) + + item = { + "path": path, + "bucket": obj["Bucket"], + "name": name, + "extension": extension, + "size": obj.get("Size"), + "last_modified": obj.get("LastModified"), + "etag": (obj.get("ETag") or "").strip('"'), + "storage_class": obj.get("StorageClass", "STANDARD"), + } + + if targets and "public_url" in targets: + item["public_url"] = self.generate_sas_url(raw_key, obj["Bucket"]) + + if self.include_metadata and (not targets or "metadata" in targets): + metadata = obj.get("_Metadata") + if metadata is None: + metadata = self.get_object_metadata(obj["Bucket"], raw_key) + item["metadata"] = json.dumps(metadata) + + return item + def read_as_table(self, key, chunk_size: int = None, chunk_overlap: int = None) -> pd.DataFrame: """ Read object as dataframe. Uses duckdb for structured files, FileReader for text files @@ -413,13 +513,74 @@ def native_query(self, query: str) -> Response: query_ast = parse_sql(query) return self.query(query_ast) - def get_objects(self, limit=None, buckets=None) -> List[dict]: + def get_objects_by_paths(self, paths, buckets=None) -> List[dict]: + client = self.connect() + objects = [] + for path in paths: + bucket, key = self._get_bucket(path) + if buckets is not None and bucket not in buckets: + continue + + try: + obj = client.head_object(Bucket=bucket, Key=key) + except ClientError as e: + logger.error(f"Error querying the file {key} in the bucket {bucket}, {e}!") + continue + + objects.append( + { + "Bucket": bucket, + "Key": self._display_path(bucket, key), + "_S3Key": key, + "Size": obj.get("ContentLength"), + "LastModified": obj.get("LastModified"), + "ETag": obj.get("ETag"), + "StorageClass": obj.get("StorageClass", "STANDARD"), + "_Metadata": obj.get("Metadata"), + } + ) + + return objects + + def get_object_metadata(self, bucket: str, key: str) -> dict: + client = self.connect() + return client.head_object(Bucket=bucket, Key=key).get("Metadata", {}) + + def _get_cached_bucket_objects(self, bucket: str, prefix=None) -> List[dict]: + client = self.connect() + prefix = self._normalize_prefix(prefix) + cache_key = (bucket, prefix or "") + cache_entry = self._objects_cache.get(cache_key) + now = time.time() + if ( + cache_entry is not None + and self.list_cache_ttl_seconds > 0 + and now - cache_entry["created_at"] < self.list_cache_ttl_seconds + ): + return cache_entry["objects"] + + paginator = client.get_paginator("list_objects_v2") + pagination_kwargs = {"Bucket": bucket} + if prefix: + pagination_kwargs["Prefix"] = prefix + + objects = [] + for page in paginator.paginate(**pagination_kwargs): + for obj in page.get("Contents", []): + if obj.get("StorageClass", "STANDARD") != "STANDARD": + continue + objects.append(dict(obj)) + + if self.list_cache_ttl_seconds > 0: + self._objects_cache[cache_key] = {"created_at": now, "objects": objects} + + return objects + + def get_objects(self, limit=None, buckets=None, prefix=None) -> List[dict]: client = self.connect() if self.bucket is not None: - add_bucket_to_name = False scan_buckets = [self.bucket] else: - add_bucket_to_name = True scan_buckets = [b["Name"] for b in client.list_buckets()["Buckets"]] objects = [] @@ -427,21 +588,15 @@ def get_objects(self, limit=None, buckets=None) -> List[dict]: if buckets is not None and bucket not in buckets: continue - resp = client.list_objects_v2(Bucket=bucket) - if "Contents" not in resp: - continue - - for obj in resp["Contents"]: - if obj.get("StorageClass", "STANDARD") != "STANDARD": - continue - + for obj in self._get_cached_bucket_objects(bucket, prefix=prefix): + obj = dict(obj) + raw_key = obj["Key"] obj["Bucket"] = bucket - if add_bucket_to_name: - # bucket is part of the name - obj["Key"] = f"{bucket}/{obj['Key']}" + obj["_S3Key"] = raw_key + obj["Key"] = self._display_path(bucket, raw_key) objects.append(obj) - if limit is not None and len(objects) >= limit: - break + if limit is not None and len(objects) >= limit: + return objects[: int(limit)] return objects From e7dc72b233687942970e76ecc029dade1be5610b Mon Sep 17 00:00:00 2001 From: Gabriel Bressan Date: Sat, 6 Jun 2026 14:34:05 -0400 Subject: [PATCH 2/2] fix region name --- .../handlers/s3_handler/connection_args.py | 7 +++++++ .../handlers/s3_handler/s3_handler.py | 21 +++++++++++-------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/mindsdb/integrations/handlers/s3_handler/connection_args.py b/mindsdb/integrations/handlers/s3_handler/connection_args.py index 37f3fdb4304..873296a0ed0 100644 --- a/mindsdb/integrations/handlers/s3_handler/connection_args.py +++ b/mindsdb/integrations/handlers/s3_handler/connection_args.py @@ -29,6 +29,12 @@ 'required': False, 'label': 'S3 Path Prefix' }, + prefix={ + 'type': ARG_TYPE.STR, + 'description': 'Alias for path_prefix. Optional S3 key prefix used to pre-filter objects returned by the files table.', + 'required': False, + 'label': 'S3 Prefix' + }, include_metadata={ 'type': ARG_TYPE.BOOL, 'description': 'Whether to include custom S3 object metadata in the files table. Defaults to false and requires one HeadObject request per file when enabled.', @@ -63,6 +69,7 @@ region_name='us-east-2', bucket='my-bucket', path_prefix='rules/country=US/', + prefix='rules/country=US/', include_metadata=False, list_cache_ttl_seconds=300, ) diff --git a/mindsdb/integrations/handlers/s3_handler/s3_handler.py b/mindsdb/integrations/handlers/s3_handler/s3_handler.py index 8b696a9512a..92ce7cbd1cf 100644 --- a/mindsdb/integrations/handlers/s3_handler/s3_handler.py +++ b/mindsdb/integrations/handlers/s3_handler/s3_handler.py @@ -123,7 +123,7 @@ def __init__(self, name: Text, connection_data: Optional[Dict], **kwargs): self.is_connected = False self.cache_thread_safe = True self.bucket = self.connection_data.get("bucket") - self.path_prefix = self._normalize_prefix(self.connection_data.get("path_prefix")) + self.path_prefix = self._normalize_prefix(self.connection_data.get("path_prefix") or self.connection_data.get("prefix")) self.include_metadata = self._coerce_bool(self.connection_data.get("include_metadata", False)) self.list_cache_ttl_seconds = int(self.connection_data.get("list_cache_ttl_seconds") or 300) self._objects_cache = {} @@ -173,14 +173,17 @@ def _connect_duckdb(self, bucket): pass duckdb_conn.execute("LOAD httpfs") - # detect region for bucket - if bucket not in self._regions: - client = self.connect() - bucket_region = client.get_bucket_location(Bucket=bucket).get("LocationConstraint") - # LocationConstraint pode ser None para us-east-1 - self._regions[bucket] = bucket_region or "us-east-1" - - region = self.connection_data.get("region_name", self._regions.get(bucket, "us-east-1")) + # detect region for bucket only when region_name not provided + configured_region = self.connection_data.get("region_name") + if configured_region: + region = configured_region + else: + if bucket not in self._regions: + client = self.connect() + bucket_region = client.get_bucket_location(Bucket=bucket).get("LocationConstraint") + # LocationConstraint pode ser None para us-east-1 + self._regions[bucket] = bucket_region or "us-east-1" + region = self._regions.get(bucket, "us-east-1") region = region or "us-east-1" credentials = None