class TigerGraphConnection(Connection):
"""
TigerGraph database connection implementation.
Key conceptual differences from ArangoDB:
1. TigerGraph uses GSQL (Graph Query Language) instead of AQL
2. Schema must be defined explicitly before data insertion
3. No automatic vertex/edge class creation - vertices and edges must be pre-defined
4. Different query syntax and execution model
5. Token-based authentication recommended for TigerGraph 4+
Authentication (recommended for TG 4+):
For best results, provide BOTH username/password AND secret:
- username/password: Required for initial connection and GSQL operations
- secret: Generates token that works for both GSQL and REST API operations
Token-based authentication using secrets is the most robust and recommended
approach for TigerGraph 4+. The connection will:
1. Use username/password for initial connection
2. Generate a token from the secret
3. Use the token for both GSQL operations (via REST API) and REST API calls
Example:
>>> config = TigergraphConfig(
... uri="http://localhost:14240",
... username="tigergraph", # Required for initial connection
... password="tigergraph", # Required for initial connection
... secret="your_secret_here", # Generates token for GSQL + REST API
... database="my_graph"
... )
>>> conn = TigerGraphConnection(config)
Port Configuration for TigerGraph 4+:
TigerGraph 4.1+ uses port 14240 (GSQL server) as the primary interface.
Port 9000 (REST++) is for internal use only in TG 4.1+.
Standard ports:
- Port 14240: GSQL server (primary interface for all API requests)
- Port 9000: REST++ (internal-only in TG 4.1+)
For custom Docker deployments with port mapping, ports are configured via
environment variables (e.g., TG_WEB, TG_REST) and loaded automatically
when using TigergraphConfig.from_docker_env().
Version Compatibility:
- All TigerGraph versions use /restpp prefix for REST++ endpoints
- Version is auto-detected, or can be manually specified in config
"""
flavor = DBType.TIGERGRAPH
def __init__(self, config: TigergraphConfig):
super().__init__()
self.config = config
self.ssl_verify = getattr(config, "ssl_verify", True)
# Keep legacy and unified graph namespace fields aligned.
if self.config.database is None and self.config.schema_name is not None:
self.config.database = self.config.schema_name
elif self.config.schema_name is None and self.config.database is not None:
self.config.schema_name = self.config.database
# Store connection configuration (no longer using pyTigerGraph)
# For TigerGraph 4+, both ports typically route through the GSQL server
# Port 9000 (REST++) is internal-only in TG 4.1+
configured_graph = self._configured_graph_name()
self.graphname: str = (
configured_graph if configured_graph is not None else "DefaultGraph"
)
self._installed_clear_data_queries: dict[str, set[str]] = {}
# Initialize URLs (ports come from config, no hardcoded defaults)
# Set GSQL URL first as it's needed for token generation
# For TigerGraph 4+, gs_port is the primary port (extracted from URI if not explicitly set)
# Fall back to port from URI if gs_port is not set
gs_port: int | str | None = config.gs_port
if gs_port is None:
# Try to get port from URI
uri_port = config.port
if uri_port:
try:
gs_port = int(uri_port)
logger.debug(f"Using port {gs_port} from URI for GSQL endpoint")
except (ValueError, TypeError):
pass
if gs_port is None:
raise ValueError(
"gs_port or URI with port must be set in TigergraphConfig. "
"Standard ports: 14240 (GSQL), 9000 (REST++)."
)
self.gsql_url = f"{config.url_without_port}:{gs_port}"
# Detect TigerGraph version for compatibility (needed before token generation)
self.tg_version: str | None = None
self._use_restpp_prefix = False # Default for 4.2.2+
# Check if version is manually configured first
if hasattr(config, "version") and config.version:
version_str = config.version
logger.info(f"Using manually configured TigerGraph version: {version_str}")
else:
# Auto-detect version using REST API
try:
version_str = self._get_version()
except Exception as e:
logger.warning(
f"Failed to detect TigerGraph version: {e}. "
f"Defaulting to 4.2.2+ behavior (no /restpp prefix)"
)
version_str = None
# Parse version string if we have one
if version_str:
# Extract version from strings like "release_4.2.2_09-29-2025" or "4.2.1" or "v4.2.1"
import re
version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", version_str)
if version_match:
major = int(version_match.group(1))
minor = int(version_match.group(2))
patch = int(version_match.group(3))
self.tg_version = f"{major}.{minor}.{patch}"
# All TigerGraph versions use /restpp prefix for REST++ endpoints
# Even 4.2.2+ requires /restpp prefix (despite some documentation suggesting otherwise)
self._use_restpp_prefix = True
logger.info(
f"TigerGraph version {self.tg_version} detected, "
f"using /restpp prefix for REST API"
)
else:
logger.warning(
f"Could not extract version number from '{version_str}'. "
f"Defaulting to using /restpp prefix for REST API"
)
self._use_restpp_prefix = True
# Store base URLs for REST++ and GSQL endpoints
# For TigerGraph 4.1+, REST++ endpoints use the GSQL port with /restpp prefix
# Port 9000 is internal-only in TG 4.1+, so we use the same port as GSQL
# Use the GSQL port we already determined to ensure consistency
base_url = f"{config.url_without_port}:{gs_port}"
# Always use /restpp prefix for REST++ endpoints (required for all TG versions)
self.restpp_url = f"{base_url}/restpp"
# Get authentication token if secret is provided
# Token-based auth is the recommended approach for TigerGraph 4+
# IMPORTANT: You should provide BOTH username/password AND secret:
# - username/password: Used for initial connection and GSQL operations
# - secret: Generates token that works for both GSQL and REST API operations
# Use graph-specific token (is_global=False) for better security
self.api_token: str | None = None
if config.secret:
try:
token, expiration = self._get_token_from_secret(
config.secret,
self.graphname, # Pass graph name for graph-specific token
)
self.api_token = token
if expiration:
logger.info(
f"Successfully obtained API token for graph '{self.graphname}' "
f"(expires: {expiration})"
)
else:
logger.info(
f"Successfully obtained API token for graph '{self.graphname}'"
)
except Exception as e:
# Log and fall back to username/password authentication
logger.warning(f"Failed to get authentication token: {e}")
logger.warning("Falling back to username/password authentication")
logger.warning(
"Note: For best results, provide both username/password AND secret. "
"Username/password is used for GSQL operations, secret generates token for REST API."
)
def _configured_graph_name(self) -> str | None:
"""Return the configured TigerGraph graph name from either config field."""
return self.config.database or self.config.schema_name
def _require_configured_graph_name(self) -> str:
"""Return graph name or raise if neither database nor schema_name is set."""
graph_name = self._configured_graph_name()
if not graph_name:
raise ValueError(
"Graph name must be configured via config.database or config.schema_name"
)
return graph_name
def _get_auth_headers(self, use_basic_auth: bool = False) -> dict[str, str]:
"""Get authentication headers for REST API calls.
Args:
use_basic_auth: If True, always use Basic Auth (required for GSQL endpoints).
If False, prioritize token-based auth for REST++ endpoints.
Prioritizes token-based authentication over Basic Auth for REST++ endpoints:
1. If API token is available (from secret), use Bearer token (recommended for TG 4+)
2. Otherwise, fall back to HTTP Basic Auth with username/password
For GSQL endpoints, always use Basic Auth as they don't support Bearer tokens.
Returns:
Dictionary with Authorization header
"""
headers = {}
# GSQL endpoints require Basic Auth, not Bearer tokens
if use_basic_auth or not self.api_token:
# Use default username "tigergraph" if username is None but password is set
username = self.config.username if self.config.username else "tigergraph"
password = self.config.password
if password:
import base64
credentials = f"{username}:{password}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
headers["Authorization"] = f"Basic {encoded_credentials}"
else:
logger.warning(
f"No password configured for Basic Auth. "
f"Username: {username}, Password: {password}"
)
else:
# Use Bearer token for REST++ endpoints
headers["Authorization"] = f"Bearer {self.api_token}"
return headers
def _get_token_from_secret(
self, secret: str, graph_name: str | None = None, lifetime: int = 3600 * 24 * 30
) -> tuple[str, str | None]:
"""
Generate authentication token from secret using TigerGraph REST API.
Implements robust token generation with fallback logic for different TG 4.x versions:
- TigerGraph 4.2.2+: POST /gsql/v1/tokens (lifetime in milliseconds)
- TigerGraph 4.0-4.2.1: POST /gsql/v1/auth/token (lifetime in seconds)
Based on pyTigerGraph's token generation mechanism with version-specific endpoint handling.
Args:
secret: Secret string created via CREATE SECRET in GSQL
graph_name: Name of the graph (None for global token)
lifetime: Token lifetime in seconds (default: 30 days)
Returns:
Tuple of (token, expiration_timestamp) or (token, None) if expiration not provided
Raises:
RuntimeError: If token generation fails after all retry attempts
"""
auth_headers = self._get_auth_headers(use_basic_auth=True)
headers = {
"Content-Type": "application/json",
**auth_headers,
}
# Determine which endpoint to try based on version
# For TG 4.2.2+, use /gsql/v1/tokens (lifetime in milliseconds)
# For TG 4.0-4.2.1, use /gsql/v1/auth/token (lifetime in seconds)
use_new_endpoint = False
if self.tg_version:
import re
version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", self.tg_version)
if version_match:
major = int(version_match.group(1))
minor = int(version_match.group(2))
patch = int(version_match.group(3))
# Use new endpoint for 4.2.2+
use_new_endpoint = (major, minor, patch) >= (4, 2, 2)
# Try endpoints in order: new endpoint first (if version >= 4.2.2), then fallback
endpoints_to_try = []
if use_new_endpoint:
# Try new endpoint first for 4.2.2+
endpoints_to_try.append(
(
f"{self.gsql_url}/gsql/v1/tokens",
{
"secret": secret,
"graph": graph_name,
"lifetime": lifetime * 1000, # Convert to milliseconds
},
True, # lifetime in milliseconds
)
)
# Fallback to old endpoint if new one fails
endpoints_to_try.append(
(
f"{self.gsql_url}/gsql/v1/auth/token",
{
"secret": secret,
"graph": graph_name,
"lifetime": lifetime, # In seconds
},
False, # lifetime in seconds
)
)
else:
# For older versions or unknown version, try old endpoint first
endpoints_to_try.append(
(
f"{self.gsql_url}/gsql/v1/auth/token",
{
"secret": secret,
"graph": graph_name,
"lifetime": lifetime, # In seconds
},
False, # lifetime in seconds
)
)
# Fallback to new endpoint (in case version detection was wrong)
endpoints_to_try.append(
(
f"{self.gsql_url}/gsql/v1/tokens",
{
"secret": secret,
"graph": graph_name,
"lifetime": lifetime * 1000, # Convert to milliseconds
},
True, # lifetime in milliseconds
)
)
last_error: Exception | None = None
all_404_errors = True # Track if all failures were 404 errors
for url, payload, _is_milliseconds in endpoints_to_try:
try:
# Remove None values from payload
clean_payload = {k: v for k, v in payload.items() if v is not None}
response = requests.post(
url,
headers=headers,
json=clean_payload, # Use json parameter instead of data
timeout=30,
verify=self.ssl_verify,
)
# Check for 404 - might indicate wrong endpoint or port issue
if response.status_code == 404:
# Try port fallback (similar to pyTigerGraph's _req method)
# If using wrong port, try GSQL port
if (
"/gsql" in url
and self.config.port is not None
and self.config.gs_port is not None
and self.config.port != self.config.gs_port
):
logger.debug(f"404 on {url}, trying GSQL port fallback...")
# Replace port in URL with GSQL port
fallback_url = url.replace(
f":{self.config.port}", f":{self.config.gs_port}"
)
try:
response = requests.post(
fallback_url,
headers=headers,
json=clean_payload,
timeout=30,
verify=self.ssl_verify,
)
if response.status_code == 200:
url = fallback_url # Update URL for logging
except Exception:
pass # Continue to next endpoint
response.raise_for_status()
result = response.json()
# Parse response (both endpoints return similar format)
# Format: {"token": "...", "expiration": "...", "error": false, "message": "..."}
# or {"token": "..."} for older versions
if result.get("error") is True:
error_msg = result.get("message", "Unknown error")
raise RuntimeError(f"Token generation failed: {error_msg}")
token = result.get("token")
expiration = result.get("expiration")
if token:
logger.debug(
f"Successfully obtained token from {url} "
f"(expiration: {expiration or 'not provided'})"
)
return (token, expiration)
else:
raise ValueError(f"No token in response: {result}")
except requests.exceptions.HTTPError as e:
# Track if this was a 404 error
if e.response.status_code != 404:
all_404_errors = False
# If 404 and we have more endpoints to try, continue
if e.response.status_code == 404 and len(endpoints_to_try) > 1:
logger.debug(
f"Endpoint {url} returned 404, trying next endpoint..."
)
last_error = e
continue
# For other HTTP errors, log and try next endpoint if available
logger.debug(
f"HTTP error {e.response.status_code} on {url}: {e.response.text}"
)
last_error = e
continue
except Exception as e:
all_404_errors = False # Non-HTTP errors are not 404s
logger.debug(f"Error trying {url}: {e}")
last_error = e
continue
# All graph-specific endpoints failed
# If all failures were 404 errors and we have a graph_name, try generating a global token
# This handles cases where the graph doesn't exist yet (e.g., "DefaultGraph" at init time)
# For TigerGraph 4.2.1, /gsql/v1/tokens requires the graph to exist, but /gsql/v1/auth/token
# can generate a global token without a graph parameter
if all_404_errors and graph_name is not None and last_error:
logger.debug(
f"All graph-specific token attempts failed with 404. "
f"Graph '{graph_name}' may not exist yet. "
f"Trying to generate a global token (without graph parameter)..."
)
# Try generating a global token using /gsql/v1/auth/token (works for TG 4.0-4.2.1)
global_token_endpoints = [
(
f"{self.gsql_url}/gsql/v1/auth/token",
{
"secret": secret,
"lifetime": lifetime, # In seconds
# No graph parameter = global token
},
False, # lifetime in seconds
)
]
# Also try /gsql/v1/tokens without graph parameter (for TG 4.2.2+)
global_token_endpoints.append(
(
f"{self.gsql_url}/gsql/v1/tokens",
{
"secret": secret,
"lifetime": lifetime * 1000, # In milliseconds
# No graph parameter = global token
},
True, # lifetime in milliseconds
)
)
for url, payload, _is_milliseconds in global_token_endpoints:
try:
clean_payload = {k: v for k, v in payload.items() if v is not None}
response = requests.post(
url,
headers=headers,
json=clean_payload,
timeout=30,
verify=self.ssl_verify,
)
response.raise_for_status()
result = response.json()
if result.get("error") is True:
error_msg = result.get("message", "Unknown error")
logger.debug(f"Global token generation failed: {error_msg}")
continue
token = result.get("token")
expiration = result.get("expiration")
if token:
logger.info(
f"Successfully obtained global token from {url} "
f"(graph '{graph_name}' may not exist yet, using global token). "
f"Expiration: {expiration or 'not provided'}"
)
return (token, expiration)
except Exception as e:
logger.debug(f"Error trying global token endpoint {url}: {e}")
continue
# All endpoints failed (including global token fallback)
error_msg = f"Failed to get token from secret after trying {len(endpoints_to_try)} endpoint(s)"
if all_404_errors and graph_name:
error_msg += f" (all returned 404, graph '{graph_name}' may not exist yet)"
if last_error:
error_msg += f": {last_error}"
logger.error(error_msg)
raise RuntimeError(error_msg)
def _get_version(self) -> str | None:
"""
Get TigerGraph version using REST API.
Tries multiple endpoints in order:
1. GET /gsql/v1/version (GSQL server, port 14240) - primary for TG 4+
2. GET /version (REST++ server, port 9000) - fallback for older versions
Note: The /version endpoint does NOT exist on GSQL port (14240).
It only exists on REST++ port (9000) for older versions.
Returns:
Version string (e.g., "4.2.1") or None if detection fails
"""
import re
if self.config.gs_port is None:
raise ValueError("gs_port must be set in config for version detection")
# Try GSQL endpoint first (primary for TigerGraph 4+)
# Note: /gsql/v1/version exists on GSQL port, but /version does NOT
# Response format: plain text like "GSQL version: 4.2.2\n"
gsql_url = f"{self.gsql_url}/gsql/v1/version"
headers = self._get_auth_headers(use_basic_auth=True)
try:
response = requests.get(
gsql_url, headers=headers, timeout=10, verify=self.ssl_verify
)
response.raise_for_status()
if not response.text.strip():
# Empty response
logger.debug("GSQL version endpoint returned empty response")
raise ValueError("Empty response from GSQL version endpoint")
# GSQL /gsql/v1/version returns plain text, not JSON
# Format: "GSQL version: 4.2.2\n" or similar
response_text = response.text.strip()
# Try to parse version from text response
# Format: "GSQL version: 4.2.2" or "version: 4.2.2" or "4.2.2"
version_match = re.search(
r"version:\s*(\d+)\.(\d+)\.(\d+)", response_text, re.IGNORECASE
)
if version_match:
version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
logger.debug(
f"Detected TigerGraph version: {version_str} from GSQL endpoint (text format)"
)
return version_str
# Try alternative: just look for version number pattern
version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", response_text)
if version_match:
version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
logger.debug(
f"Detected TigerGraph version: {version_str} from GSQL endpoint (text format)"
)
return version_str
# If text parsing failed, try JSON as fallback (some versions might return JSON)
try:
result = response.json()
message = result.get("message", "")
if message:
version_match = re.search(r"release_(\d+)\.(\d+)\.(\d+)", message)
if version_match:
version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
logger.debug(
f"Detected TigerGraph version: {version_str} from GSQL endpoint (JSON format)"
)
return version_str
except ValueError:
# Not JSON, that's fine - we already tried text parsing
pass
except Exception as e:
logger.debug(f"Failed to get version from GSQL endpoint: {e}")
# Fallback: Try REST++ /version endpoint (for older versions or if GSQL endpoint fails)
# Note: /version only exists on REST++ port (9000), not GSQL port (14240)
try:
# Use REST++ port if different from GSQL port
restpp_port = self.config.port if self.config.port else self.config.gs_port
if restpp_port is None:
return None
restpp_url = f"{self.config.url_without_port}:{restpp_port}/version"
headers = self._get_auth_headers(use_basic_auth=True)
response = requests.get(
restpp_url, headers=headers, timeout=10, verify=self.ssl_verify
)
response.raise_for_status()
# Check content type and response
if not response.text.strip():
logger.debug("REST++ version endpoint returned empty response")
return None
try:
result = response.json()
except ValueError:
logger.debug(
f"REST++ version endpoint returned non-JSON response: "
f"status={response.status_code}, text={response.text[:200]}"
)
return None
# Parse version from REST++ response
message = result.get("message", "")
if message:
version_match = re.search(r"release_(\d+)\.(\d+)\.(\d+)", message)
if version_match:
version_str = f"{version_match.group(1)}.{version_match.group(2)}.{version_match.group(3)}"
logger.debug(
f"Detected TigerGraph version: {version_str} from REST++ endpoint"
)
return version_str
except Exception as e:
logger.debug(f"Failed to get version from REST++ endpoint: {e}")
return None
def _execute_gsql(self, gsql_command: str) -> str:
"""
Execute GSQL command using REST API.
For TigerGraph 4.0-4.2.1, uses POST /gsql/v1/statements endpoint.
Note: GSQL endpoints require Basic Auth (username/password), not Bearer tokens.
Args:
gsql_command: GSQL command string to execute
Returns:
Response string from GSQL execution
"""
url = f"{self.gsql_url}/gsql/v1/statements"
auth_headers = self._get_auth_headers(use_basic_auth=True)
headers = {
"Content-Type": "text/plain",
**auth_headers,
}
# Debug: Log if Authorization header is missing
if "Authorization" not in headers:
logger.error(
f"No Authorization header generated. "
f"Username: {self.config.username}, Password: {'***' if self.config.password else None}"
)
try:
response = requests.post(
url,
headers=headers,
data=gsql_command,
timeout=120,
verify=self.ssl_verify,
)
response.raise_for_status()
# Try to parse JSON response, fallback to text
try:
result = response.json()
# Extract message or result from JSON response
if isinstance(result, dict):
return result.get("message", str(result))
return str(result)
except ValueError:
# Not JSON, return text
return response.text
except requests_exceptions.HTTPError as e:
error_msg = str(e)
# Try to extract error message from response
try:
error_details = e.response.json() if e.response else {}
error_msg = error_details.get("message", error_msg)
except Exception:
pass
raise RuntimeError(f"GSQL execution failed: {error_msg}") from e
def _get_vertex_types(self, graph_name: str | None = None) -> list[str]:
"""
Get list of vertex types using GSQL.
Args:
graph_name: Name of the graph (defaults to self.graphname)
Returns:
List of vertex type names
"""
graph_name = graph_name or self.graphname
try:
result = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW VERTEX *")
# Parse GSQL output using the proper parser
if isinstance(result, str):
return self._parse_show_output(result, "VERTEX")
return []
except Exception as e:
logger.debug(f"Failed to get vertex types via GSQL: {e}")
return []
def _parse_show_edge_output_with_vertices(
self, output: str
) -> dict[str, list[tuple[str, str]]]:
"""
Parse SHOW EDGE * output (compact TigerGraph format).
Returns:
dict mapping edge_name -> list of (source_vertex, target_vertex)
"""
edge_map: dict[str, list[tuple[str, str]]] = defaultdict(list)
# Match lines like:
# - DIRECTED EDGE contains(FROM Author, TO ResearchField|FROM ResearchField, TO ResearchField)
edge_line_pattern = re.compile(
r"-\s+(?:DIRECTED|UNDIRECTED)\s+EDGE\s+(\w+)\(([^)]+)\)"
)
# Match FROM X, TO Y
from_to_pattern = re.compile(r"FROM\s+(\w+)\s*,\s*TO\s+(\w+)")
for line in output.splitlines():
line = line.strip()
if not line.startswith("-"):
continue
edge_match = edge_line_pattern.search(line)
if not edge_match:
continue
edge_name = edge_match.group(1)
endpoints_blob = edge_match.group(2)
# Split multiple vertex pairs
for endpoint in endpoints_blob.split("|"):
ft_match = from_to_pattern.search(endpoint)
if ft_match:
source, target = ft_match.groups()
edge_map[edge_name].append((source, target))
return dict(edge_map)
def _get_edge_types(
self, graph_name: str | None = None
) -> dict[str, list[tuple[str, str]]]:
"""
Get edge types and their (source, target) vertex pairs using GSQL.
Args:
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Dict mapping edge_type -> list of (source_vertex, target_vertex)
"""
graph_name = graph_name or self.graphname
try:
result = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW EDGE *")
if isinstance(result, str):
return self._parse_show_edge_output_with_vertices(result)
return {}
except Exception as e:
logger.error(f"Failed to get edge types via GSQL: {e}")
return {}
def _get_installed_queries(self, graph_name: str | None = None) -> list[str]:
"""
Get list of installed queries using REST API.
Uses the /endpoints endpoint with dynamic=true to get all installed query endpoints,
then extracts query names from the endpoint paths.
Args:
graph_name: Name of the graph (defaults to self.graphname)
Returns:
List of query names
"""
graph_name = graph_name or self.graphname
try:
# Use REST API endpoint to get dynamic endpoints (installed queries)
# Format: GET /endpoints?dynamic=true
endpoint = "/endpoints"
params = {"dynamic": "true"}
result = self._call_restpp_api(endpoint, method="GET", params=params)
# Parse the response to extract query names
# The response is a dict where keys are endpoint paths like:
# "POST /query/{graph_name}/{query_name}" or "GET /query/{graph_name}/{query_name}"
queries = []
if isinstance(result, dict):
query_prefix = f"/query/{graph_name}/"
for endpoint_path in result.keys():
# Extract query name from endpoint path
# Format: "POST /query/{graph_name}/{query_name}" or "GET /query/{graph_name}/{query_name}"
if query_prefix in endpoint_path:
# Extract the query name after the graph name
# Handle both "POST /query/..." and "/query/..." formats
idx = endpoint_path.find(query_prefix)
if idx >= 0:
query_part = endpoint_path[idx + len(query_prefix) :]
# Extract query name (everything up to first space, newline, or end)
query_name = query_part.split()[0] if query_part else ""
# Remove any trailing slashes or special characters
query_name = query_name.rstrip("/").strip()
if query_name and query_name not in queries:
queries.append(query_name)
return queries
except Exception as e:
logger.debug(f"Failed to get installed queries via REST API: {e}")
return []
def _drop_installed_queries_for_graph(self, graph_name: str) -> None:
"""Drop all installed queries that belong to the provided graph.
Uses GSQL ``DROP QUERY *`` as the primary mechanism — this removes every
installed query in the graph in one shot and does not require prior
discovery. The REST-API-based individual-drop path runs afterwards as
a best-effort cleanup for any stragglers.
TigerGraph will not DROP GRAPH while installed queries exist; this step
must succeed before the graph can be removed.
"""
# Primary: bulk drop via GSQL — works regardless of what the REST API reports.
try:
self._execute_gsql(f"USE GRAPH {graph_name}\nDROP QUERY *")
logger.debug(f"Bulk-dropped all queries from graph '{graph_name}'")
except Exception as e:
logger.debug(
f"Bulk DROP QUERY * for graph '{graph_name}' failed (may have no queries): {e}"
)
# Secondary: REST-API discovery + individual drops for any stragglers.
queries = self._get_installed_queries(graph_name=graph_name)
if queries:
logger.info(
f"Dropping {len(queries)} remaining queries from graph '{graph_name}'"
)
for query_name in queries:
try:
self._execute_gsql(
f"USE GRAPH {graph_name}\nDROP QUERY {query_name} IF EXISTS"
)
logger.debug(
f"Dropped query '{query_name}' from graph '{graph_name}'"
)
except Exception:
try:
self._execute_gsql(
f"USE GRAPH {graph_name}\nDROP QUERY {query_name}"
)
except Exception as query_error:
logger.debug(
f"Could not drop query '{query_name}' from graph "
f"'{graph_name}': {query_error}"
)
self._installed_clear_data_queries.pop(graph_name, None)
def _drop_global_schema_types(
self, schema: "Schema", surviving_graphs: list[str]
) -> None:
"""Drop global vertex and edge types that belong to *schema*.
TigerGraph stores vertex/edge types globally. When a graph is dropped
the types may linger as orphans and block subsequent schema creation for
a graph with the same name. This method cleans them up idempotently.
Types that still appear in *surviving_graphs* (other graphs on the
server) are **not** dropped: a global ``DROP VERTEX`` / ``DROP EDGE``
can cascade-invalidate installed queries in unrelated graphs.
Order: edges first (they depend on vertices), then vertices.
"""
in_use_vertices: set[str] = set()
in_use_edges: set[str] = set()
for g in surviving_graphs:
verts, edges = self._get_graph_type_names(g)
in_use_vertices |= verts
in_use_edges |= edges
db_schema = schema.resolve_db_aware(DBType.TIGERGRAPH)
edge_config = schema.core_schema.edge_config
# Collect unique edge relation names
edge_names: set[str] = set()
for edge in edge_config.values():
runtime = db_schema.edge_config.runtime(edge)
rel = runtime.relation_name or f"{edge.source}_{edge.target}"
if rel:
edge_names.add(rel)
for edge_name in edge_names:
if edge_name in in_use_edges:
logger.warning(
f"Skipping DROP EDGE '{edge_name}' — still referenced by "
"surviving graphs"
)
continue
try:
result = self._execute_gsql(f"DROP EDGE {edge_name}")
logger.warning(f"Dropped global edge type '{edge_name}': {result}")
except Exception as e:
logger.debug(
f"Could not drop global edge type '{edge_name}' "
f"(may not exist or still referenced): {e}"
)
# Collect unique vertex db-names
vertex_names: set[str] = set()
for vertex in schema.core_schema.vertex_config.vertices:
try:
dbname = db_schema.vertex_config.vertex_dbname(vertex.name)
vertex_names.add(dbname if dbname else vertex.name)
except Exception:
vertex_names.add(vertex.name)
for vertex_name in vertex_names:
if vertex_name in in_use_vertices:
logger.warning(
f"Skipping DROP VERTEX '{vertex_name}' — still referenced by "
"surviving graphs"
)
continue
try:
result = self._execute_gsql(f"DROP VERTEX {vertex_name}")
logger.warning(f"Dropped global vertex type '{vertex_name}': {result}")
except Exception as e:
logger.debug(
f"Could not drop global vertex type '{vertex_name}' "
f"(may not exist or still referenced): {e}"
)
def _drop_jobs_for_graph(self, graph_name: str) -> None:
"""Drop jobs visible in the given graph context."""
try:
result = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW JOB *")
except Exception as e:
logger.debug(f"Could not list jobs for graph '{graph_name}': {e}")
return
job_names = self._parse_show_job_output(str(result))
if not job_names:
logger.debug(f"No jobs found for graph '{graph_name}'")
return
logger.info(f"Dropping {len(job_names)} jobs from graph '{graph_name}'")
for job_name in job_names:
try:
self._execute_gsql(f"USE GRAPH {graph_name}\nDROP JOB {job_name}")
logger.debug(f"Dropped job '{job_name}' from graph '{graph_name}'")
except Exception as e:
logger.debug(
f"Could not drop job '{job_name}' from graph '{graph_name}': {e}"
)
def _run_installed_query(
self, query_name: str, graph_name: str | None = None, **kwargs: Any
) -> dict[str, Any] | list[dict]:
"""
Run an installed query using REST API.
Args:
query_name: Name of the installed query
graph_name: Name of the graph (defaults to self.graphname)
**kwargs: Query parameters
Returns:
Query result (dict or list)
"""
graph_name = graph_name or self.graphname
endpoint = f"/query/{graph_name}/{query_name}"
result = self._call_restpp_api(endpoint, method="POST", data=kwargs)
if (
isinstance(result, dict)
and result.get("error") is True
and self._is_missing_query_endpoint_error(result)
):
# Some TigerGraph environments expose installed query endpoints as GET-only.
return self._call_restpp_api(endpoint, method="GET", params=kwargs)
return result
@staticmethod
def _is_missing_query_endpoint_error(result: dict[str, Any]) -> bool:
"""Return True when REST++ reports an installed query endpoint is missing."""
message = str(result.get("message", "")).lower()
details = str(result.get("details", "")).lower()
return (
"endpoint is not found" in message
or "endpoint is not found" in details
or "no such endpoint" in message
or "no such endpoint" in details
)
def _build_clear_data_query_name(self, vertex_types: tuple[str, ...]) -> str:
"""Build a stable, schema-aware query name for clear-data operations."""
signature = "|".join(sorted(vertex_types))
digest = hashlib.sha1(signature.encode("utf-8")).hexdigest()[:12]
return f"graflo_clear_data_{digest}"
def _install_clear_data_query(
self, graph_name: str, query_name: str, vertex_types: tuple[str, ...]
) -> None:
"""Create and install a pre-compiled query that deletes all schema vertex types."""
delete_stmts = "\n".join(
f" DELETE FROM {vertex_type};" for vertex_type in sorted(vertex_types)
)
create_query = "\n".join(
[
f"USE GRAPH {graph_name}",
f"CREATE QUERY {query_name}() FOR GRAPH {graph_name} {{",
delete_stmts,
"}",
]
)
install_query = "\n".join(
[
f"USE GRAPH {graph_name}",
f"INSTALL QUERY {query_name}",
]
)
self._execute_gsql(create_query)
self._execute_gsql(install_query)
def _clear_data_via_installed_query(
self, graph_name: str, vertex_types: tuple[str, ...]
) -> None:
"""Run clear-data through an installed GSQL query for faster cluster cleanup."""
query_name = self._build_clear_data_query_name(vertex_types)
installed_queries = self._installed_clear_data_queries.get(graph_name)
if installed_queries is None:
installed_queries = set(self._get_installed_queries(graph_name=graph_name))
self._installed_clear_data_queries[graph_name] = installed_queries
if query_name not in installed_queries:
self._install_clear_data_query(
graph_name=graph_name,
query_name=query_name,
vertex_types=vertex_types,
)
installed_queries.add(query_name)
try:
result = self._execute_gsql(
f"USE GRAPH {graph_name}\nRUN QUERY {query_name}()"
)
except Exception as run_error:
raise RuntimeError(
f"Installed clear_data query '{query_name}' failed: {run_error}"
) from run_error
result_text = str(result).lower()
if "error" in result_text or "failed" in result_text:
raise RuntimeError(
f"Installed clear_data query '{query_name}' failed: {result}"
)
def _upsert_vertex(
self,
vertex_type: str,
vertex_id: str,
attributes: dict[str, Any],
graph_name: str | None = None,
) -> dict[str, Any] | list[dict]:
"""
Upsert a single vertex using REST API.
Args:
vertex_type: Vertex type name
vertex_id: Vertex ID
attributes: Vertex attributes
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Response from API
"""
graph_name = graph_name or self.graphname
endpoint = f"/graph/{graph_name}/vertices/{vertex_type}/{quote(str(vertex_id))}"
return self._call_restpp_api(endpoint, method="POST", data=attributes)
def _upsert_edge(
self,
source_type: str,
source_id: str,
edge_type: str,
target_type: str,
target_id: str,
attributes: dict[str, Any] | None = None,
graph_name: str | None = None,
) -> dict[str, Any] | list[dict]:
"""
Upsert a single edge using REST API.
Args:
source_type: Source vertex type
source_id: Source vertex ID
edge_type: Edge type name
target_type: Target vertex type
target_id: Target vertex ID
attributes: Edge attributes (optional)
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Response from API
"""
graph_name = graph_name or self.graphname
# TigerGraph 4.2+: .../edges/{source_type}/{source_id}/{edge_type}/{target_type}/{target_id}
endpoint = (
f"/graph/{graph_name}/edges/"
f"{source_type}/{quote(str(source_id))}/"
f"{edge_type}/"
f"{target_type}/{quote(str(target_id))}"
)
data = attributes if attributes else {}
return self._call_restpp_api(endpoint, method="POST", data=data)
def _get_edges(
self,
source_type: str,
source_id: str,
edge_type: str | None = None,
graph_name: str | None = None,
) -> list[dict[str, Any]]:
"""
Get edges from a vertex using REST API.
Based on pyTigerGraph's getEdges() implementation.
Uses GET /graph/{graph}/edges/{source_vertex_type}/{source_vertex_id} endpoint.
Args:
source_type: Source vertex type
source_id: Source vertex ID
edge_type: Edge type to filter by (optional, filtered client-side)
graph_name: Name of the graph (defaults to self.graphname)
Returns:
List of edge dictionaries
"""
graph_name = graph_name or self.graphname
# Use the correct endpoint format matching pyTigerGraph's _prep_get_edges:
# GET /graph/{graph}/edges/{source_type}/{source_id}
# If edge_type is specified, append it: /graph/{graph}/edges/{source_type}/{source_id}/{edge_type}
if edge_type:
endpoint = f"/graph/{graph_name}/edges/{source_type}/{quote(str(source_id))}/{edge_type}"
else:
endpoint = (
f"/graph/{graph_name}/edges/{source_type}/{quote(str(source_id))}"
)
result = self._call_restpp_api(endpoint, method="GET")
# Parse REST++ API response format
# Response format: {"version": {...}, "error": false, "message": "", "results": [...]}
if isinstance(result, dict):
# Check for error first
if result.get("error") is True:
error_msg = result.get("message", "Unknown error")
logger.error(f"Error fetching edges: {error_msg}")
return []
# Extract results array
if "results" in result:
edges = result["results"]
else:
logger.debug(
f"Unexpected response format from edges endpoint: {result.keys()}"
)
return []
elif isinstance(result, list):
edges = result
else:
logger.debug(
f"Unexpected response type from edges endpoint: {type(result)}"
)
return []
# Filter by edge_type if specified (client-side filtering)
# REST API endpoint doesn't support edge_type filtering directly
if edge_type and isinstance(edges, list):
edges = [
e for e in edges if isinstance(e, dict) and e.get("e_type") == edge_type
]
return edges
def _get_vertices_by_id(
self, vertex_type: str, vertex_id: str, graph_name: str | None = None
) -> dict[str, dict[str, Any]]:
"""
Get vertex by ID using REST API.
Args:
vertex_type: Vertex type name
vertex_id: Vertex ID
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Dictionary mapping vertex_id to vertex data
"""
graph_name = graph_name or self.graphname
endpoint = f"/graph/{graph_name}/vertices/{vertex_type}/{quote(str(vertex_id))}"
result = self._call_restpp_api(endpoint, method="GET")
# Parse response format to match expected format
# Returns {vertex_id: {"attributes": {...}}}
if isinstance(result, dict):
if "results" in result:
# REST API format
results = result["results"]
if results and isinstance(results, list) and len(results) > 0:
vertex_data = results[0]
return {
vertex_id: {"attributes": vertex_data.get("attributes", {})}
}
elif vertex_id in result:
return {vertex_id: result[vertex_id]}
else:
# Try to extract vertex data
return {vertex_id: {"attributes": result.get("attributes", {})}}
return {}
def _get_vertex_count(self, vertex_type: str, graph_name: str | None = None) -> int:
"""
Get vertex count using REST API.
Args:
vertex_type: Vertex type name
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Number of vertices
"""
graph_name = graph_name or self.graphname
endpoint = f"/graph/{graph_name}/vertices/{vertex_type}"
params = {"limit": "1", "count": "true"}
result = self._call_restpp_api(endpoint, method="GET", params=params)
# Parse count from response
if isinstance(result, dict):
return result.get("count", 0)
return 0
def _delete_vertices(
self, vertex_type: str, where: str | None = None, graph_name: str | None = None
) -> dict[str, Any] | list[dict]:
"""
Delete vertices using REST API.
Args:
vertex_type: Vertex type name
where: WHERE clause for filtering (optional)
graph_name: Name of the graph (defaults to self.graphname)
Returns:
Response from API
"""
graph_name = graph_name or self.graphname
endpoint = f"/graph/{graph_name}/vertices/{vertex_type}"
params = {}
if where:
params["filter"] = where
return self._call_restpp_api(endpoint, method="DELETE", params=params)
def _call_restpp_api(
self,
endpoint: str,
method: str = "GET",
data: dict[str, Any] | None = None,
params: dict[str, str] | None = None,
) -> dict[str, Any] | list[dict]:
"""Call TigerGraph REST++ API endpoint.
Args:
endpoint: REST++ API endpoint (e.g., "/graph/{graph_name}/vertices/{vertex_type}")
method: HTTP method (GET, POST, etc.)
data: Optional data to send in request body (for POST)
params: Optional query parameters
Returns:
Response data (dict or list)
"""
url = f"{self.restpp_url}{endpoint}"
headers = {
"Content-Type": "application/json",
**self._get_auth_headers(),
}
logger.debug(f"REST++ API call: {method} {url}")
try:
if method.upper() == "GET":
response = requests.get(
url,
headers=headers,
params=params,
timeout=120,
verify=self.ssl_verify,
)
elif method.upper() == "POST":
response = requests.post(
url,
headers=headers,
data=json.dumps(data, default=_json_serializer) if data else None,
params=params,
timeout=120,
verify=self.ssl_verify,
)
elif method.upper() == "DELETE":
response = requests.delete(
url,
headers=headers,
params=params,
timeout=120,
verify=self.ssl_verify,
)
else:
raise ValueError(f"Unsupported HTTP method: {method}")
response.raise_for_status()
return response.json()
except requests_exceptions.HTTPError as errh:
# For TigerGraph 4.2.1, if token auth fails with 401/REST-10018, try Basic Auth fallback
if (
errh.response.status_code == 401
and self.api_token
and self.config.username
and self.config.password
and "REST-10018" in str(errh)
):
logger.warning(
"Token authentication failed with REST-10018, "
"falling back to Basic Auth for TigerGraph 4.2.1 compatibility"
)
# Retry with Basic Auth
import base64
credentials = f"{self.config.username}:{self.config.password}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
headers["Authorization"] = f"Basic {encoded_credentials}"
try:
if method.upper() == "GET":
response = requests.get(
url,
headers=headers,
params=params,
timeout=120,
verify=self.ssl_verify,
)
elif method.upper() == "POST":
response = requests.post(
url,
headers=headers,
data=json.dumps(data, default=_json_serializer)
if data
else None,
params=params,
timeout=120,
verify=self.ssl_verify,
)
elif method.upper() == "DELETE":
response = requests.delete(
url,
headers=headers,
params=params,
timeout=120,
verify=self.ssl_verify,
)
else:
raise ValueError(f"Unsupported HTTP method: {method}")
response.raise_for_status()
logger.info("Successfully authenticated using Basic Auth fallback")
return response.json()
except requests_exceptions.HTTPError as errh2:
logger.error(f"HTTP Error (after Basic Auth fallback): {errh2}")
error_response = {"error": True, "message": str(errh2)}
try:
error_json = response.json()
if isinstance(error_json, dict):
error_response.update(error_json)
else:
error_response["details"] = response.text
except Exception:
error_response["details"] = response.text
return error_response
logger.error(f"HTTP Error: {errh}")
error_response = {"error": True, "message": str(errh)}
try:
# Try to parse error response for more details
error_json = response.json()
if isinstance(error_json, dict):
error_response.update(error_json)
else:
error_response["details"] = response.text
except Exception:
error_response["details"] = response.text
return error_response
except requests_exceptions.ConnectionError as errc:
logger.error(f"Error Connecting: {errc}")
return {"error": True, "message": str(errc)}
except requests_exceptions.Timeout as errt:
logger.error(f"Timeout Error: {errt}")
return {"error": True, "message": str(errt)}
except requests_exceptions.RequestException as err:
logger.error(f"An unexpected error occurred: {err}")
return {"error": True, "message": str(err)}
@contextlib.contextmanager
def _ensure_graph_context(self, graph_name: str | None = None):
"""
Context manager that ensures graph context for metadata operations.
Stores graph name for operations that need it.
Args:
graph_name: Name of the graph to use. If None, uses configured graph name.
Yields:
The graph name that was set.
"""
graph_name = graph_name or self._configured_graph_name()
if not graph_name:
raise ValueError(
"Graph name must be provided via graph_name parameter "
"or config.database/config.schema_name"
)
old_graphname = self.graphname
self.graphname = graph_name
try:
yield graph_name
finally:
# Restore original graphname
self.graphname = old_graphname
def _get_all_graph_names(self) -> list[str]:
"""Return all graph names currently in TigerGraph (SHOW GRAPH * + parser)."""
try:
result = self._execute_gsql("USE GLOBAL\nSHOW GRAPH *")
return self._parse_show_graph_output(str(result))
except Exception as e:
logger.warning(f"Could not list graphs for orphan check: {e}")
return []
def _get_graph_type_names(self, graph_name: str) -> tuple[set[str], set[str]]:
"""Return ``(vertex_names, edge_names)`` visible in *graph_name* context."""
vertex_names: set[str] = set()
edge_names: set[str] = set()
try:
r = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW VERTEX *")
vertex_names = set(self._parse_show_vertex_output(str(r)))
except Exception as e:
logger.debug(f"Could not list vertices for graph '{graph_name}': {e}")
try:
r = self._execute_gsql(f"USE GRAPH {graph_name}\nSHOW EDGE *")
edge_names = {name for name, _ in self._parse_show_edge_output(str(r))}
except Exception as e:
logger.debug(f"Could not list edges for graph '{graph_name}': {e}")
return vertex_names, edge_names
def _snapshot_all_queries(self) -> dict[str, list[str]]:
"""Return ``{graph_name: [installed_query_names]}`` for every graph.
Used before/after destructive operations to detect accidental query loss
in graphs that were not the direct target of a recreate.
"""
snapshot: dict[str, list[str]] = {}
for graph_name in self._get_all_graph_names():
snapshot[graph_name] = self._get_installed_queries(graph_name=graph_name)
return snapshot
def graph_exists(self, name: str) -> bool:
"""
Check if a graph with the given name exists.
Prefers `SHOW GRAPH *` parsing for deterministic existence checks,
with a best-effort fallback to `USE GRAPH` output heuristics.
Args:
name: Name of the graph to check
Returns:
bool: True if the graph exists, False otherwise
"""
normalized_name = name.strip().lower()
if not normalized_name:
return False
try:
result = self._execute_gsql("USE GLOBAL\nSHOW GRAPH *")
graph_names = self._parse_show_graph_output(str(result))
if graph_names:
return any(g.lower() == normalized_name for g in graph_names)
logger.debug(
"SHOW GRAPH * returned no parsed graphs; falling back to USE GRAPH check for '%s'",
name,
)
except Exception as e:
logger.debug(f"SHOW GRAPH check failed for graph '{name}': {e}")
try:
result = self._execute_gsql(f"USE GRAPH {name}")
result_str = str(result).lower()
return (
"does not exist" not in result_str and "doesn't exist" not in result_str
)
except Exception as e:
logger.debug(f"Fallback USE GRAPH check failed for '{name}': {e}")
error_str = str(e).lower()
if "does not exist" in error_str or "doesn't exist" in error_str:
return False
return False
@_wrap_tg_exception
def create_database(
self,
name: str,
vertex_names: list[str] | None = None,
edge_names: list[str] | None = None,
):
"""
Create a TigerGraph database (graph) using GSQL commands.
This method creates a graph with explicitly attached vertices and edges.
Example: CREATE GRAPH researchGraph (author, paper, wrote)
This method uses direct REST API calls to execute GSQL commands
that create and use the graph. Supported in TigerGraph version 4.2.2+.
Args:
name: Name of the graph to create
vertex_names: Optional list of vertex type names to attach to the graph
edge_names: Optional list of edge type names to attach to the graph
Raises:
RuntimeError: If graph already exists or creation fails
"""
# Check if graph already exists first
if self.graph_exists(name):
raise RuntimeError(f"Graph '{name}' already exists")
try:
# Build the list of types to include in CREATE GRAPH
all_types = []
if vertex_names:
all_types.extend(vertex_names)
if edge_names:
all_types.extend(edge_names)
# Format the CREATE GRAPH command with types
if all_types:
types_str = ", ".join(all_types)
gsql_commands = f"CREATE GRAPH {name} ({types_str})\nUSE GRAPH {name}"
else:
# Fallback to empty graph if no types provided
gsql_commands = f"CREATE GRAPH {name}()\nUSE GRAPH {name}"
# Execute using direct GSQL REST API which handles authentication
logger.debug(f"Creating graph '{name}' via GSQL: {gsql_commands}")
try:
result = self._execute_gsql(gsql_commands)
logger.info(
f"Successfully created graph '{name}' with types {all_types}: {result}"
)
# Verify the result doesn't indicate the graph already existed
result_str = str(result).lower()
if (
"already exists" in result_str
or "duplicate" in result_str
or "graph already exists" in result_str
):
raise RuntimeError(f"Graph '{name}' already exists")
return result
except RuntimeError:
# Re-raise RuntimeError as-is (already handled)
raise
except Exception as e:
error_msg = str(e).lower()
# Check if graph already exists - raise exception in this case
# TigerGraph may return various error messages for existing graphs
if (
"already exists" in error_msg
or "duplicate" in error_msg
or "graph already exists" in error_msg
or "already exist" in error_msg
):
logger.warning(f"Graph '{name}' already exists: {e}")
raise RuntimeError(f"Graph '{name}' already exists") from e
logger.error(f"Failed to create graph '{name}': {e}")
raise
except RuntimeError:
# Re-raise RuntimeError as-is
raise
except Exception as e:
logger.error(f"Error creating graph '{name}' via GSQL: {e}")
raise
def _gsql_result_has_error(self, result: str) -> bool:
"""Return True when a GSQL response text signals a semantic/runtime failure."""
lowered = result.lower()
return (
"semantic check fails" in lowered
or "failed to" in lowered
or "parse error" in lowered
or "syntax error" in lowered
)
@_wrap_tg_exception
def delete_database(self, name: str):
"""
Delete a TigerGraph database (graph).
Teardown sequence:
1) Drop installed queries for the graph
2) Drop jobs scoped to the graph
3) DROP GRAPH
The GSQL endpoint returns HTTP 200 even for logical failures, so we
inspect the response text for GSQL-level error markers rather than
relying on a follow-up graph_exists() call (which can produce false
positives when SHOW GRAPH * is unavailable or slow to propagate).
Args:
name: Name of the graph to delete
"""
logger.debug(f"Attempting to drop graph '{name}'")
self._drop_installed_queries_for_graph(name)
self._drop_jobs_for_graph(name)
result = self._execute_gsql(f"USE GLOBAL\nDROP GRAPH {name}")
result_str = str(result) if result else ""
result_lower = result_str.lower()
# Treat "does not exist" as a success: graph is already gone.
if "does not exist" in result_lower or "doesn't exist" in result_lower:
logger.info(
f"Graph '{name}' did not exist; treating as successful deletion"
)
return result
if self._gsql_result_has_error(result_str):
error_msg = f"DROP GRAPH '{name}' failed: {result_str}"
logger.error(error_msg)
raise RuntimeError(error_msg)
logger.info(f"Successfully dropped graph '{name}'")
return result
@_wrap_tg_exception
def execute(self, query, **kwargs):
"""
Execute GSQL query or installed query based on content.
"""
try:
# Check if this is an installed query call
if query.strip().upper().startswith("RUN "):
# Extract query name and parameters
query_name = query.strip()[4:].split("(")[0].strip()
result = self._run_installed_query(query_name, **kwargs)
else:
# Execute as raw GSQL
result = self._execute_gsql(query)
return result
except Exception as e:
logger.error(f"Error executing query '{query}': {e}")
raise
def close(self):
"""Close connection - no cleanup needed (using direct REST API calls)."""
pass
def bulk_load_begin(
self, schema: Schema, bulk_cfg: TigergraphBulkLoadConfig
) -> str:
"""Start CSV staging session under ``bulk_cfg.staging_dir /<session_id>``."""
if not bulk_cfg.enabled:
raise ValueError(
"bulk_load_begin requires TigergraphBulkLoadConfig.enabled=True"
)
if not bulk_cfg.staging_dir:
raise ValueError(
"TigergraphBulkLoadConfig.staging_dir is required for bulk load"
)
schema_db = schema.resolve_db_aware(DBType.TIGERGRAPH)
if schema_db.vertex_config.blank_vertices:
raise ValueError(
"TigerGraph bulk_load does not support blank_vertices in this release; "
"use REST ingest or remove blank vertex placeholders."
)
session_id = uuid.uuid4().hex[:12]
staging_root = Path(bulk_cfg.staging_dir) / session_id
staging_root.mkdir(parents=True, exist_ok=True)
appender = BulkCsvAppender(
staging_dir=staging_root,
bulk_cfg=bulk_cfg,
schema_db=schema_db,
)
with _tiger_bulk_sessions_lock:
_tiger_bulk_sessions[session_id] = (
appender,
bulk_cfg,
schema_db,
staging_root,
)
return session_id
def bulk_load_append(
self, session_id: str, gc: GraphContainer, schema: Schema
) -> None:
with _tiger_bulk_sessions_lock:
if session_id not in _tiger_bulk_sessions:
raise KeyError(f"Unknown TigerGraph bulk session {session_id!r}")
appender, _, _, _ = _tiger_bulk_sessions[session_id]
appender.append_graph_container(gc, schema)
def bulk_load_finalize( # noqa: PLR0912
self,
session_id: str,
schema: Schema,
*,
bindings: Bindings | None = None,
connection_provider: ConnectionProvider | None = None,
) -> str:
"""Upload to S3 when configured, then CREATE/RUN/DROP LOADING JOB."""
_ = schema
with _tiger_bulk_sessions_lock:
if session_id not in _tiger_bulk_sessions:
raise KeyError(f"Unknown TigerGraph bulk session {session_id!r}")
appender, bulk_cfg, schema_db, _staging_root = _tiger_bulk_sessions.pop(
session_id
)
appender.close()
staged = appender.staged_file_paths
if not staged:
return ""
graph_name = self._require_configured_graph_name()
job_name = f"{bulk_cfg.loading_job.job_name_prefix}_{session_id}"
path_for_gsql: dict[str, str] = {k: str(v.resolve()) for k, v in staged.items()}
proxy = bulk_cfg.resolve_s3_conn_proxy(bindings)
bucket = bulk_cfg.s3_bucket
tigergraph_s3_loader: S3GeneralizedConnConfig | None = None
if proxy and connection_provider is not None:
from graflo.hq.connection_provider import S3GeneralizedConnConfig
gen = connection_provider.get_generalized_config_by_proxy(proxy)
if isinstance(gen, S3GeneralizedConnConfig):
tigergraph_s3_loader = gen
resolved_bucket = bucket or gen.bucket
if not resolved_bucket:
raise ValueError(
"S3 bulk staging requires TigergraphBulkLoadConfig.s3_bucket "
"or S3GeneralizedConnConfig.bucket"
)
path_for_gsql = upload_staged_csvs(
staged_files=staged,
bucket=resolved_bucket,
key_prefix=bulk_cfg.s3_key_prefix,
session_id=session_id,
s3_cfg=gen,
)
if bulk_cfg.loading_job.run_mode == "run_only":
gsql = build_run_loading_job_only(
job_name=job_name, opts=bulk_cfg.loading_job
)
else:
gsql = build_create_and_run_loading_job(
graph_name=graph_name,
job_name=job_name,
schema_db=schema_db,
staged_files=staged,
bulk_cfg=bulk_cfg,
path_for_gsql=path_for_gsql,
tigergraph_s3_loader=tigergraph_s3_loader,
tigergraph_s3_data_source_name=f"gf_s3_{session_id}",
)
return str(self._execute_gsql(gsql))
def _gsql_vertex_field_def(
self,
*,
logical_vertex_name: str,
field_name: str,
tg_type: str,
db_profile: DatabaseProfile | None,
) -> str:
"""Single attribute fragment: ``name TYPE`` or ``name TYPE DEFAULT ...``."""
line = f"{field_name} {tg_type}"
if db_profile is None or not db_profile.has_vertex_property_default(
logical_vertex_name, field_name
):
return line
raw = db_profile.vertex_property_default(logical_vertex_name, field_name)
if raw is None:
return line
lit = gsql_default_literal(raw)
return f"{line} DEFAULT {lit}"
def _get_vertex_add_statement(
self,
vertex: Vertex,
vertex_config,
*,
db_profile: DatabaseProfile | None = None,
) -> str:
"""Generate ADD VERTEX statement for a schema change job.
Args:
vertex: Vertex object to generate statement for
vertex_config: Vertex configuration
db_profile: Optional profile for ``default_property_values`` (GSQL DEFAULT clauses).
Returns:
str: GSQL ADD VERTEX statement
"""
profile = db_profile
if profile is None and hasattr(vertex_config, "db_profile"):
profile = getattr(vertex_config, "db_profile", None)
vertex_dbname = vertex_config.vertex_dbname(vertex.name)
logical = vertex.name
index_fields = vertex_config.identity_fields(vertex.name)
if len(index_fields) == 0:
raise ValueError(
f"Vertex '{vertex_dbname}' must have at least one index field"
)
# Get field type for primary key field(s) - convert FieldType enum to string
field_type_map = {}
for f in vertex.properties:
if f.type:
field_type_map[f.name] = (
f.type.value if hasattr(f.type, "value") else str(f.type)
)
else:
field_type_map[f.name] = FieldType.STRING.value
# Format all fields
all_fields = []
for field in vertex.properties:
if field.type:
field_type = (
field.type.value
if hasattr(field.type, "value")
else str(field.type)
)
else:
field_type = FieldType.STRING.value
all_fields.append((field.name, field_type))
if len(index_fields) == 1:
# Single field: PRIMARY_ID when no DEFAULT on the id; otherwise PRIMARY KEY
# (GSQL does not allow DEFAULT on the PRIMARY_ID id_name id_type fragment).
primary_field_name = index_fields[0]
primary_field_type = field_type_map.get(
primary_field_name, FieldType.STRING.value
)
other_fields = [
(name, ftype)
for name, ftype in all_fields
if name != primary_field_name
]
primary_attr = self._gsql_vertex_field_def(
logical_vertex_name=logical,
field_name=primary_field_name,
tg_type=primary_field_type,
db_profile=profile,
)
# PRIMARY_ID form is `PRIMARY_ID id_name id_type` — DEFAULT is only valid on
# attributes in attribute_list, not on the PRIMARY_ID fragment (GSQL parse error).
# When the identity field needs DEFAULT, use PRIMARY KEY form instead:
# `name TYPE [DEFAULT ...] PRIMARY KEY`.
primary_default_val = (
profile.vertex_property_default(logical, primary_field_name)
if profile is not None
else None
)
if primary_default_val is not None:
field_parts = [f"{primary_attr} PRIMARY KEY"]
vertex_with = 'WITH STATS="OUTDEGREE_BY_EDGETYPE"'
else:
field_parts = [f"PRIMARY_ID {primary_attr}"]
vertex_with = (
'WITH STATS="OUTDEGREE_BY_EDGETYPE", PRIMARY_ID_AS_ATTRIBUTE="true"'
)
for name, ftype in other_fields:
field_parts.append(
self._gsql_vertex_field_def(
logical_vertex_name=logical,
field_name=name,
tg_type=ftype,
db_profile=profile,
)
)
field_definitions = ",\n ".join(field_parts)
return (
f"ADD VERTEX {vertex_dbname} (\n"
f" {field_definitions}\n"
f" ) {vertex_with}"
)
else:
# Composite key: use PRIMARY KEY syntax
field_parts = [
self._gsql_vertex_field_def(
logical_vertex_name=logical,
field_name=name,
tg_type=ftype,
db_profile=profile,
)
for name, ftype in all_fields
]
vindex = "(" + ", ".join(index_fields) + ")"
field_parts.append(f"PRIMARY KEY {vindex}")
field_definitions = ",\n ".join(field_parts)
return (
f"ADD VERTEX {vertex_dbname} (\n"
f" {field_definitions}\n"
f' ) WITH STATS="OUTDEGREE_BY_EDGETYPE"'
)
def _edge_for_tigergraph_ddl(self, edge: Edge, ec_db: EdgeConfigDBAware) -> Edge:
"""Deep-copy edge with TigerGraph-effective weights for GSQL (non-mutating on schema)."""
ew = ec_db.effective_weights(edge)
edge_copy = edge.model_copy(deep=True)
if ew is not None:
edge_copy.properties = [f.model_copy(deep=True) for f in ew.direct]
else:
edge_copy.properties = []
return edge_copy
def _validate_tigergraph_vertex_properties(self, vertex: Vertex) -> None:
"""Reject reserved or invalid TigerGraph names on vertex attributes."""
for field in vertex.properties:
_validate_tigergraph_schema_name(field.name, "vertex property")
def _validate_tigergraph_edge_property_names(
self, edge: Edge, edge_config_db: EdgeConfigDBAware
) -> None:
"""Reject reserved or invalid names on edge attributes and discriminators."""
ddl_edge = self._edge_for_tigergraph_ddl(edge, edge_config_db)
names: set[str] = {f.name for f in ddl_edge.properties}
names.update(self._edge_identity_discriminator_fields(ddl_edge))
for attr in names:
_validate_tigergraph_schema_name(attr, "edge attribute")
def _format_edge_attributes(
self,
edge: Edge,
exclude_fields: set[str] | None = None,
*,
db_profile: DatabaseProfile | None = None,
edge_id: EdgeId | None = None,
) -> str:
"""Format edge attributes for GSQL ADD DIRECTED EDGE statement.
Args:
edge: Edge object to format attributes for
exclude_fields: Optional set of field names to exclude from attributes
db_profile: Optional profile for ``default_property_values`` (GSQL DEFAULT).
edge_id: Logical edge identity; defaults to ``edge.edge_id``.
Returns:
str: Formatted attribute string (e.g., " date STRING,\n relation STRING")
"""
if not edge.properties:
return ""
if exclude_fields is None:
exclude_fields = set()
eid = edge_id if edge_id is not None else edge.edge_id
attr_parts = []
for field in edge.properties:
field_name = field.name
if field_name not in exclude_fields:
field_type = self._get_tigergraph_type(field.type)
segment = f"{field_name} {field_type}"
if db_profile is not None and db_profile.has_edge_property_default(
eid, field_name
):
raw = db_profile.edge_property_default(eid, field_name)
if raw is not None:
lit = gsql_default_literal(raw)
segment = f"{segment} DEFAULT {lit}"
attr_parts.append(f" {segment}")
return ",\n".join(attr_parts)
def _edge_identity_discriminator_fields(self, edge: Edge) -> set[str]:
"""Return TigerGraph discriminator fields from logical edge identities."""
fields: set[str] = set()
for identity_key in edge.identities:
for token in identity_key:
if token in {"source", "target"}:
continue
if token == "relation":
fields.add(DEFAULT_TIGERGRAPH_RELATION_WEIGHTNAME)
continue
if token not in {"_from", "_to"}:
fields.add(token)
return fields
def _get_edge_add_statement(
self,
edge: Edge,
*,
relation_name: str,
source_vertex: str,
target_vertex: str,
db_profile: DatabaseProfile | None = None,
) -> str:
"""Generate ADD DIRECTED EDGE statement for a schema change job.
Args:
edge: Edge object to generate statement for
Returns:
str: GSQL ADD DIRECTED EDGE statement
"""
# TigerGraph discriminators are derived from logical edge identity.
indexed_field_names = self._edge_identity_discriminator_fields(edge)
# IMPORTANT: In TigerGraph, discriminator fields MUST also be edge attributes.
# If an indexed field is not in attributes, we need to add it.
existing_weight_names = {f.name for f in edge.properties}
# Add any indexed fields that are missing from attributes
for field_name in indexed_field_names:
if field_name not in existing_weight_names:
from graflo.architecture.schema.edge import Field
edge.properties.append(Field(name=field_name, type=FieldType.STRING))
existing_weight_names.add(field_name)
logger.info(
f"Added indexed field '{field_name}' to edge attributes for discriminator compatibility"
)
# Format edge attributes, excluding discriminator fields (they're in DISCRIMINATOR clause)
edge_attrs = self._format_edge_attributes(
edge,
exclude_fields=indexed_field_names,
db_profile=db_profile,
edge_id=edge.edge_id,
)
# Build discriminator clause with all indexed fields
# DISCRIMINATOR goes INSIDE parentheses, on same line as FROM/TO, with types
# Format: FROM company, TO company, DISCRIMINATOR(relation STRING), date STRING, ...
# Get field types for discriminator fields
field_types = {}
if edge.properties:
for field in edge.properties:
field_types[field.name] = self._get_tigergraph_type(field.type)
# Use sanitized dbname for schema names when available
relation_db = relation_name
# Build FROM/TO line with discriminator
from_to_parts = [
f" FROM {source_vertex}",
f" TO {target_vertex}",
]
if indexed_field_names:
# Format discriminator with types: DISCRIMINATOR(field1 TYPE1, field2 TYPE2)
discriminator_parts = []
for field_name in sorted(indexed_field_names):
field_type = field_types.get(field_name, "STRING") # Default to STRING
discriminator_parts.append(f"{field_name} {field_type}")
discriminator_str = f"DISCRIMINATOR({', '.join(discriminator_parts)})"
from_to_parts.append(f" {discriminator_str}")
logger.info(
f"Added discriminator for edge {relation_db}: {', '.join(discriminator_parts)}"
)
else:
logger.debug(
f"No identity discriminator fields found for edge {relation_db}. "
f"Identities: {edge.identities}, relation: {edge.relation}"
)
# Combine FROM/TO and discriminator with commas
from_to_line = ",\n".join(from_to_parts)
# Build the complete statement
if edge_attrs:
# Has attributes - add comma after FROM/TO line (which may include discriminator)
# edge_attrs already has proper indentation, so we just need to add it after a comma
return (
f"ADD DIRECTED EDGE {relation_db} (\n"
f"{from_to_line},\n"
f"{edge_attrs}\n"
f" )"
)
else:
# No attributes - FROM/TO line (which may include discriminator) is the last thing
# No trailing comma needed
return f"ADD DIRECTED EDGE {relation_db} (\n{from_to_line}\n )"
def _get_edge_group_create_statement(
self,
edges: list[Edge],
*,
relation_name: str,
source_vertices: dict[int, str],
target_vertices: dict[int, str],
db_profile: DatabaseProfile | None = None,
) -> str:
"""Generate ADD DIRECTED EDGE statement for a group of edges with the same relation.
TigerGraph requires edges of the same type to be created in a single statement
with multiple FROM/TO pairs separated by |.
Args:
edges: List of Edge objects with the same relation (edge type)
Returns:
str: GSQL ADD DIRECTED EDGE statement with multiple FROM/TO pairs
"""
if not edges:
raise ValueError("Cannot create edge statement from empty edge list")
# Use the first edge to determine attributes and discriminator
# (all edges of the same relation should have the same schema)
first_edge = edges[0]
relation = relation_name
# Collect identity discriminator fields (same logic as _get_edge_add_statement)
indexed_field_names = self._edge_identity_discriminator_fields(first_edge)
# Ensure indexed fields are in attributes (same logic as _get_edge_add_statement)
existing_weight_names = {f.name for f in first_edge.properties}
for field_name in indexed_field_names:
if field_name not in existing_weight_names:
from graflo.architecture.schema.edge import Field
first_edge.properties.append(
Field(name=field_name, type=FieldType.STRING)
)
existing_weight_names.add(field_name)
# Format edge attributes, excluding discriminator fields
edge_attrs = self._format_edge_attributes(
first_edge,
exclude_fields=indexed_field_names,
db_profile=db_profile,
edge_id=first_edge.edge_id,
)
# Get field types for discriminator fields
field_types = {}
if first_edge.properties:
for field in first_edge.properties:
field_types[field.name] = self._get_tigergraph_type(field.type)
# Build FROM/TO pairs for all edges, separated by |
from_to_lines = []
for edge in edges:
# Build FROM/TO line: "FROM A, TO B" or "FROM A, TO B, DISCRIMINATOR(...)"
from_to_parts = [
f"FROM {source_vertices[id(edge)]}",
f"TO {target_vertices[id(edge)]}",
]
# Add discriminator if needed (same for all edges of the same relation)
if indexed_field_names:
discriminator_parts = []
for field_name in sorted(indexed_field_names):
field_type = field_types.get(field_name, "STRING")
discriminator_parts.append(f"{field_name} {field_type}")
discriminator_str = f"DISCRIMINATOR({', '.join(discriminator_parts)})"
from_to_parts.append(discriminator_str)
# Combine FROM/TO and discriminator with commas on one line
from_to_line = ", ".join(from_to_parts)
from_to_lines.append(f" {from_to_line}")
# Join all FROM/TO pairs with |
all_from_to = " |\n".join(from_to_lines)
# Build the complete statement
if edge_attrs:
# Has attributes - add comma after FROM/TO section
return (
f"ADD DIRECTED EDGE {relation} (\n{all_from_to},\n{edge_attrs}\n )"
)
else:
# No attributes - FROM/TO section is the last thing
return f"ADD DIRECTED EDGE {relation} (\n{all_from_to}\n )"
def _batch_schema_statements(
self, schema_change_stmts: list[str], graph_name: str, max_job_size: int
) -> list[list[str]]:
"""Batch schema change statements into groups that fit within max_job_size.
Intelligently merges small statements together while ensuring no batch
exceeds the maximum job size limit.
Args:
schema_change_stmts: List of schema change statements to batch
graph_name: Name of the graph (used for size estimation)
max_job_size: Maximum size in characters for a single job
Returns:
List of batches, where each batch is a list of statements
"""
if not schema_change_stmts:
return []
# Calculate base overhead for a job
# Use worst-case job name length (multi-batch format) for conservative estimation
worst_case_job_name = (
f"schema_change_{graph_name}_batch_999" # Use large number for worst case
)
base_template = (
f"USE GRAPH {graph_name}\n"
f"CREATE SCHEMA_CHANGE JOB {worst_case_job_name} FOR GRAPH {graph_name} {{\n"
f"}}\n"
f"RUN SCHEMA_CHANGE JOB {worst_case_job_name}"
)
base_overhead = len(base_template)
# Each statement adds 5 characters: first gets " " (4) + ";" (1),
# subsequent get ";\n " (5) between statements, final ";" (1) is included
# For N statements: 4 (first indent) + (N-1)*5 (separators) + 1 (final semicolon) = 5*N
def estimate_batch_size(stmts: list[str]) -> int:
"""Estimate the total size of a batch of statements."""
if not stmts:
return base_overhead
total_stmt_size = sum(len(stmt) for stmt in stmts)
return base_overhead + total_stmt_size + 5 * len(stmts)
# Calculate total estimated size for all statements
num_statements = len(schema_change_stmts)
total_stmt_size = sum(len(stmt) for stmt in schema_change_stmts)
estimated_size = base_overhead + total_stmt_size + 5 * num_statements
# If everything fits in one batch, return single batch
if estimated_size <= max_job_size:
logger.info(
f"Applying schema change as single job (estimated size: {estimated_size} chars)"
)
return [schema_change_stmts]
# Need to split into multiple batches
# Strategy: Use a greedy bin-packing approach that merges small statements
# Start by creating batches, trying to pack as many statements as possible
# into each batch without exceeding max_job_size
batches: list[list[str]] = []
# Sort statements by size (smallest first) to help pack efficiently
# We'll process them in order and try to add to existing batches
stmt_with_size = [(stmt, len(stmt)) for stmt in schema_change_stmts]
stmt_with_size.sort(key=lambda x: x[1]) # Sort by statement size
for stmt, stmt_size in stmt_with_size:
# Calculate overhead for adding this statement: 5 chars (indent + semicolon)
stmt_overhead = 5
# Try to add to an existing batch
added = False
for batch in batches:
current_batch_size = estimate_batch_size(batch)
# Check if adding this statement would exceed the limit
if current_batch_size + stmt_size + stmt_overhead <= max_job_size:
batch.append(stmt)
added = True
break
# If couldn't add to existing batch, create a new one
if not added:
# Check if statement itself is too large
single_stmt_size = estimate_batch_size([stmt])
if single_stmt_size > max_job_size:
logger.warning(
f"Statement exceeds max_job_size ({single_stmt_size} > {max_job_size}). "
f"Will attempt to execute anyway, but may fail."
)
batches.append([stmt])
logger.info(
f"Large schema detected (estimated size: {estimated_size} chars). "
f"Splitting into {len(batches)} batches."
)
return batches
@_wrap_tg_exception
def _define_schema_local(self, schema: Schema) -> None:
"""Define TigerGraph schema locally for the current graph using a SCHEMA_CHANGE job.
Args:
schema: Schema definition
"""
graph_name = self._require_configured_graph_name()
# Validate graph name
_validate_tigergraph_schema_name(graph_name, "graph")
vertex_config = schema.core_schema.vertex_config
edge_config = schema.core_schema.edge_config
db_schema = schema.resolve_db_aware(DBType.TIGERGRAPH)
vertex_stmts = []
edge_stmts = []
# Vertices
for vertex in vertex_config.vertices:
# Validate vertex name
vertex_dbname = db_schema.vertex_config.vertex_dbname(vertex.name)
_validate_tigergraph_schema_name(vertex_dbname, "vertex")
self._validate_tigergraph_vertex_properties(vertex)
stmt = self._get_vertex_add_statement(
vertex,
db_schema.vertex_config,
db_profile=db_schema.db_profile,
)
vertex_stmts.append(stmt)
# Edges - group by relation since TigerGraph requires edges of the same type
# to be created in a single statement with multiple FROM/TO pairs
edges_to_create = list(edge_config.values())
source_vertices: dict[int, str] = {}
target_vertices: dict[int, str] = {}
relation_names: dict[int, str] = {}
for edge in edges_to_create:
runtime = db_schema.edge_config.runtime(edge)
source_vertices[id(edge)] = runtime.source_storage
target_vertices[id(edge)] = runtime.target_storage
edge_dbname = runtime.relation_name or f"{edge.source}_{edge.target}"
relation_names[id(edge)] = edge_dbname
_validate_tigergraph_schema_name(edge_dbname, "edge")
self._validate_tigergraph_edge_property_names(edge, db_schema.edge_config)
# Group edges by relation
edges_by_relation: dict[str, list[Edge]] = defaultdict(list)
for edge in edges_to_create:
key = relation_names[id(edge)]
edges_by_relation[key].append(edge)
# Create one statement per relation with all FROM/TO pairs
for relation, edge_group in edges_by_relation.items():
ddl_edges = [
self._edge_for_tigergraph_ddl(e, db_schema.edge_config)
for e in edge_group
]
ddl_source_vertices = {
id(de): source_vertices[id(og)]
for de, og in zip(ddl_edges, edge_group, strict=True)
}
ddl_target_vertices = {
id(de): target_vertices[id(og)]
for de, og in zip(ddl_edges, edge_group, strict=True)
}
stmt = self._get_edge_group_create_statement(
ddl_edges,
relation_name=relation,
source_vertices=ddl_source_vertices,
target_vertices=ddl_target_vertices,
db_profile=db_schema.db_profile,
)
edge_stmts.append(stmt)
if not vertex_stmts and not edge_stmts:
logger.debug(f"No schema changes to apply for graph '{graph_name}'")
return
# Estimate the size of the GSQL command to determine if we need to split it
# Large SCHEMA_CHANGE JOBs (>30k chars) can cause parser failures with misleading errors
# like "Missing return statement" (which is actually a parser size limit issue)
# We'll split into batches based on configurable max_job_size
# Batch vertices and edges separately, then concatenate
vertex_batches = (
self._batch_schema_statements(
vertex_stmts, graph_name, self.config.max_job_size
)
if vertex_stmts
else []
)
edge_batches = (
self._batch_schema_statements(
edge_stmts, graph_name, self.config.max_job_size
)
if edge_stmts
else []
)
batches = vertex_batches + edge_batches
# Execute batches sequentially
for batch_idx, batch_stmts in enumerate(batches):
job_name = (
f"schema_change_{graph_name}_batch_{batch_idx}"
if len(batches) > 1
else f"schema_change_{graph_name}"
)
# First, try to drop the job if it exists (ignore errors if it doesn't)
try:
drop_job_cmd = f"USE GRAPH {graph_name}\nDROP JOB {job_name}"
self._execute_gsql(drop_job_cmd)
logger.debug(f"Dropped existing schema change job '{job_name}'")
except Exception as e:
err_str = str(e).lower()
# Ignore errors if job doesn't exist
if "not found" in err_str or "could not be found" in err_str:
logger.debug(
f"Schema change job '{job_name}' does not exist, skipping drop"
)
else:
logger.debug(f"Could not drop schema change job '{job_name}': {e}")
# Create and run SCHEMA_CHANGE job for this batch
gsql_commands = [
f"USE GRAPH {graph_name}",
f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
" " + ";\n ".join(batch_stmts) + ";",
"}",
f"RUN SCHEMA_CHANGE JOB {job_name}",
]
full_gsql = "\n".join(gsql_commands)
actual_size = len(full_gsql)
# Safety check: warn if actual size exceeds limit (indicates estimation error)
if actual_size > self.config.max_job_size:
logger.warning(
f"Batch {batch_idx + 1} actual size ({actual_size} chars) exceeds limit ({self.config.max_job_size} chars). "
f"This may cause parser errors. Consider reducing max_job_size or improving estimation."
)
logger.info(
f"Applying schema change batch {batch_idx + 1}/{len(batches)} for graph '{graph_name}' "
f"({len(batch_stmts)} statements, {actual_size} chars)"
)
if actual_size < 5000: # Only log full command if it's reasonably small
logger.debug(f"GSQL command:\n{full_gsql}")
else:
logger.debug(f"GSQL command size: {actual_size} characters")
try:
result = self._execute_gsql(full_gsql)
logger.debug(f"Schema change batch {batch_idx + 1} result: {result}")
# Check if result indicates success - should contain "Local schema change succeeded." near the end
result_str = str(result) if result else ""
if result_str:
# Check for success message near the end (last 500 characters to handle long outputs)
result_tail = (
result_str[-500:] if len(result_str) > 500 else result_str
)
if "Local schema change succeeded." not in result_tail:
error_msg = (
f"Schema change job batch {batch_idx + 1} did not report success. "
f"Expected 'Local schema change succeeded.' near the end of the result. "
f"Result (last 500 chars): {result_tail}"
)
logger.error(error_msg)
logger.error(f"Full result: {result_str}")
raise RuntimeError(error_msg)
# Check if result indicates an error - be more lenient with error detection
# Only treat as error if result explicitly contains error indicators
if (
result
and result_str
and (
"Encountered" in result_str
or "syntax error" in result_str.lower()
or "parse error" in result_str.lower()
or "missing return statement" in result_str.lower()
)
):
# "Missing return statement" is a misleading error - it's actually a parser size limit
# SCHEMA_CHANGE JOB doesn't require RETURN statements, so this indicates parser failure
if "missing return statement" in result_str.lower():
error_msg = (
f"Schema change job batch {batch_idx + 1} failed with parser error. "
f"This is likely due to the GSQL command size ({actual_size} chars) exceeding "
f"TigerGraph's parser limit (~30-40K chars). The 'Missing return statement' error "
f"is misleading - SCHEMA_CHANGE JOB doesn't require RETURN statements. "
f"Original error: {result}"
)
else:
error_msg = f"Schema change job batch {batch_idx + 1} reported an error: {result}"
logger.error(error_msg)
logger.error(
f"GSQL command that failed (first 1000 chars):\n{full_gsql[:1000]}..."
)
raise RuntimeError(error_msg)
except Exception as e:
logger.error(
f"Failed to execute schema change batch {batch_idx + 1}: {e}"
)
raise
# Verify that the schema was actually created by checking vertex and edge types
# Wait a moment for schema changes to propagate (after all batches)
import time
time.sleep(1.0) # Increased wait time
with self._ensure_graph_context(graph_name):
vertex_types = self._get_vertex_types()
edge_types = self._get_edge_types()
# Use vertex_dbname instead of v.name to match what TigerGraph actually creates
# vertex_dbname returns dbname if set, otherwise None - fallback to v.name if None
expected_vertex_types = set()
for v in vertex_config.vertices:
try:
dbname = db_schema.vertex_config.vertex_dbname(v.name)
# If dbname is None, use vertex name
expected_name = dbname if dbname is not None else v.name
except (KeyError, AttributeError):
# Fallback to vertex name if vertex_dbname fails
expected_name = v.name
expected_vertex_types.add(expected_name)
expected_edge_types = {
relation_names[id(e)]
for e in edges_to_create
if relation_names.get(id(e))
}
# Convert to sets for case-insensitive comparison
# TigerGraph may capitalize vertex names, so compare case-insensitively
vertex_types_lower = {vt.lower() for vt in vertex_types}
expected_vertex_types_lower = {evt.lower() for evt in expected_vertex_types}
missing_vertices_lower = expected_vertex_types_lower - vertex_types_lower
# Convert back to original case for error message
missing_vertices = {
evt
for evt in expected_vertex_types
if evt.lower() in missing_vertices_lower
}
missing_edges = expected_edge_types - set(edge_types)
if missing_vertices or missing_edges:
error_msg = (
f"Schema change job completed but types were not created correctly. "
f"Missing vertex types: {missing_vertices}, "
f"Missing edge types: {missing_edges}. "
f"Created vertex types: {vertex_types}, "
f"Created edge types: {edge_types}."
)
logger.error(error_msg)
raise RuntimeError(error_msg)
logger.info(
f"Schema verified: {len(vertex_types)} vertex types, {len(edge_types)} edge types created"
)
@_wrap_tg_exception
def init_db(self, schema: Schema, recreate_schema: bool = False) -> None:
"""
Initialize database with schema definition.
If the graph already exists and recreate_schema is False, raises
SchemaExistsError and the script halts.
Follows the same pattern as ArangoDB:
1. Halt if graph exists and recreate_schema is False
2. Clean (drop graph) if recreate_schema
3. Create graph if not exists
4. Define schema locally within the graph
5. Define indexes
If any step fails, the graph will be cleaned up gracefully.
"""
# Use schema.metadata.name for graph creation
graph_created = False
# Determine graph name from config; fallback to schema.metadata.name.
graph_name = self._configured_graph_name()
if not graph_name:
graph_name = schema.metadata.name
# Update config for subsequent operations
self.config.database = graph_name
self.config.schema_name = graph_name
logger.info(f"Using schema name '{graph_name}' from schema.metadata.name")
# Validate graph name
_validate_tigergraph_schema_name(graph_name, "graph")
try:
if self.graph_exists(graph_name) and not recreate_schema:
raise SchemaExistsError(
f"Schema/graph already exists: graph '{graph_name}'. "
"Set recreate_schema=True to replace, or use clear_data=True before ingestion."
)
if recreate_schema:
pre_query_snapshot = self._snapshot_all_queries()
logger.info(
"Pre-recreate installed-query snapshot for graph '%s': %s",
graph_name,
pre_query_snapshot,
)
try:
graph_existed_before = self.graph_exists(graph_name)
# Drop the graph (queries and jobs are dropped first inside delete_database).
self.delete_database(graph_name)
# TigerGraph stores vertex/edge types globally. Dropping the graph
# does NOT remove those types; they linger as orphans and cause
# "used by another object" failures when we try to re-create them.
# Clean them up explicitly before re-creating the schema.
if graph_existed_before:
surviving_graphs = self._get_all_graph_names()
normalized = graph_name.strip().lower()
surviving_graphs = [
g
for g in surviving_graphs
if g.strip().lower() != normalized
]
logger.debug(
f"Dropping global schema types for graph '{graph_name}' "
f"(surviving graphs for orphan check: {surviving_graphs})"
)
self._drop_global_schema_types(schema, surviving_graphs)
logger.debug(f"Cleaned up graph '{graph_name}' for fresh start")
except Exception as clean_error:
error_msg = (
f"Error during recreate_schema for graph '{graph_name}': "
f"{clean_error}"
)
logger.error(error_msg, exc_info=True)
raise RuntimeError(error_msg) from clean_error
post_query_snapshot = self._snapshot_all_queries()
normalized_graph = graph_name.strip().lower()
for other_graph, pre_queries in pre_query_snapshot.items():
if other_graph.strip().lower() == normalized_graph:
continue
post_queries = post_query_snapshot.get(other_graph, [])
lost = set(pre_queries) - set(post_queries)
if lost:
logger.error(
"QUERY LOSS DETECTED in graph '%s' after recreating '%s': %s",
other_graph,
graph_name,
sorted(lost),
)
# Step 1: Create graph first if it doesn't exist
if not self.graph_exists(graph_name):
logger.debug(f"Creating empty graph '{graph_name}'")
try:
# Create empty graph
self.create_database(graph_name)
graph_created = True
logger.info(f"Successfully created empty graph '{graph_name}'")
except Exception as create_error:
logger.error(
f"Failed to create graph '{graph_name}': {create_error}",
exc_info=True,
)
raise
else:
logger.debug(f"Graph '{graph_name}' already exists in init_db")
# Step 2: Define schema locally for the graph
# This uses a SCHEMA_CHANGE job which is the standard way to define local types
logger.info(f"Defining local schema for graph '{graph_name}'")
try:
self._define_schema_local(schema)
except Exception as schema_error:
logger.error(
f"Failed to define local schema for graph '{graph_name}': {schema_error}",
exc_info=True,
)
raise
# Step 3: Define indexes
try:
self.define_indexes(schema)
logger.info(f"Index definition completed for graph '{graph_name}'")
except Exception as index_error:
logger.error(
f"Failed to define indexes for graph '{graph_name}': {index_error}",
exc_info=True,
)
raise
except Exception as e:
logger.error(f"Error initializing database: {e}")
# Graceful teardown: if graph was created in this session, clean it up
if graph_created:
try:
logger.info(
f"Cleaning up graph '{graph_name}' after initialization failure"
)
self.delete_database(graph_name)
except Exception as cleanup_error:
logger.warning(
f"Failed to clean up graph '{graph_name}': {cleanup_error}"
)
raise
@_wrap_tg_exception
def define_schema(self, schema: Schema):
"""
Define TigerGraph schema locally for the current graph.
Assumes graph already exists (created in init_db).
"""
try:
self._define_schema_local(schema)
except Exception as e:
logger.error(f"Error defining schema: {e}")
raise
def define_vertex_classes( # type: ignore[override]
self, vertex_config: VertexConfig
) -> None:
"""Define TigerGraph vertex types locally for the current graph.
Args:
vertex_config: Vertex configuration containing vertices to create
"""
graph_name = self._require_configured_graph_name()
schema_change_stmts = []
db_vertex = (
VertexConfigDBAware(
vertex_config, DatabaseProfile(db_flavor=DBType.TIGERGRAPH)
)
if not isinstance(vertex_config, VertexConfigDBAware)
else vertex_config
)
for vertex in vertex_config.vertices:
stmt = self._get_vertex_add_statement(vertex, db_vertex)
schema_change_stmts.append(stmt)
if not schema_change_stmts:
return
job_name = f"add_vertices_{graph_name}"
gsql_commands = [
f"USE GRAPH {graph_name}",
f"DROP JOB {job_name}",
f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
" " + ";\n ".join(schema_change_stmts) + ";",
"}",
f"RUN SCHEMA_CHANGE JOB {job_name}",
]
logger.info(f"Adding vertices locally to graph '{graph_name}'")
self._execute_gsql("\n".join(gsql_commands))
def define_edge_classes(self, edges: list[Edge]):
"""Define TigerGraph edge types locally for the current graph.
Args:
edges: List of edges to create
"""
graph_name = self._require_configured_graph_name()
# Need vertex_config for dbname lookup if finish_init hasn't been called
# But edges should ideally already be initialized.
# If not, this might fail or needs a vertex_config.
schema_change_stmts = []
for edge in edges:
stmt = self._get_edge_add_statement(
edge,
relation_name=edge.relation or f"{edge.source}_{edge.target}",
source_vertex=edge.source,
target_vertex=edge.target,
)
schema_change_stmts.append(stmt)
if not schema_change_stmts:
return
job_name = f"add_edges_{graph_name}"
gsql_commands = [
f"USE GRAPH {graph_name}",
f"DROP JOB {job_name}",
f"CREATE SCHEMA_CHANGE JOB {job_name} FOR GRAPH {graph_name} {{",
" " + ";\n ".join(schema_change_stmts) + ";",
"}",
f"RUN SCHEMA_CHANGE JOB {job_name}",
]
logger.info(f"Adding edges locally to graph '{graph_name}'")
self._execute_gsql("\n".join(gsql_commands))
def _format_vertex_fields(self, vertex: Vertex) -> str:
"""
Format vertex fields for GSQL CREATE VERTEX statement.
Uses Field objects with types, applying TigerGraph defaults (STRING for None types).
Formats fields as: field_name TYPE
Args:
vertex: Vertex object with Field definitions
Returns:
str: Formatted field definitions for GSQL CREATE VERTEX statement
"""
fields = vertex.properties
if not fields:
# Default fields if none specified
return 'name STRING DEFAULT "",\n properties MAP<STRING, STRING> DEFAULT (map())'
field_list = []
for field in fields:
# Field type should already be set (STRING if was None)
field_type = field.type or FieldType.STRING.value
# Format as: field_name TYPE (DEFAULT clauses live in schema.db_profile.default_property_values)
field_list.append(f"{field.name} {field_type}")
return ",\n ".join(field_list)
def _format_edge_attributes_for_create(self, edge: Edge) -> str:
"""
Format edge attributes for GSQL CREATE EDGE statement.
Edge properties come from edge.properties (list of Field objects).
Each attribute field needs to be included in the CREATE EDGE statement with its type.
"""
attrs = []
if edge.properties:
for field in edge.properties:
# Field objects have name and type attributes
field_name = field.name
# Get TigerGraph type - FieldType enum values are already in TigerGraph format
tg_type = self._get_tigergraph_type(field.type)
attrs.append(f"{field_name} {tg_type}")
return ",\n " + ",\n ".join(attrs) if attrs else ""
def _get_tigergraph_type(self, field_type: FieldType | str | None) -> str:
"""
Convert field type to TigerGraph type string.
FieldType enum values are already in TigerGraph format (e.g., "INT", "STRING", "DATETIME").
This method normalizes various input formats to the correct TigerGraph type.
Args:
field_type: FieldType enum, string, or None
Returns:
str: TigerGraph type string (e.g., "INT", "STRING", "DATETIME")
"""
if field_type is None:
return FieldType.STRING.value
# If it's a FieldType enum, use its value directly (already in TigerGraph format)
if isinstance(field_type, FieldType):
return field_type.value
# If it's an enum-like object with a value attribute
if hasattr(field_type, "value"):
enum_value = field_type.value
# Convert to string and normalize
enum_value_str = str(enum_value).upper()
# Check if the value matches a FieldType enum value
if enum_value_str in VALID_TIGERGRAPH_TYPES:
return enum_value_str
# Return as string (normalized to uppercase)
return enum_value_str
# If it's a string, normalize and check against FieldType values
field_type_str = str(field_type).upper()
# Check if it matches a FieldType enum value directly
if field_type_str in VALID_TIGERGRAPH_TYPES:
return field_type_str
# Handle TigerGraph-specific type aliases
return TIGERGRAPH_TYPE_ALIASES.get(field_type_str, FieldType.STRING.value)
def define_vertex_indexes(
self, vertex_config: VertexConfig, schema: Schema | None = None
):
"""
TigerGraph automatically indexes primary keys.
Secondary indexes are less common but can be created.
"""
db_vertex = (
schema.resolve_db_aware(DBType.TIGERGRAPH).vertex_config
if schema is not None
else None
)
for vertex_class in vertex_config.vertex_set:
vertex_dbname = (
db_vertex.vertex_dbname(vertex_class) if db_vertex else vertex_class
)
index_list = (
schema.db_profile.vertex_secondary_indexes(vertex_class)
if schema is not None
else []
)
for index_obj in index_list:
self._add_index(vertex_dbname, index_obj)
def define_edge_indexes(self, edges: list[Edge], schema: Schema | None = None):
"""Define indexes for edges if specified.
Note: TigerGraph does not support creating indexes on edge attributes.
Edge indexes are skipped with a warning. Only vertex indexes are supported.
"""
for edge in edges:
index_list = (
schema.db_profile.edge_secondary_indexes(edge.edge_id)
if schema is not None
else []
)
if index_list:
edge_db = (
schema.resolve_db_aware(
DBType.TIGERGRAPH
).edge_config.relation_dbname(edge)
if schema is not None
else (edge.relation or f"{edge.source}_{edge.target}")
)
logger.info(
f"Skipping {len(index_list)} index(es) on edge '{edge_db}': "
f"TigerGraph does not support indexes on edge attributes. "
f"Only vertex indexes are supported."
)
def _add_index(self, obj_name, index: Index, is_vertex_index=True):
"""
Create an index on a vertex type using GSQL schema change jobs.
TigerGraph requires indexes to be created through schema change jobs.
This implementation creates a local schema change job for the current graph.
Note: TigerGraph only supports secondary indexes on vertex attributes, not on edge attributes.
Indexes on edges are not supported and should be skipped.
TigerGraph only supports indexes on a single field.
Indexes with multiple fields will be skipped with a warning.
Args:
obj_name: Name of the vertex type
index: Index configuration object
is_vertex_index: Whether this is a vertex index (True) or edge index (False)
"""
# TigerGraph does not support indexes on edge attributes
if not is_vertex_index:
logger.warning(
f"Skipping index creation on edge '{obj_name}': "
f"TigerGraph does not support indexes on edge attributes. "
f"Only vertex indexes are supported."
)
return
try:
if not index.fields:
logger.warning(f"No fields specified for index on {obj_name}, skipping")
return
# TigerGraph only supports secondary indexes on a single field
if len(index.fields) > 1:
logger.warning(
f"TigerGraph only supports indexes on a single field. "
f"Skipping multi-field index on {obj_name} with fields {index.fields}"
)
return
# We have exactly one field - proceed with index creation
field_name = index.fields[0]
# Generate index name if not provided
if index.name:
index_name = index.name
else:
# Generate name from obj_name and field name
index_name = f"{obj_name}_{field_name}_index"
# Generate job name from obj_name and field name
job_name = f"add_{obj_name}_{field_name}_index"
# Build the ALTER command (single field only)
graph_name = self._configured_graph_name()
if not graph_name:
logger.warning(
f"No graph name configured, cannot create index on {obj_name}"
)
return
# Build the ALTER statement inside the job
# Note: For edges, use "EDGE" not "DIRECTED EDGE" in ALTER statements
obj_type = "VERTEX" if is_vertex_index else "EDGE"
alter_stmt = (
f"ALTER {obj_type} {obj_name} ADD INDEX {index_name} ON ({field_name})"
)
# Step 1: Drop existing job if it exists (ignore errors)
try:
drop_job_cmd = f"USE GRAPH {graph_name}\nDROP JOB {job_name}"
self._execute_gsql(drop_job_cmd)
logger.debug(f"Dropped existing job '{job_name}'")
except Exception as e:
err_str = str(e).lower()
# Ignore errors if job doesn't exist
if "not found" in err_str or "could not be found" in err_str:
logger.debug(f"Job '{job_name}' does not exist, skipping drop")
else:
logger.debug(f"Could not drop job '{job_name}': {e}")
# Step 2: Create the schema change job
# Use local schema change for the graph
create_job_cmd = (
f"USE GRAPH {graph_name}\n"
f"CREATE SCHEMA_CHANGE job {job_name} FOR GRAPH {graph_name} {{{alter_stmt};}}"
)
logger.debug(f"Executing GSQL (create job): {create_job_cmd}")
try:
result = self._execute_gsql(create_job_cmd)
logger.debug(f"Created schema change job '{job_name}': {result}")
except Exception as e:
err = str(e).lower()
# Check if job already exists
if (
"already exists" in err
or "duplicate" in err
or "used by another object" in err
):
logger.debug(f"Schema change job '{job_name}' already exists")
else:
logger.error(
f"Failed to create schema change job '{job_name}': {e}"
)
raise
# Step 2: Run the schema change job
run_job_cmd = f"RUN SCHEMA_CHANGE job {job_name}"
logger.debug(f"Executing GSQL (run job): {run_job_cmd}")
try:
result = self._execute_gsql(run_job_cmd)
logger.debug(
f"Ran schema change job '{job_name}', created index '{index_name}' on {obj_name}: {result}"
)
except Exception as e:
err = str(e).lower()
# Check if index already exists or job was already run
if (
"already exists" in err
or "duplicate" in err
or "used by another object" in err
or "already applied" in err
):
logger.debug(
f"Index '{index_name}' on {obj_name} already exists or job already run, skipping"
)
else:
logger.error(f"Failed to run schema change job '{job_name}': {e}")
raise
except Exception as e:
logger.warning(f"Could not create index for {obj_name}: {e}")
def _parse_show_output(self, result_str: str, prefix: str) -> list[str]:
"""
Parse SHOW * output to extract type names.
Looks for lines matching: "- PREFIX name(" or "PREFIX name("
Args:
result_str: String output from SHOW * GSQL command
prefix: The prefix to look for (e.g., "VERTEX", "EDGE")
Returns:
List of extracted names
"""
import re
names = []
# Pattern: "- VERTEX name(" or "VERTEX name("
# Match lines that contain the prefix followed by a word (the name) and then "("
pattern = rf"(?:^|\s)-?\s*{re.escape(prefix)}\s+(\w+)\s*\("
for line in result_str.split("\n"):
line = line.strip()
if not line:
continue
# Use regex to find matches
match = re.search(pattern, line, re.IGNORECASE)
if match:
name = match.group(1)
if name and name not in names:
names.append(name)
return names
def _parse_show_edge_output(self, result_str: str) -> list[tuple[str, bool]]:
"""
Parse SHOW EDGE * output to extract edge type names and direction.
Format: "- DIRECTED EDGE belongsTo(FROM Author, TO ResearchField, ...)"
or "- UNDIRECTED EDGE edgeName(...)"
Args:
result_str: String output from SHOW EDGE * GSQL command
Returns:
List of tuples (edge_name, is_directed)
"""
import re
edge_types = []
# Pattern for DIRECTED EDGE: "- DIRECTED EDGE name("
directed_pattern = r"(?:^|\s)-?\s*DIRECTED\s+EDGE\s+(\w+)\s*\("
# Pattern for UNDIRECTED EDGE: "- UNDIRECTED EDGE name("
undirected_pattern = r"(?:^|\s)-?\s*UNDIRECTED\s+EDGE\s+(\w+)\s*\("
for line in result_str.split("\n"):
line = line.strip()
if not line:
continue
# Check for DIRECTED EDGE
match = re.search(directed_pattern, line, re.IGNORECASE)
if match:
edge_name = match.group(1)
if edge_name:
edge_types.append((edge_name, True))
continue
# Check for UNDIRECTED EDGE
match = re.search(undirected_pattern, line, re.IGNORECASE)
if match:
edge_name = match.group(1)
if edge_name:
edge_types.append((edge_name, False))
return edge_types
def _is_not_found_error(self, error: Exception | str) -> bool:
"""
Check if an error indicates that an object doesn't exist.
Args:
error: Exception object or error string
Returns:
True if the error indicates "not found" or "does not exist"
"""
err_str = str(error).lower()
return "does not exist" in err_str or "not found" in err_str
def _clean_document(self, doc: dict[str, Any]) -> dict[str, Any]:
"""
Remove internal keys that shouldn't be stored in the database.
Removes keys starting with "_" except "_key".
Args:
doc: Document dictionary to clean
Returns:
Cleaned document dictionary
"""
return {k: v for k, v in doc.items() if not k.startswith("_") or k == "_key"}
def _parse_show_vertex_output(self, result_str: str) -> list[str]:
"""Parse SHOW VERTEX * output to extract vertex type names."""
return self._parse_show_output(result_str, "VERTEX")
def _parse_show_graph_output(self, result_str: str) -> list[str]:
"""Parse SHOW GRAPH * output to extract graph names."""
import re
names: list[str] = []
# Accept multiple TigerGraph output styles:
# - GRAPH name
# - GRAPH name(
# - - Graph name
graph_pattern = re.compile(
r"(?:^|\s)-?\s*GRAPH\s+([A-Za-z_][A-Za-z0-9_]*)\s*(?:\(|$)",
re.IGNORECASE,
)
for line in result_str.split("\n"):
line = line.strip()
if not line:
continue
match = graph_pattern.search(line)
if not match:
continue
graph_name = match.group(1)
if graph_name not in names:
names.append(graph_name)
return names
def _parse_show_job_output(self, result_str: str) -> list[str]:
"""Parse SHOW JOB * output to extract job names."""
return self._parse_show_output(result_str, "JOB")
def delete_graph_structure(
self,
vertex_types: tuple[str, ...] | list[str] = (),
graph_names: tuple[str, ...] | list[str] = (),
delete_all: bool = False,
*,
confirm_global_teardown: bool = False,
) -> None:
"""
Delete graph structure (graphs, vertex types, edge types) from TigerGraph.
In TigerGraph:
- Graph: Top-level container (functions like a database in ArangoDB)
- Vertex Types: Global vertex type definitions (can be shared across graphs)
- Edge Types: Global edge type definitions (can be shared across graphs)
- Vertex and edge types are associated with graphs
Teardown order (``delete_all=True``):
1. Drop targeted graphs (``delete_database`` drops graph-scoped queries and jobs).
2. Drop global edge types that are not still referenced by any surviving graph.
3. Drop global vertex types that are not still referenced by any surviving graph.
Global ``DROP VERTEX`` / ``DROP EDGE`` can silently invalidate installed queries
in unrelated graphs. ``delete_all=True`` therefore requires
``confirm_global_teardown=True``.
Args:
vertex_types: Vertex type names to delete (not used in TigerGraph teardown)
graph_names: Graph names to delete (if empty and delete_all=True, deletes all)
delete_all: If True, perform full teardown of targeted graphs and orphaned
global types.
confirm_global_teardown: Must be True when ``delete_all=True``; otherwise
this method raises without issuing GSQL.
"""
cnames = vertex_types
gnames = graph_names
if delete_all and not confirm_global_teardown:
raise ValueError(
"delete_all=True requires confirm_global_teardown=True. "
"Global teardown can silently drop installed queries in unrelated "
"graphs when shared vertex/edge types are removed."
)
try:
if delete_all:
# Step 1: Drop all graphs
graphs_to_drop = list(gnames) if gnames else []
# Guardrail: never auto-discover and drop every graph by default.
# If no explicit graph target is provided, constrain to current graph.
if not graphs_to_drop:
current_graph = self._configured_graph_name()
if current_graph:
graphs_to_drop = [current_graph]
logger.warning(
"delete_all=True without explicit graph_names: limiting TigerGraph teardown to current graph '%s'",
current_graph,
)
else:
raise ValueError(
"Refusing global TigerGraph teardown without explicit "
"graph_names or config.database/config.schema_name"
)
# Drop each graph
logger.info(
f"Found {len(graphs_to_drop)} graphs to drop: {graphs_to_drop}"
)
for graph_name in graphs_to_drop:
try:
self.delete_database(graph_name)
logger.info(f"Successfully dropped graph '{graph_name}'")
except Exception as e:
if self._is_not_found_error(e):
logger.debug(
f"Graph '{graph_name}' already dropped or doesn't exist"
)
else:
logger.warning(f"Failed to drop graph '{graph_name}': {e}")
logger.warning(
f"Error details: {type(e).__name__}: {str(e)}"
)
dropped_set = {g.strip().lower() for g in graphs_to_drop}
surviving_graphs = [
g
for g in self._get_all_graph_names()
if g.strip().lower() not in dropped_set
]
in_use_vertices: set[str] = set()
in_use_edges: set[str] = set()
for g in surviving_graphs:
verts, edges = self._get_graph_type_names(g)
in_use_vertices |= verts
in_use_edges |= edges
# Step 2: Drop global edge types not still referenced by surviving graphs.
# Edges before vertices (dependencies).
try:
show_edges_cmd = "SHOW EDGE *"
result = self._execute_gsql(show_edges_cmd)
result_str = str(result)
edge_types = self._parse_show_edge_output(result_str)
logger.info(
"Found %s edge types in global catalog: %s",
len(edge_types),
[name for name, _ in edge_types],
)
for e_type, _is_directed in edge_types:
if e_type in in_use_edges:
logger.warning(
"Skipping DROP EDGE '%s' — still referenced by surviving graphs",
e_type,
)
continue
try:
drop_edge_cmd = f"DROP EDGE {e_type}"
logger.debug(f"Executing: {drop_edge_cmd}")
result = self._execute_gsql(drop_edge_cmd)
logger.info(
f"Successfully dropped edge type '{e_type}': {result}"
)
except Exception as e:
if self._is_not_found_error(e):
logger.debug(
f"Edge type '{e_type}' already dropped or doesn't exist"
)
else:
logger.warning(
f"Failed to drop edge type '{e_type}': {e}"
)
logger.warning(
f"Error details: {type(e).__name__}: {str(e)}"
)
except Exception as e:
logger.warning(f"Could not list or drop edge types: {e}")
logger.warning(f"Error details: {type(e).__name__}: {str(e)}")
# Step 3: Drop global vertex types not still referenced by surviving graphs.
try:
show_vertices_cmd = "SHOW VERTEX *"
result = self._execute_gsql(show_vertices_cmd)
result_str = str(result)
listed_vertex_types = self._parse_show_vertex_output(result_str)
logger.info(
"Found %s vertex types in global catalog: %s",
len(listed_vertex_types),
listed_vertex_types,
)
for v_type in listed_vertex_types:
if v_type in in_use_vertices:
logger.warning(
"Skipping DROP VERTEX '%s' — still referenced by "
"surviving graphs",
v_type,
)
continue
try:
try:
result = self._delete_vertices(v_type)
logger.debug(
f"Cleared data from vertex type '{v_type}': {result}"
)
except Exception as clear_err:
logger.debug(
f"Could not clear data from vertex type '{v_type}': {clear_err}"
)
drop_vertex_cmd = f"DROP VERTEX {v_type}"
logger.debug(f"Executing: {drop_vertex_cmd}")
result = self._execute_gsql(drop_vertex_cmd)
logger.info(
f"Successfully dropped vertex type '{v_type}': {result}"
)
except Exception as e:
if self._is_not_found_error(e):
logger.debug(
f"Vertex type '{v_type}' already dropped or doesn't exist"
)
else:
logger.warning(
f"Failed to drop vertex type '{v_type}': {e}"
)
logger.warning(
f"Error details: {type(e).__name__}: {str(e)}"
)
except Exception as e:
logger.warning(f"Could not list or drop vertex types: {e}")
logger.warning(f"Error details: {type(e).__name__}: {str(e)}")
elif gnames:
# Drop specific graphs
for graph_name in gnames:
try:
self.delete_database(graph_name)
except Exception as e:
logger.error(f"Error deleting graph '{graph_name}': {e}")
elif cnames:
# Delete vertices from specific vertex types (data only, not schema)
with self._ensure_graph_context():
for class_name in cnames:
try:
result = self._delete_vertices(class_name)
logger.debug(
f"Deleted vertices from {class_name}: {result}"
)
except Exception as e:
logger.error(
f"Error deleting vertices from {class_name}: {e}"
)
except Exception as e:
logger.error(f"Error in delete_graph_structure: {e}")
def clear_data(self, schema: Schema) -> None:
"""Remove all data from the graph without dropping the schema.
Deletes vertices (and their edges) for all vertex types in the schema.
"""
vc = schema.resolve_db_aware(DBType.TIGERGRAPH).vertex_config
graph_name = self._configured_graph_name() or schema.metadata.name
vertex_types = tuple(vc.vertex_dbname(v) for v in vc.vertex_set)
if not vertex_types:
return
try:
self._clear_data_via_installed_query(
graph_name=graph_name, vertex_types=vertex_types
)
remaining = [
vertex_type
for vertex_type in vertex_types
if len(self.fetch_docs(vertex_type, limit=1)) > 0
]
if remaining:
raise RuntimeError(
"Installed clear_data query completed but left data in: "
+ ", ".join(remaining)
)
logger.info(
"Cleared data via installed query for graph '%s' (%d vertex types)",
graph_name,
len(vertex_types),
)
return
except Exception as query_error:
logger.info(
"Installed clear-data query path failed for graph '%s': %s. "
"Falling back to GSQL vertex deletion.",
graph_name,
query_error,
)
gsql_failures: list[str] = []
for vertex_type in vertex_types:
try:
result = self._execute_gsql(
f"USE GRAPH {graph_name}\nDELETE FROM {vertex_type}"
)
if isinstance(result, dict) and result.get("error") is True:
raise RuntimeError(str(result))
result_text = str(result).lower()
if '"error": true' in result_text or "failed" in result_text:
raise RuntimeError(str(result))
logger.debug(
"Deleted vertices via GSQL from %s in graph %s: %s",
vertex_type,
graph_name,
result,
)
except Exception as gsql_error:
gsql_failures.append(f"{vertex_type}: {gsql_error}")
if not gsql_failures:
logger.info(
"Cleared data via direct GSQL deletion for graph '%s' (%d vertex types)",
graph_name,
len(vertex_types),
)
return
logger.warning(
"Direct GSQL delete path failed for graph '%s': %s. "
"Falling back to REST vertex deletion.",
graph_name,
"; ".join(gsql_failures),
)
failures: list[str] = []
for vertex_type in vertex_types:
try:
result = self._delete_vertices(
vertex_type=vertex_type,
graph_name=graph_name,
)
if isinstance(result, dict) and result.get("error") is True:
raise RuntimeError(
result.get("message", "Unknown TigerGraph error")
)
logger.debug(
"Deleted vertices from %s in graph %s: %s",
vertex_type,
graph_name,
result,
)
except Exception as e:
logger.error(
"Error deleting vertices from %s in graph %s: %s",
vertex_type,
graph_name,
e,
)
failures.append(f"{vertex_type}: {e}")
if failures:
raise RuntimeError(
"TigerGraph clear_data failed for vertex types: " + "; ".join(failures)
)
def _generate_upsert_payload(
self, data: list[dict[str, Any]], vname: str, vindex: tuple[str, ...]
) -> dict[str, Any]:
"""
Transforms a list of dictionaries into the TigerGraph REST++ batch upsert JSON format.
The composite Primary ID is created by concatenating the values of the fields
specified in vindex with an underscore '_'. Index fields are included in the
vertex attributes since PRIMARY KEY fields are automatically accessible as
attributes in TigerGraph queries.
Attribute values are wrapped in {"value": ...} format as required by TigerGraph REST++ API.
Args:
data: List of document dictionaries to upsert
vname: Target vertex name
vindex: Tuple of index fields used to create the composite Primary ID
Returns:
Dictionary in TigerGraph REST++ batch upsert format:
{"vertices": {vname: {vertex_id: {attr_name: {"value": attr_value}, ...}}}}
"""
# Initialize the required JSON structure for vertices
payload: dict[str, Any] = {"vertices": {vname: {}}}
vertex_map = payload["vertices"][vname]
for record in data:
try:
# 1. Calculate the Composite Primary ID
# Assumes all index keys exist in the record
primary_id_components = [str(record[key]) for key in vindex]
vertex_id = "_".join(primary_id_components)
# 2. Clean the record (remove internal keys that shouldn't be stored)
clean_record = self._clean_document(record)
# 3. Keep index fields in attributes
# When using PRIMARY KEY (composite keys), the key fields are automatically
# accessible as attributes in queries, so we include them in the payload
# 4. Format attributes for TigerGraph REST++ API
# TigerGraph requires attribute values to be wrapped in {"value": ...}
# Include falsy but valid values (0, False, "") — only None is omitted.
formatted_attributes = {
k: {"value": v} for k, v in clean_record.items() if v is not None
}
# 5. Add the record attributes to the map using the composite ID as the key
vertex_map[vertex_id] = formatted_attributes
except KeyError as e:
logger.warning(
f"Record is missing a required index field: {e}. Skipping record: {record}"
)
continue
return payload
def _upsert_data(
self,
payload: dict[str, Any],
) -> dict[str, Any]:
"""
Sends the generated JSON payload to the TigerGraph REST++ upsert endpoint.
Args:
payload: The JSON payload in TigerGraph REST++ format
Returns:
Dictionary containing the response from TigerGraph
"""
graph_name = self._require_configured_graph_name()
# Use restpp_url which handles version-specific prefixes (e.g., /restpp for 4.2.1)
url = f"{self.restpp_url}/graph/{graph_name}"
# Use centralized auth headers (supports Bearer token for 4.2.1+)
headers = self._get_auth_headers()
headers["Content-Type"] = "application/json"
logger.debug(f"Attempting batch upsert to: {url}")
try:
response = requests.post(
url,
headers=headers,
data=json.dumps(payload, default=_json_serializer),
# Increase timeout for large batches
timeout=120,
verify=self.ssl_verify,
)
response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx)
# TigerGraph response is a JSON object
return response.json()
except requests_exceptions.HTTPError as errh:
# For TigerGraph 4.2.1, if token auth fails with 401/REST-10018, try Basic Auth fallback
if (
errh.response.status_code == 401
and self.api_token
and self.config.username
and self.config.password
and "REST-10018" in str(errh)
):
logger.warning(
"Token authentication failed with REST-10018, "
"falling back to Basic Auth for TigerGraph 4.2.1 compatibility"
)
# Retry with Basic Auth
import base64
credentials = f"{self.config.username}:{self.config.password}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
headers["Authorization"] = f"Basic {encoded_credentials}"
try:
response = requests.post(
url,
headers=headers,
data=json.dumps(payload, default=_json_serializer),
timeout=120,
verify=self.ssl_verify,
)
response.raise_for_status()
logger.info("Successfully authenticated using Basic Auth fallback")
return response.json()
except requests_exceptions.HTTPError as errh2:
logger.error(f"HTTP Error (after Basic Auth fallback): {errh2}")
error_details = ""
try:
error_details = response.text
except Exception:
pass
return {
"error": True,
"message": str(errh2),
"details": error_details,
}
logger.error(f"HTTP Error: {errh}")
error_details = ""
try:
error_details = response.text
except Exception:
pass
return {"error": True, "message": str(errh), "details": error_details}
except requests_exceptions.ConnectionError as errc:
logger.error(f"Error Connecting: {errc}")
return {"error": True, "message": str(errc)}
except requests_exceptions.Timeout as errt:
logger.error(f"Timeout Error: {errt}")
return {"error": True, "message": str(errt)}
except requests_exceptions.RequestException as err:
logger.error(f"An unexpected error occurred: {err}")
return {"error": True, "message": str(err)}
@_wrap_tg_exception
def upsert_docs_batch(self, docs, class_name, match_keys, **kwargs):
"""
Batch upsert documents as vertices using TigerGraph REST++ API.
Creates a GSQL job and formats the payload for batch upsert operations.
Uses composite Primary IDs constructed from match_keys.
"""
dry = kwargs.pop("dry", False)
if dry:
logger.debug(f"Dry run: would upsert {len(docs)} documents to {class_name}")
return
try:
# Convert match_keys to tuple if it's a list
vindex = tuple(match_keys) if isinstance(match_keys, list) else match_keys
# Generate the upsert payload
payload = self._generate_upsert_payload(docs, class_name, vindex)
# Check if payload has any vertices
if not payload.get("vertices", {}).get(class_name):
logger.warning(f"No valid vertices to upsert for {class_name}")
return
# Send the upsert request
result = self._upsert_data(payload)
if result.get("error"):
logger.error(
f"Error upserting vertices to {class_name}: {result.get('message')}"
)
else:
num_vertices = len(payload["vertices"][class_name])
logger.debug(
f"Upserted {num_vertices} vertices to {class_name}: {result}"
)
return result
except Exception as e:
logger.error(f"Error upserting vertices to {class_name}: {e}")
def _generate_edge_upsert_payloads(
self,
edges_data: list[tuple[dict, dict, dict]],
source_class: str,
target_class: str,
edge_type: str,
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
) -> list[dict[str, Any]]:
"""
Transforms edge data into multiple TigerGraph REST++ batch upsert JSON payloads.
Groups edges by (source_id, target_id, edge_type) and collects all weight combinations
for each triple. Then creates separate payloads by "zipping" the weight lists across
all (source_id, target_id, edge_type) groups.
Args:
edges_data: List of tuples (source_doc, target_doc, edge_props)
source_class: Source vertex type name
target_class: Target vertex type name
edge_type: Edge type/relation name (e.g., "relates")
match_keys_source: Tuple of index fields for source vertex
match_keys_target: Tuple of index fields for target vertex
Returns:
List of payload dictionaries in TigerGraph REST++ format:
[{"edges": {source_v_type: {source_id: {edge_type: {target_v_type: {target_id: attributes}}}}}}, ...]
"""
from collections import defaultdict
# Step 1: Group edges by (source_id, target_id, edge_type) and collect weight combinations
# Structure: {(source_id, target_id, edge_type): [weight_dict1, weight_dict2, ...]}
uvr_weights_map: defaultdict[tuple[str, str, str], list[dict]] = defaultdict(
list
)
# Also track original edge data for fallback
uvr_edges_map: defaultdict[
tuple[str, str, str], list[tuple[dict, dict, dict]]
] = defaultdict(list)
for source_doc, target_doc, edge_props in edges_data:
try:
# Extract IDs
source_id = self._extract_id(source_doc, match_keys_source)
target_id = self._extract_id(target_doc, match_keys_target)
if not source_id or not target_id:
logger.warning(
f"Missing source_id ({source_id}) or target_id ({target_id}) for edge"
)
continue
# Clean and format edge attributes
clean_edge_props = self._clean_document(edge_props)
formatted_attributes = {
k: {"value": v} for k, v in clean_edge_props.items() if v
}
# Group by (source_id, target_id, edge_type)
# edge_type is the actual edge type name (e.g., "relates"), not a weight value
uvr_key = (source_id, target_id, edge_type)
uvr_weights_map[uvr_key].append(formatted_attributes)
uvr_edges_map[uvr_key].append((source_doc, target_doc, edge_props))
except Exception as e:
logger.error(f"Error processing edge: {e}")
continue
# Step 2: Find the maximum number of weights across all (u, v, r) groups
# This determines how many payloads we need to create (k payloads for k max elements)
max_weights = (
max(len(weights_list) for weights_list in uvr_weights_map.values())
if uvr_weights_map
else 0
)
if max_weights == 0:
return []
# Step 3: Create k payloads by "zipping" weight lists across all (u, v, r) groups
# Unlike Python's zip() which stops at the shortest iterable, we create k payloads
# where k is the maximum group size. Payload i contains element i from each group
# (if that group has an element at index i).
payloads = []
for weight_idx in range(max_weights):
payload: dict[str, Any] = {"edges": {source_class: {}}}
source_map = payload["edges"][source_class]
payload_original_edges = []
# Iterate through all (u, v, r) groups and take element at weight_idx
for uvr_key, weights_list in uvr_weights_map.items():
# Skip if this group doesn't have a weight at this index
if weight_idx >= len(weights_list):
continue
source_id, target_id, edge_type_key = uvr_key
weight_attrs = weights_list[weight_idx]
original_edge = uvr_edges_map[uvr_key][weight_idx]
# Build nested structure
if source_id not in source_map:
source_map[source_id] = {edge_type: {}}
if edge_type not in source_map[source_id]:
source_map[source_id][edge_type] = {target_class: {}}
if target_class not in source_map[source_id][edge_type]:
source_map[source_id][edge_type][target_class] = {}
target_map = source_map[source_id][edge_type][target_class]
# Add edge at this index from this (u, v, r) group
target_map[target_id] = weight_attrs
payload_original_edges.append(original_edge)
# Only add payload if it has edges (skip empty payloads)
if payload_original_edges:
payload["_original_edges"] = payload_original_edges
payloads.append(payload)
return payloads
def _extract_id(
self, doc: dict[str, Any], match_keys: list[str] | tuple[str, ...]
) -> str | None:
"""
Extract vertex ID from document based on match keys.
For composite keys, concatenates values with an underscore '_'.
Prefers '_key' if present.
Args:
doc: Document dictionary
match_keys: Keys used to identify the vertex
Returns:
str | None: The extracted ID or None if missing required fields
"""
if not doc:
return None
# Try _key first (common in ArangoDB style docs)
if "_key" in doc and doc["_key"]:
return str(doc["_key"])
# If multiple match keys, create a composite ID
if len(match_keys) > 1:
try:
id_parts = [str(doc[key]) for key in match_keys]
return "_".join(id_parts)
except KeyError:
return None
# Single match key
if len(match_keys) == 1:
key = match_keys[0]
if key in doc and doc[key] is not None:
return str(doc[key])
return None
def _fallback_individual_edge_upsert(
self,
edges_data: list[tuple[dict, dict, dict]],
source_class: str,
target_class: str,
edge_type: str,
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
) -> None:
"""Fallback method for individual edge upserts.
Args:
edges_data: List of tuples (source_doc, target_doc, edge_props)
source_class: Source vertex type name
target_class: Target vertex type name
edge_type: Edge type name
match_keys_source: Keys for source vertex ID
match_keys_target: Keys for target vertex ID
"""
for source_doc, target_doc, edge_props in edges_data:
try:
source_id = self._extract_id(source_doc, match_keys_source)
target_id = self._extract_id(target_doc, match_keys_target)
if source_id and target_id:
clean_edge_props = self._clean_document(edge_props)
# Serialize data for REST API
serialized_props = json.loads(
json.dumps(clean_edge_props, default=_json_serializer)
)
self._upsert_edge(
source_class,
source_id,
edge_type,
target_class,
target_id,
serialized_props,
)
except Exception as e:
logger.error(f"Error upserting individual edge: {e}")
def insert_edges_batch(
self,
docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
source_class: str,
target_class: str,
relation_name: str,
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
filter_uniques: bool = True,
head: int | None = None,
**kwargs: Any,
) -> None:
"""
Batch insert/upsert edges using TigerGraph REST++ API.
Handles edge data in tuple format: [(source_doc, target_doc, edge_props), ...]
or dict format: [{"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}, ...]
Args:
docs_edges: List of edge documents (tuples or dicts)
source_class: Source vertex type name
target_class: Target vertex type name
relation_name: Edge type/relation name
match_keys_source: Keys to match source vertices
match_keys_target: Keys to match target vertices
filter_uniques: If True, filter duplicate edges (used)
head: Optional limit on number of edges to insert (used)
**kwargs: Additional options:
- dry: If True, don't execute the query
- collection_name: Alternative edge type name (used if relation_name is None)
- uniq_weight_fields: Unused in TigerGraph (ArangoDB-specific)
- uniq_weight_collections: Unused in TigerGraph (ArangoDB-specific)
- on_duplicate: Unused in TigerGraph (ArangoDB-specific AQL policy)
- relationship_merge_properties: Unused (Cypher property-graph backends only)
"""
opts = consume_insert_edges_kwargs(kwargs)
dry = opts.dry
collection_name = opts.collection_name
if dry:
if docs_edges is not None:
logger.debug(f"Dry run: would insert {len(docs_edges)} edges")
return
# Process edges list
if isinstance(docs_edges, list):
if head is not None:
docs_edges = docs_edges[:head]
if filter_uniques:
docs_edges = pick_unique_dict(docs_edges)
# Normalize edge data format - handle both tuple and dict formats
if docs_edges is None:
return
normalized_edges = []
for edge_item in docs_edges:
try:
if isinstance(edge_item, tuple) and len(edge_item) == 3:
# Tuple format: (source_doc, target_doc, edge_props)
source_doc, target_doc, edge_props = edge_item
normalized_edges.append((source_doc, target_doc, edge_props))
elif isinstance(edge_item, dict):
# Dict format: {"_source_aux": {...}, "_target_aux": {...}, "_edge_props": {...}}
source_doc = edge_item.get("_source_aux", {})
target_doc = edge_item.get("_target_aux", {})
edge_props = edge_item.get("_edge_props", {})
normalized_edges.append((source_doc, target_doc, edge_props))
else:
logger.warning(f"Unexpected edge format: {edge_item}")
except Exception as e:
logger.error(f"Error normalizing edge item: {e}")
continue
if not normalized_edges:
logger.warning("No valid edges to insert")
return
resolved_edge_type = (relation_name or collection_name or "").strip()
if not resolved_edge_type:
logger.error(
"Edge type must be specified via relation_name or collection_name"
)
return
try:
# Convert match_keys to tuples if they're lists
match_keys_src = (
tuple(match_keys_source)
if isinstance(match_keys_source, list)
else match_keys_source
)
match_keys_tgt = (
tuple(match_keys_target)
if isinstance(match_keys_target, list)
else match_keys_target
)
edge_type = resolved_edge_type
# Generate multiple edge upsert payloads (one per unique attribute combination)
payloads = self._generate_edge_upsert_payloads(
normalized_edges,
source_class,
target_class,
edge_type,
match_keys_src,
match_keys_tgt,
)
if not payloads:
logger.warning(f"No valid edges to upsert for edge type {edge_type}")
return
# Send each payload in batch
total_edges = 0
failed_payloads = []
for i, payload in enumerate(payloads):
edges_payload = payload.get("edges", {})
if not edges_payload or source_class not in edges_payload:
continue
# Store original edges for fallback before removing metadata
original_edges = payload.pop("_original_edges", [])
# Send the batch upsert request
result = self._upsert_data(payload)
# Restore original edges for potential fallback
payload["_original_edges"] = original_edges
if result.get("error"):
logger.error(
f"Error upserting edges of type {edge_type} (payload {i + 1}/{len(payloads)}): "
f"{result.get('message')}"
)
# Collect failed payload for fallback
failed_payloads.append((payload, i))
else:
# Count edges in this payload
edge_count = 0
for source_id_map in edges_payload[source_class].values():
if edge_type in source_id_map:
for target_type_map in source_id_map[edge_type].values():
for attrs_or_list in target_type_map.values():
if isinstance(attrs_or_list, list):
edge_count += len(attrs_or_list)
else:
edge_count += 1
total_edges += edge_count
logger.debug(
f"Upserted {edge_count} edges of type {edge_type} via batch "
f"(payload {i + 1}/{len(payloads)}): {result}"
)
# Handle failed payloads with individual upserts
if failed_payloads:
logger.warning(
f"{len(failed_payloads)} payload(s) failed, falling back to individual upserts"
)
# Extract original edges from failed payloads for individual upsert
failed_edges = []
for payload, _ in failed_payloads:
# Use the stored original edges for this payload
original_edges = payload.get("_original_edges", [])
failed_edges.extend(original_edges)
if failed_edges:
logger.debug(
f"Sending {len(failed_edges)} edges from failed payloads via individual upserts"
)
self._fallback_individual_edge_upsert(
failed_edges,
source_class,
target_class,
edge_type,
match_keys_src,
match_keys_tgt,
)
logger.debug(
f"Total upserted {total_edges} edges of type {edge_type} across {len(payloads)} payloads"
)
return
except Exception as e:
logger.error(f"Error batch inserting edges: {e}")
# Fallback to individual operations
m_src = (
tuple(match_keys_source)
if isinstance(match_keys_source, list)
else match_keys_source
)
m_tgt = (
tuple(match_keys_target)
if isinstance(match_keys_target, list)
else match_keys_target
)
self._fallback_individual_edge_upsert(
normalized_edges,
source_class,
target_class,
resolved_edge_type,
m_src,
m_tgt,
)
def _extract_id(self, doc, match_keys):
"""
Extract vertex ID from document based on match keys.
"""
if not doc:
return None
# Try _key first (common in ArangoDB style docs)
if "_key" in doc and doc["_key"]:
return str(doc["_key"])
# Try other match keys
for key in match_keys:
if key in doc and doc[key] is not None:
return str(doc[key])
# Fallback: create composite ID
id_parts = []
for key in match_keys:
if key in doc and doc[key] is not None:
id_parts.append(str(doc[key]))
return "_".join(id_parts) if id_parts else None
def insert_return_batch(
self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
"""
TigerGraph doesn't have INSERT...RETURN semantics like ArangoDB.
"""
raise NotImplementedError(
"insert_return_batch not supported in TigerGraph - use upsert_docs_batch instead"
)
def _render_rest_filter(
self,
filters: list | dict | FilterExpression | None,
field_types: dict[str, FieldType] | None = None,
) -> str:
"""Convert filter expressions to REST++ filter format.
REST++ filter format: "field=value" or "field>value" etc.
Format: fieldoperatorvalue (no spaces, quotes for string values)
Example: "hindex=10" or "hindex>20" or 'name="John"'
Args:
filters: Filter expression to convert
field_types: Optional mapping of field names to FieldType enum values
Returns:
str: REST++ filter string (empty if no filters)
"""
if filters is not None:
if not isinstance(filters, FilterExpression):
ff = FilterExpression.from_dict(filters)
else:
ff = filters
# Use GSQL flavor with empty doc_name to trigger REST++ format
# Pass field_types to help with proper value quoting
result = ff(
doc_name="",
kind=self.expression_flavor(),
field_types=field_types,
)
return result if isinstance(result, str) else ""
else:
return ""
def fetch_docs(
self,
class_name: str,
filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
limit: int | None = None,
return_keys: list[str] | None = None,
unset_keys: list[str] | None = None,
**kwargs: Any,
) -> list[dict[str, Any]]:
"""
Fetch documents (vertices) with filtering and projection using REST++ API.
Args:
class_name: Vertex type name (or dbname)
filters: Filter expression (list, dict, or FilterExpression)
limit: Maximum number of documents to return
return_keys: Keys to return (projection)
unset_keys: Keys to exclude (projection)
**kwargs: Additional parameters
field_types: Optional mapping of field names to FieldType enum values
Used to properly quote string values in filters
If not provided and vertex_config is provided, will be auto-detected
vertex_config: Optional VertexConfig object to use for field type lookup
Returns:
list: List of fetched documents
"""
try:
graph_name = self._require_configured_graph_name()
# Get field_types from kwargs or auto-detect from vertex_config
field_types = kwargs.get("field_types")
vertex_config = kwargs.get("vertex_config")
if field_types is None and vertex_config is not None:
field_types = {
f.name: f.type for f in vertex_config.properties(class_name)
}
# Build REST++ filter string with field type information
filter_str = self._render_rest_filter(filters, field_types=field_types)
# Build REST++ API endpoint with query parameters manually
# Format: /graph/{graph_name}/vertices/{vertex_type}?filter=...&limit=...
# Example: /graph/g22c97325/vertices/Author?filter=hindex>20&limit=10
endpoint = f"/graph/{graph_name}/vertices/{class_name}"
query_parts = []
if filter_str:
# URL-encode the filter string to handle special characters
encoded_filter = quote(filter_str, safe="=<>!&|")
query_parts.append(f"filter={encoded_filter}")
if limit is not None:
query_parts.append(f"limit={limit}")
if query_parts:
endpoint = f"{endpoint}?{'&'.join(query_parts)}"
logger.debug(f"Calling REST++ API: {endpoint}")
# Call REST++ API directly (no params dict, we built the URL ourselves)
response = self._call_restpp_api(endpoint)
# Parse REST++ response (vertices only)
result: list[dict[str, Any]] = self._parse_restpp_response(
response, is_edge=False
)
# Check for errors
if isinstance(response, dict) and response.get("error"):
raise Exception(
f"REST++ API error: {response.get('message', response)}"
)
# Apply projection (client-side projection is acceptable for result formatting)
if return_keys is not None:
result = [
{k: doc.get(k) for k in return_keys if k in doc}
for doc in result
if isinstance(doc, dict)
]
elif unset_keys is not None:
result = [
{k: v for k, v in doc.items() if k not in unset_keys}
for doc in result
if isinstance(doc, dict)
]
return result
except Exception as e:
logger.error(f"Error fetching documents from {class_name} via REST++: {e}")
raise
def fetch_edges(
self,
from_type: str,
from_id: str,
edge_type: str | None = None,
to_type: str | None = None,
to_id: str | None = None,
filters: list[Any] | dict[str, Any] | FilterExpression | None = None,
limit: int | None = None,
return_keys: list[str] | None = None,
unset_keys: list[str] | None = None,
**kwargs: Any,
) -> list[dict[str, Any]]:
"""
Fetch edges from TigerGraph using REST API.
In TigerGraph, you must know at least one vertex ID before you can fetch edges.
Uses REST API which handles special characters in vertex IDs.
Args:
from_type: Source vertex type (required)
from_id: Source vertex ID (required)
edge_type: Optional edge type to filter by
to_type: Optional target vertex type to filter by (not used in REST API)
to_id: Optional target vertex ID to filter by (not used in REST API)
filters: Additional query filters (not supported by REST API)
limit: Maximum number of edges to return (not supported by REST API)
return_keys: Keys to return (projection)
unset_keys: Keys to exclude (projection)
**kwargs: Additional parameters
Returns:
list: List of fetched edges
"""
try:
if not from_type or not from_id:
raise ValueError(
"from_type and from_id are required for fetching edges in TigerGraph"
)
# Use REST API to get edges
# Returns: list of edge dictionaries
logger.debug(
f"Fetching edges using REST API: from_type={from_type}, from_id={from_id}, edge_type={edge_type}"
)
# Handle None edge_type
edge_type_str = edge_type if edge_type is not None else None
edges = self._get_edges(from_type, from_id, edge_type_str)
# Parse REST API response format
# _get_edges() returns list of edge dicts from REST++ API
# Format: [{"e_type": "...", "from_id": "...", "to_id": "...", "attributes": {...}}, ...]
# The REST API returns edges in a flat format with e_type, from_id, to_id, attributes
if isinstance(edges, list):
# Process each edge to normalize format
result = []
for edge in edges:
if isinstance(edge, dict):
# Normalize edge format - REST API returns flat structure
normalized_edge = {}
# Extract edge type (rename e_type to edge_type for consistency)
normalized_edge["edge_type"] = edge.get(
"e_type", edge.get("edge_type", "")
)
# Extract from/to IDs and types
normalized_edge["from_id"] = edge.get("from_id", "")
normalized_edge["from_type"] = edge.get("from_type", "")
normalized_edge["to_id"] = edge.get("to_id", "")
normalized_edge["to_type"] = edge.get("to_type", "")
# Handle nested "from"/"to" objects if present (some API versions)
if "from" in edge and isinstance(edge["from"], dict):
normalized_edge["from_id"] = edge["from"].get(
"id",
edge["from"].get("v_id", normalized_edge["from_id"]),
)
normalized_edge["from_type"] = edge["from"].get(
"type",
edge["from"].get(
"v_type", normalized_edge["from_type"]
),
)
if "to" in edge and isinstance(edge["to"], dict):
normalized_edge["to_id"] = edge["to"].get(
"id", edge["to"].get("v_id", normalized_edge["to_id"])
)
normalized_edge["to_type"] = edge["to"].get(
"type",
edge["to"].get("v_type", normalized_edge["to_type"]),
)
# Extract attributes and merge into normalized edge
attributes = edge.get("attributes", {})
if attributes:
normalized_edge.update(attributes)
else:
# If no attributes key, include all other fields as attributes
for k, v in edge.items():
if k not in (
"e_type",
"edge_type",
"from",
"to",
"from_id",
"to_id",
"from_type",
"to_type",
"directed",
):
normalized_edge[k] = v
result.append(normalized_edge)
elif isinstance(edges, dict):
# Single edge dict - normalize and wrap in list
normalized_edge = {}
normalized_edge["edge_type"] = edges.get(
"e_type", edges.get("edge_type", "")
)
normalized_edge["from_id"] = edges.get("from_id", "")
normalized_edge["to_id"] = edges.get("to_id", "")
if "from" in edges and isinstance(edges["from"], dict):
normalized_edge["from_id"] = edges["from"].get(
"id", edges["from"].get("v_id", normalized_edge["from_id"])
)
if "to" in edges and isinstance(edges["to"], dict):
normalized_edge["to_id"] = edges["to"].get(
"id", edges["to"].get("v_id", normalized_edge["to_id"])
)
attributes = edges.get("attributes", {})
if attributes:
normalized_edge.update(attributes)
else:
for k, v in edges.items():
if k not in (
"e_type",
"edge_type",
"from",
"to",
"from_id",
"to_id",
):
normalized_edge[k] = v
result = [normalized_edge]
else:
# Fallback for unexpected types
result: list[dict[str, Any]] = []
logger.debug(f"Unexpected edges type: {type(edges)}")
# Apply limit if specified (client-side since REST API doesn't support it)
if limit is not None and limit > 0:
result = result[:limit]
# Apply projection (client-side projection is acceptable for result formatting)
if return_keys is not None:
result = [
{k: doc.get(k) for k in return_keys if k in doc}
for doc in result
if isinstance(doc, dict)
]
elif unset_keys is not None:
result = [
{k: v for k, v in doc.items() if k not in unset_keys}
for doc in result
if isinstance(doc, dict)
]
return result
except Exception as e:
logger.error(f"Error fetching edges via REST API: {e}")
raise
def _parse_restpp_response(
self, response: dict | list, is_edge: bool = False
) -> list[dict]:
"""Parse REST++ API response into list of documents.
Args:
response: REST++ API response (dict or list)
is_edge: Whether this is an edge response (default: False for vertices)
Returns:
list: List of parsed documents
"""
result = []
if isinstance(response, dict):
if "results" in response:
for data in response["results"]:
if is_edge:
# Edge response format: {"e_type": "...", "from_id": "...", "to_id": "...", "attributes": {...}}
edge_type = data.get("e_type", "")
from_id = data.get("from_id", data.get("from", ""))
to_id = data.get("to_id", data.get("to", ""))
attributes = data.get("attributes", {})
doc = {
**attributes,
"edge_type": edge_type,
"from_id": from_id,
"to_id": to_id,
}
else:
# Vertex response format: {"v_id": "...", "attributes": {...}}
vertex_id = data.get("v_id", data.get("id"))
attributes = data.get("attributes", {})
doc = {**attributes, "id": vertex_id}
result.append(doc)
elif isinstance(response, list):
# Direct list response
for data in response:
if isinstance(data, dict):
if is_edge:
edge_type = data.get("e_type", "")
from_id = data.get("from_id", data.get("from", ""))
to_id = data.get("to_id", data.get("to", ""))
attributes = data.get("attributes", data)
doc = {
**attributes,
"edge_type": edge_type,
"from_id": from_id,
"to_id": to_id,
}
else:
vertex_id = data.get("v_id", data.get("id"))
attributes = data.get("attributes", data)
doc = {**attributes, "id": vertex_id}
result.append(doc)
return result
def fetch_present_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
flatten: bool = False,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
"""
Check which documents from batch are present in the database.
"""
try:
present_docs: list[dict[str, Any]] = []
keep_keys_list: list[str] | tuple[str, ...] = (
list(keep_keys) if keep_keys is not None else []
)
if isinstance(keep_keys_list, tuple):
keep_keys_list = list(keep_keys_list)
for doc in batch:
vertex_id = self._extract_id(doc, match_keys)
if not vertex_id:
continue
try:
vertex_data = self._get_vertices_by_id(class_name, vertex_id)
if vertex_data and vertex_id in vertex_data:
# Extract requested keys
vertex_attrs = vertex_data[vertex_id].get("attributes", {})
filtered_doc: dict[str, Any] = {}
if keep_keys_list:
for key in keep_keys_list:
if key == "id":
filtered_doc[key] = vertex_id
elif key in vertex_attrs:
filtered_doc[key] = vertex_attrs[key]
else:
# If no keep_keys specified, return all attributes
filtered_doc = vertex_attrs.copy()
filtered_doc["id"] = vertex_id
present_docs.append(filtered_doc)
except Exception:
# Vertex doesn't exist or error occurred
continue
return present_docs
except Exception as e:
logger.error(f"Error fetching present documents: {e}")
return []
def aggregate(
self,
class_name,
aggregation_function: AggregationType,
discriminant: str | None = None,
aggregated_field: str | None = None,
filters: list | dict | None = None,
):
"""
Perform aggregation operations.
"""
try:
if aggregation_function == AggregationType.COUNT and discriminant is None:
# Simple vertex count
count = self._get_vertex_count(class_name)
return [{"_value": count}]
else:
# Complex aggregations require custom GSQL queries
logger.warning(
f"Complex aggregation {aggregation_function} requires custom GSQL implementation"
)
return []
except Exception as e:
logger.error(f"Error in aggregation: {e}")
return []
def keep_absent_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
"""
Return documents from batch that are NOT present in database.
"""
present_docs = self.fetch_present_documents(
batch=batch,
class_name=class_name,
match_keys=match_keys,
keep_keys=keep_keys,
flatten=False,
filters=filters,
)
# Create a set of IDs from present documents for efficient lookup
present_ids = set()
for present_doc in present_docs:
# Extract ID from present document (it should have 'id' key)
if "id" in present_doc:
present_ids.add(present_doc["id"])
# Find documents that are not present
absent_docs: list[dict[str, Any]] = []
keep_keys_list: list[str] | tuple[str, ...] = (
list(keep_keys) if keep_keys is not None else []
)
if isinstance(keep_keys_list, tuple):
keep_keys_list = list(keep_keys_list)
for doc in batch:
vertex_id = self._extract_id(doc, match_keys)
if not vertex_id or vertex_id not in present_ids:
if keep_keys_list:
# Filter to keep only requested keys
filtered_doc = {k: doc.get(k) for k in keep_keys_list if k in doc}
absent_docs.append(filtered_doc)
else:
absent_docs.append(doc)
return absent_docs
@_wrap_tg_exception
def define_indexes(self, schema: Schema):
"""Define all indexes from schema."""
try:
self.define_vertex_indexes(schema.core_schema.vertex_config, schema=schema)
edges_for_indexes = list(schema.core_schema.edge_config.values())
self.define_edge_indexes(edges_for_indexes, schema=schema)
except Exception as e:
logger.error(f"Error defining indexes: {e}")
def fetch_indexes(self, vertex_type: str | None = None):
"""
Fetch indexes for vertex types using GSQL.
In TigerGraph, indexes are associated with vertex types.
Use DESCRIBE VERTEX to get index information.
Args:
vertex_type: Optional vertex type name to fetch indexes for.
If None, fetches indexes for all vertex types.
Returns:
dict: Mapping of vertex type names to their indexes.
Format: {vertex_type: [{"name": "index_name", "fields": ["field1", ...]}, ...]}
"""
try:
with self._ensure_graph_context():
result = {}
if vertex_type:
vertex_types = [vertex_type]
else:
vertex_types = self._get_vertex_types()
for v_type in vertex_types:
try:
# Parse indexes from the describe output
indexes = []
try:
indexes.append(
{"name": "stat_index", "source": "show_stat"}
)
except Exception:
# If SHOW STAT INDEX doesn't work, try alternative methods
pass
result[v_type] = indexes
except Exception as e:
logger.debug(
f"Could not fetch indexes for vertex type {v_type}: {e}"
)
result[v_type] = []
return result
except Exception as e:
logger.error(f"Error fetching indexes: {e}")
return {}