feat: add strict type checks for column_name formatting and raise errors on invalid inputs

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 Note: version releases in the 0.x.y range may include both bug fixes and new features, not strictly limited to patches.
 
+## 0.2.13
+- feat: rename parameter `secret` to `secret_alias` in `create_token` method for clarity
+- feat: add strict type checks for column_name formatting and raise errors on invalid inputs
+
 ## 0.2.12
 - feat: add validation for reserved keywords in graph schema
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tigergraphx"
-version = "0.2.12"
+version = "0.2.13"
 description = "TigerGraphX is a high-level Python library offering a unified, Python-native interface for graph databases, advanced analytics, and GraphRAG workflows. Combining the simplicity of NetworkX with the advanced capabilities of TigerGraph, including tgCloud, it empowers Python developers to harness the power of graphs without the need to learn query languages like Cypher or GSQL."
 authors = ["Xuanlei Lin <xuanlei.lin@tigergraph.com>"]
 license = "MIT"

tests/integration/core/tigergraph_database/tigergraph_database_test.py

@@ -116,10 +116,10 @@ def test_tokens(self):
             assert "has been created for user" in result
 
             match = re.search(r"The secret:\s*([a-z0-9]+)\s*has been created", result)
-            secret = match.group(1) if match else None
-            assert secret, "Secret not extracted from response"
+            secret_alias = match.group(1) if match else None
+            assert secret_alias, "Secret not extracted from response"
 
-            token = self.db.create_token(secret)
+            token = self.db.create_token(secret_alias)
 
             # Check token is a non-empty string and matches JWT pattern (3 parts separated by dots)
             assert isinstance(token, str)

tests/unit/core/managers/data_manager_test.py

@@ -40,3 +40,29 @@ def test_load_data_failure(self):
         # Assert that RuntimeError is raised on failure
         with pytest.raises(RuntimeError):
             self.data_manager.load_data(loading_job_config)
+
+    def test_format_none(self):
+        assert DataManager._format_column_name(None) == "_"
+
+    def test_format_int(self):
+        assert DataManager._format_column_name(3) == "$3"
+
+    def test_format_str_simple(self):
+        assert DataManager._format_column_name("person") == '$"person"'
+
+    def test_format_str_with_space(self):
+        assert DataManager._format_column_name("person name") == '$"person name"'
+
+    def test_format_function_dict(self):
+        assert (
+            DataManager._format_column_name({"func": "gsql_uuid_v4()"})
+            == "gsql_uuid_v4()"
+        )
+
+    def test_format_function_dict_invalid_type(self):
+        with pytest.raises(TypeError):
+            DataManager._format_column_name({"func": 123})
+
+    def test_format_unsupported_type(self):
+        with pytest.raises(TypeError):
+            DataManager._format_column_name([1, 2, 3])  # pyright: ignore

tigergraphx/config/graph_db/loading_job_config.py

@@ -46,7 +46,7 @@ class NodeMappingConfig(BaseConfig):
     """
 
     target_name: str = Field(description="The name of the target node type.")
-    attribute_column_mappings: Dict[str, str | int] = Field(
+    attribute_column_mappings: Dict[str, str | int | Dict] = Field(
         default={}, description="Mapping file columns to node attributes."
     )
 

tigergraphx/core/managers/data_manager.py

@@ -22,7 +22,9 @@ class DataManager(BaseManager):
     def __init__(self, context: GraphContext):
         super().__init__(context)
 
-    def load_data(self, loading_job_config: LoadingJobConfig | Dict | str | Path) -> str:
+    def load_data(
+        self, loading_job_config: LoadingJobConfig | Dict | str | Path
+    ) -> str:
         loading_job_config = LoadingJobConfig.ensure_config(loading_job_config)
         logger.info(
             f"Initiating data load for job: {loading_job_config.loading_job_name}...",
@@ -229,13 +231,23 @@ def _create_gsql_load_data(
         return gsql_script.strip()
 
     @staticmethod
-    def _format_column_name(column_name: str | int | None) -> str:
+    def _format_column_name(column_name: str | int | Dict | None) -> str:
         """Format column names as $number, $"variable", or _ for empty names."""
         if column_name is None:
             return "_"
         if isinstance(column_name, int):
             return f"${column_name}"
-        if isinstance(column_name, str) and column_name.isidentifier():
+        if isinstance(column_name, str):
             return f'$"{column_name}"'
-        # Return the original name as string if it doesn't match any of the specified formats
-        return str(column_name)
+        if isinstance(column_name, dict) and "func" in column_name:
+            func = column_name["func"]
+            if not isinstance(func, str):
+                raise TypeError(
+                    f"Invalid function reference: {func!r}. Expected a string."
+                )
+            return func
+
+        # Unknown type → raise explicit error
+        raise TypeError(
+            f"Unsupported column name type: {type(column_name).__name__}, value={column_name!r}"
+        )

tigergraphx/core/tigergraph_api/api/security_api.py

@@ -13,12 +13,12 @@
 class SecurityAPI(BaseAPI):
     def create_token(
         self,
-        secret: str,
+        secret_alias: str,
         graph_name: Optional[str] = None,
         lifetime_seconds: Optional[int] = None,
     ) -> str:
         payload: Dict[str, Any] = {
-            "secret": secret,
+            "secret": secret_alias,
         }
         if graph_name:
             payload["graph"] = graph_name

tigergraphx/core/tigergraph_api/tigergraph_api.py

@@ -149,21 +149,21 @@ def gsql(self, command: str) -> str:
     # ------------------------------ Security ------------------------------
     def create_token(
         self,
-        secret: str,
+        secret_alias: str,
         graph_name: Optional[str] = None,
         lifetime_seconds: Optional[int] = None,
     ) -> str:
         """Create an auth token using a secret.
 
         Args:
-            secret: The secret alias to use for token generation.
+            secret_alias: The secret alias to use for token generation.
             graph_name: The name of the graph to scope the token.
             lifetime_seconds: Duration in seconds before the token expires.
 
         Returns:
             The generated authentication token as a string.
         """
-        return self._security_api.create_token(secret, graph_name, lifetime_seconds)
+        return self._security_api.create_token(secret_alias, graph_name, lifetime_seconds)
 
     def drop_token(
         self,

tigergraphx/core/tigergraph_database.py

@@ -126,21 +126,21 @@ def drop_secret(self, alias: str) -> str:
 
     def create_token(
         self,
-        secret: str,
+        secret_alias: str,
         graph_name: Optional[str] = None,
         lifetime_seconds: Optional[int] = None,
     ) -> str:
         """Create an auth token using a secret.
 
         Args:
-            secret: The secret alias to use for token generation.
+            secret_alias: The secret alias to use for token generation.
             graph_name: The name of the graph to scope the token.
             lifetime_seconds: Duration in seconds before the token expires.
 
         Returns:
             The generated authentication token as a string.
         """
-        return self._tigergraph_api.create_token(secret, graph_name, lifetime_seconds)
+        return self._tigergraph_api.create_token(secret_alias, graph_name, lifetime_seconds)
 
     def drop_token(
         self,