orioledb_s3_loader 'endpoint' arg now can be same as orioledb.s3_host

JoshInnis · Dec 11, 2023 · 3b6c7ce · 3b6c7ce
1 parent 0315775
commit 3b6c7ce
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/doc/usage.md b/doc/usage.md
@@ -218,7 +218,7 @@ Run the script with the same parameters as from your S3 Postgres cluster config:
 * `AWS_ACCESS_KEY_ID` - same as `orioledb.s3_accesskey`
 * `AWS_SECRET_ACCESS_KEY` - same as `orioledb.s3_secretkey`
 * `AWS_DEFAULT_REGION` - same as `orioledb.s3_region`
-* `--endpoint` - prefix from `orioledb.s3_host` (full URL with `https://` but without a bucket name) Should be either `--endpoint=https://s3-accelerate.amazonaws.com` or `--endpoint=https://s3.amazonaws.com`
+* `--endpoint` - same as `orioledb.s3_host` (full URL with `https://` prefix) E.g `--endpoint=https://mybucket.s3-accelerate.amazonaws.com` or `--endpoint=https://mybucket.s3.amazonaws.com`
 * `--bucket-name` - S3 bucket name from `orioledb.s3_host` E.g `--bucket-name=mybucket`
 * `--data-dir` - destination directory on the local machine you want to write data to. E.g. `--data-dir=mydata/`
 * `--verbose` - optionally print extended info.

diff --git a/orioledb_s3_loader.py b/orioledb_s3_loader.py
@@ -12,6 +12,7 @@
 from botocore.exceptions import ClientError
 from threading import Event
 from typing import Callable
+from urllib.parse import urlparse
 
 class OrioledbS3ObjectLoader:
 	def parse_args(self):
@@ -47,6 +48,10 @@ def parse_args(self):
 		else:
 			verify = None
 
+		parsed_url = urlparse(args.endpoint)
+		bucket = parsed_url.netloc.split('.')[0]
+		if bucket == args.bucket_name:
+			args.endpoint = f"{parsed_url.scheme}://{'.'.join(parsed_url.netloc.split('.')[1:])}"
 		self.s3 = boto3.client("s3", endpoint_url=args.endpoint,
 							   verify=verify)
 		self._error_occurred = Event()