Fix iceberg naming issues

Author: joocer

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1862
+__build__ = 1871
 __author__ = "@joocer"
-__version__ = "0.26.2-beta.1862"
+__version__ = "0.26.2-beta.1871"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/connectors/__init__.py

@@ -181,6 +181,11 @@ def register_store(prefix, connector, *, remove_prefix: bool = False, **kwargs):
         # uninstantiated classes aren't a type
         raise ValueError("connectors registered with `register_store` must be uninstantiated.")
 
+    if connector.__name__ == "IcebergConnector" and not remove_prefix:
+        raise ValueError(
+            "IcebergConnector requires remove_prefix=True so catalog prefixes don't leak into table names."
+        )
+
     # Store connector class directly (not as a string)
     _storage_prefixes[prefix] = {
         "connector": connector,  # type: ignore
@@ -289,6 +294,9 @@ def connector_factory(dataset, statistics, **config):
         dataset = dataset[len(prefix) :]
         if dataset.startswith(".") or dataset.startswith("//"):
             dataset = dataset[1:] if dataset.startswith(".") else dataset[2:]
+        if connector.__name__ == "IcebergConnector" and dataset and "." not in dataset:
+            # Default to a namespace matching the prefix when one isn't provided
+            dataset = f"{prefix}.{dataset}"
 
     return connector(dataset=dataset, statistics=statistics, **connector_entry)
 

opteryx/connectors/aws_s3_connector.py

@@ -176,7 +176,7 @@ def get_dataset_schema(self) -> RelationSchema:
         self.schema = next(self.read_dataset(just_schema=True), None)
 
         if self.schema is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         return self.schema
 

opteryx/connectors/disk_connector.py

@@ -376,6 +376,6 @@ def get_dataset_schema(self) -> RelationSchema:
         if self.schema is None:
             if os.path.isdir(self.dataset):
                 raise EmptyDatasetError(dataset=self.dataset.replace(OS_SEP, "."))
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         return self.schema

opteryx/connectors/file_connector.py

@@ -135,7 +135,7 @@ def __init__(self, *args, **kwargs):
 
         if ".." in self.dataset or self.dataset[0] in ("\\", "/", "~"):
             # Don't find any datasets which look like path traversal
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         # Check if dataset contains wildcards
         self.has_wildcards = any(char in self.dataset for char in ["*", "?", "["])
@@ -144,7 +144,7 @@ def __init__(self, *args, **kwargs):
             # Expand wildcards to get list of files
             self.files = self._expand_wildcards(self.dataset)
             if not self.files:
-                raise DatasetNotFoundError(dataset=self.dataset)
+                raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
             # Use the first file to determine the decoder
             self.decoder = get_decoder(self.files[0])
         else:
@@ -168,7 +168,7 @@ def _expand_wildcards(self, pattern: str) -> List[str]:
         """
         # Additional path traversal check after expansion
         if ".." in pattern:
-            raise DatasetNotFoundError(dataset=pattern)
+            raise DatasetNotFoundError(dataset=pattern, connector=self.__type__)
 
         # Use glob to expand the pattern
         matched_files = glob.glob(pattern, recursive=False)

opteryx/connectors/gcp_cloudstorage_connector.py

@@ -318,7 +318,7 @@ def get_dataset_schema(self) -> RelationSchema:
             break
 
         if self.schema is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         # if we have more than one blob we need to estimate the row count
         if self.schema.row_count_metric and number_of_blobs > 1:

opteryx/connectors/gcp_firestore_connector.py

@@ -126,7 +126,7 @@ def get_dataset_schema(self) -> RelationSchema:
         record = next(self.read_dataset(chunk_size=10), None)
 
         if record is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         arrow_schema = record.schema
 

opteryx/connectors/iceberg_connector.py

@@ -148,12 +148,11 @@ def __init__(self, *args, catalog=None, io=DiskConnector, **kwargs):
 
         import pyiceberg
 
-        self.dataset = self.dataset.lower()
         try:
             self.table = catalog.load_table(self.dataset)
             self.io_connector = io(**kwargs)
         except pyiceberg.exceptions.NoSuchTableError:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__) from None
 
     def get_dataset_schema(self) -> RelationSchema:
         iceberg_schema = self.table.schema()

opteryx/connectors/mongodb_connector.py

@@ -76,7 +76,7 @@ def get_dataset_schema(self) -> RelationSchema:
         record = next(self.read_dataset(chunk_size=25), None)
 
         if record is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         arrow_schema = record.schema
 

opteryx/connectors/virtual_data.py

@@ -93,7 +93,9 @@ def read_dataset(self, columns: list = None, **kwargs) -> "DatasetReader":
     def get_dataset_schema(self) -> RelationSchema:
         if self.dataset not in WELL_KNOWN_DATASETS:
             suggestion = suggest(self.dataset)
-            raise DatasetNotFoundError(suggestion=suggestion, dataset=self.dataset)
+            raise DatasetNotFoundError(
+                suggestion=suggestion, dataset=self.dataset, connector=self.__type__
+            )
         data_provider, _ = _load_provider(self.dataset)
         self.relation_statistics = data_provider.statistics()
         return data_provider.schema()
@@ -138,6 +140,8 @@ def __next__(self) -> "pyarrow.Table":
         data_provider, _ = _load_provider(self.dataset_name)
         if data_provider is None:
             suggestion = suggest(self.dataset_name.lower())
-            raise DatasetNotFoundError(suggestion=suggestion, dataset=self.dataset_name)
+            raise DatasetNotFoundError(
+                suggestion=suggestion, dataset=self.dataset_name, connector="SAMPLE"
+            )
         table = data_provider.read(self.date, self.variables)
         return arrow.post_read_projector(table, self.columns)