awslabs.timestream-for-influxdb-mcp-server 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

awslabs/timestream_for_influxdb_mcp_server/__init__.py

@@ -15,4 +15,4 @@
 
 """awslabs.timestream-for-influxdb-mcp-server"""
 
-__version__ = '0.0.11'
+__version__ = '0.0.12'

awslabs/timestream_for_influxdb_mcp_server/server.py

@@ -20,6 +20,7 @@ import os
 from influxdb_client.client.influxdb_client import InfluxDBClient
 from influxdb_client.client.write.point import Point
 from influxdb_client.client.write_api import ASYNCHRONOUS, SYNCHRONOUS
+from influxdb_client.domain.bucket_retention_rules import BucketRetentionRules
 from influxdb_client.domain.write_precision import WritePrecision
 from loguru import logger
 from mcp.server.fastmcp import FastMCP
@@ -28,6 +29,11 @@ from typing import Any, Dict, List, Optional
 from urllib.parse import urlparse
 
 
+# InfluxDB environment variables for data plane operations
+INFLUXDB_TOKEN = os.environ.get('INFLUXDB_TOKEN')
+INFLUXDB_URL = os.environ.get('INFLUXDB_URL')
+INFLUXDB_ORG = os.environ.get('INFLUXDB_ORG')
+
 # Define Field parameters as global variables to avoid duplication
 # Common fields
 REQUIRED_FIELD_DB_CLUSTER_ID = Field(
@@ -180,10 +186,19 @@ REQUIRED_FIELD_STATUS_CLUSTER = Field(
 )
 
 # InfluxDB fields
-REQUIRED_FIELD_URL = Field(..., description='The URL of the InfluxDB server.')
-REQUIRED_FIELD_TOKEN = Field(..., description='The authentication token.')
+OPTIONAL_FIELD_URL = Field(
+    None,
+    description='The URL of the InfluxDB server. Falls back to INFLUXDB_URL env var if not provided.',
+)
+OPTIONAL_FIELD_TOKEN = Field(
+    None,
+    description='The authentication token. Falls back to INFLUXDB_TOKEN env var if not provided.',
+)
 REQUIRED_FIELD_BUCKET_INFLUX = Field(..., description='The destination bucket for writes.')
-REQUIRED_FIELD_ORG = Field(..., description='The organization name.')
+OPTIONAL_FIELD_ORG = Field(
+    None,
+    description='The organization name. Falls back to INFLUXDB_ORG env var if not provided.',
+)
 REQUIRED_FIELD_POINTS = Field(
     ...,
     description='List of data points to write. Each point should be a dictionary with measurement, tags, fields, and optional time.',
@@ -242,6 +257,46 @@ def get_timestream_influxdb_client():
     return client
 
 
+def resolve_influxdb_config(
+    url: Optional[str] = None,
+    token: Optional[str] = None,
+    org: Optional[str] = None,
+    require_org: bool = True,
+) -> tuple:
+    """Resolve InfluxDB configuration from parameters or environment variables.
+
+    Args:
+        url: The URL of the InfluxDB server (optional, falls back to INFLUXDB_URL env var).
+        token: The authentication token (optional, falls back to INFLUXDB_TOKEN env var).
+        org: The organization name (optional, falls back to INFLUXDB_ORG env var).
+        require_org: Whether the organization is required (default True).
+
+    Returns:
+        A tuple of (resolved_url, resolved_token, resolved_org).
+
+    Raises:
+        ValueError: If required parameters are not provided.
+    """
+    resolved_url = url or INFLUXDB_URL
+    resolved_token = token or INFLUXDB_TOKEN
+    resolved_org = org or INFLUXDB_ORG
+
+    if not resolved_url:
+        raise ValueError(
+            'URL must be provided either as parameter or via INFLUXDB_URL environment variable'
+        )
+    if not resolved_token:
+        raise ValueError(
+            'Token must be provided either as parameter or via INFLUXDB_TOKEN environment variable'
+        )
+    if require_org and not resolved_org:
+        raise ValueError(
+            'Organization must be provided either as parameter or via INFLUXDB_ORG environment variable'
+        )
+
+    return resolved_url, resolved_token, resolved_org
+
+
 def get_influxdb_client(url, token, org=None, timeout=10000, verify_ssl: bool = True):
     """Get an InfluxDB client.
 
@@ -1054,10 +1109,10 @@ async def create_db_parameter_group(
 
 @mcp.tool(name='InfluxDBWritePoints', description='Write data points to InfluxDB endpoint.')
 async def influxdb_write_points(
-    url: str = REQUIRED_FIELD_URL,
-    token: str = REQUIRED_FIELD_TOKEN,
-    bucket: str = REQUIRED_FIELD_BUCKET,
-    org: str = REQUIRED_FIELD_ORG,
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    bucket: str = REQUIRED_FIELD_BUCKET_INFLUX,
+    org: Optional[str] = OPTIONAL_FIELD_ORG,
     points: List[Dict[str, Any]] = REQUIRED_FIELD_POINTS,
     time_precision: str = OPTIONAL_FIELD_WRITE_PRECISION,
     sync_mode: Optional[str] = OPTIONAL_FIELD_SYNC_MODE,
@@ -1084,10 +1139,12 @@ async def influxdb_write_points(
             'InfluxDBWritePoints tool invocation not allowed when tool-write-mode is set to False'
         )
 
-    try:
-        client = get_influxdb_client(url, token, org, verify_ssl)
+    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)
 
-        # Set write mode
+    try:
+        client = get_influxdb_client(
+            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
+        )
         if sync_mode and sync_mode.lower() == 'synchronous':
             write_api = client.write_api(write_options=SYNCHRONOUS)
         else:
@@ -1117,7 +1174,7 @@ async def influxdb_write_points(
         # Write points
         write_api.write(
             bucket=bucket,
-            org=org,
+            org=resolved_org,
             record=influx_points,
             write_precision=getattr(WritePrecision, time_precision.upper()),
             verify_ssl=verify_ssl,
@@ -1137,10 +1194,10 @@ async def influxdb_write_points(
 
 @mcp.tool(name='InfluxDBWriteLP', description='Write data in Line Protocol format to InfluxDB.')
 async def influxdb_write_line_protocol(
-    url: str = REQUIRED_FIELD_URL,
-    token: str = REQUIRED_FIELD_TOKEN,
-    bucket: str = REQUIRED_FIELD_BUCKET,
-    org: str = REQUIRED_FIELD_ORG,
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    bucket: str = REQUIRED_FIELD_BUCKET_INFLUX,
+    org: Optional[str] = OPTIONAL_FIELD_ORG,
     data_line_protocol: str = REQUIRED_FIELD_DATA_LINE_PROTOCOL,
     time_precision: str = OPTIONAL_FIELD_WRITE_PRECISION,
     sync_mode: str = OPTIONAL_FIELD_SYNC_MODE,
@@ -1157,8 +1214,12 @@ async def influxdb_write_line_protocol(
             'InfluxDBWriteLineProtocol tool invocation not allowed when tool-write-mode is set to False'
         )
 
+    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)
+
     try:
-        client = get_influxdb_client(url, token, org)
+        client = get_influxdb_client(
+            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
+        )
 
         # Set write mode
         if sync_mode and sync_mode.lower() == 'synchronous':
@@ -1169,7 +1230,7 @@ async def influxdb_write_line_protocol(
         # Write line protocol
         write_api.write(
             bucket=bucket,
-            org=org,
+            org=resolved_org,
             record=data_line_protocol,
             write_precision=getattr(WritePrecision, time_precision.upper()),
             verify_ssl=verify_ssl,
@@ -1189,9 +1250,9 @@ async def influxdb_write_line_protocol(
 
 @mcp.tool(name='InfluxDBQuery', description='Query data from InfluxDB using Flux query language.')
 async def influxdb_query(
-    url: str = REQUIRED_FIELD_URL,
-    token: str = REQUIRED_FIELD_TOKEN,
-    org: str = REQUIRED_FIELD_ORG,
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    org: Optional[str] = OPTIONAL_FIELD_ORG,
     query: str = REQUIRED_FIELD_QUERY,
     verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
 ) -> Dict[str, Any]:
@@ -1200,24 +1261,47 @@ async def influxdb_query(
     Returns:
         Query results in the specified format.
     """
+    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)
+
     try:
-        client = get_influxdb_client(url, token, org, verify_ssl)
+        client = get_influxdb_client(
+            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
+        )
         query_api = client.query_api()
 
         # Return as JSON
-        tables = query_api.query(org=org, query=query)
+        tables = query_api.query(org=resolved_org, query=query)
 
         # Process the tables into a more usable format
+        # System keys that are not tags
+        system_keys = {
+            'result',
+            'table',
+            '_start',
+            '_stop',
+            '_time',
+            '_value',
+            '_field',
+            '_measurement',
+        }
+
         result = []
         for table in tables:
             for record in table.records:
+                # Extract tags by filtering out system keys
+                tags = {
+                    k: v
+                    for k, v in record.values.items()
+                    if k not in system_keys and v is not None
+                }
+
                 result.append(
                     {
                         'measurement': record.get_measurement(),
                         'field': record.get_field(),
                         'value': record.get_value(),
                         'time': record.get_time().isoformat() if record.get_time() else None,
-                        'tags': record.values.get('tags', {}),
+                        'tags': tags,
                     }
                 )
 
@@ -1229,6 +1313,203 @@ async def influxdb_query(
         return {'status': 'error', 'message': str(e)}
 
 
+@mcp.tool(name='InfluxDBListBuckets', description='List all buckets in InfluxDB.')
+async def influxdb_list_buckets(
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    org: Optional[str] = OPTIONAL_FIELD_ORG,
+    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
+) -> Dict[str, Any]:
+    """List all buckets in InfluxDB.
+
+    Returns:
+        List of buckets with their details.
+    """
+    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)
+
+    try:
+        client = get_influxdb_client(
+            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
+        )
+        buckets_api = client.buckets_api()
+        buckets = buckets_api.find_buckets().buckets
+
+        result = []
+        for bucket in buckets:
+            result.append(
+                {
+                    'id': bucket.id,
+                    'name': bucket.name,
+                    'org_id': bucket.org_id,
+                    'retention_period': bucket.retention_rules[0].every_seconds
+                    if bucket.retention_rules
+                    else None,
+                    'created_at': bucket.created_at.isoformat() if bucket.created_at else None,
+                    'updated_at': bucket.updated_at.isoformat() if bucket.updated_at else None,
+                }
+            )
+
+        client.close()
+        return {'status': 'success', 'buckets': result}
+
+    except Exception as e:
+        logger.error(f'Error listing InfluxDB buckets: {str(e)}')
+        return {'status': 'error', 'message': str(e)}
+
+
+@mcp.tool(name='InfluxDBCreateBucket', description='Create a new bucket in InfluxDB.')
+async def influxdb_create_bucket(
+    bucket_name: str = Field(..., description='The name of the bucket to create.'),
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    org: Optional[str] = OPTIONAL_FIELD_ORG,
+    retention_seconds: Optional[int] = Field(
+        None, description='Retention period in seconds. 0 or None means infinite retention.'
+    ),
+    description: Optional[str] = Field(None, description='Description of the bucket.'),
+    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
+    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
+) -> Dict[str, Any]:
+    """Create a new bucket in InfluxDB.
+
+    Returns:
+        Details of the created bucket.
+    """
+    if not tool_write_mode:
+        raise Exception(
+            'InfluxDBCreateBucket tool invocation not allowed when tool-write-mode is set to False'
+        )
+
+    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)
+
+    try:
+        client = get_influxdb_client(
+            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
+        )
+        buckets_api = client.buckets_api()
+
+        # Get org ID from org name
+        orgs_api = client.organizations_api()
+        orgs = orgs_api.find_organizations(org=resolved_org)
+        if not orgs:
+            raise ValueError(f'Organization "{resolved_org}" not found')
+        org_id = orgs[0].id
+
+        retention_rules = []
+        if retention_seconds and retention_seconds > 0:
+            retention_rules.append(
+                BucketRetentionRules(type='expire', every_seconds=retention_seconds)
+            )
+
+        bucket = buckets_api.create_bucket(
+            bucket_name=bucket_name,
+            org_id=org_id,
+            retention_rules=retention_rules,
+            description=description,
+        )
+
+        client.close()
+        return {
+            'status': 'success',
+            'bucket': {
+                'id': bucket.id,
+                'name': bucket.name,
+                'org_id': bucket.org_id,
+                'retention_period': bucket.retention_rules[0].every_seconds
+                if bucket.retention_rules
+                else None,
+                'created_at': bucket.created_at.isoformat() if bucket.created_at else None,
+            },
+        }
+
+    except Exception as e:
+        logger.error(f'Error creating InfluxDB bucket: {str(e)}')
+        return {'status': 'error', 'message': str(e)}
+
+
+@mcp.tool(name='InfluxDBListOrgs', description='List all organizations in InfluxDB.')
+async def influxdb_list_orgs(
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
+) -> Dict[str, Any]:
+    """List all organizations in InfluxDB.
+
+    Returns:
+        List of organizations with their details.
+    """
+    resolved_url, resolved_token, _ = resolve_influxdb_config(
+        url, token, org=None, require_org=False
+    )
+
+    try:
+        client = get_influxdb_client(url=resolved_url, token=resolved_token, verify_ssl=verify_ssl)
+        orgs_api = client.organizations_api()
+        orgs = orgs_api.find_organizations()
+
+        result = []
+        for org in orgs:
+            result.append(
+                {
+                    'id': org.id,
+                    'name': org.name,
+                    'description': org.description,
+                    'created_at': org.created_at.isoformat() if org.created_at else None,
+                    'updated_at': org.updated_at.isoformat() if org.updated_at else None,
+                }
+            )
+
+        client.close()
+        return {'status': 'success', 'organizations': result}
+
+    except Exception as e:
+        logger.error(f'Error listing InfluxDB organizations: {str(e)}')
+        return {'status': 'error', 'message': str(e)}
+
+
+@mcp.tool(name='InfluxDBCreateOrg', description='Create a new organization in InfluxDB.')
+async def influxdb_create_org(
+    org_name: str = Field(..., description='The name of the organization to create.'),
+    url: Optional[str] = OPTIONAL_FIELD_URL,
+    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
+    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
+    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
+) -> Dict[str, Any]:
+    """Create a new organization in InfluxDB.
+
+    Returns:
+        Details of the created organization.
+    """
+    if not tool_write_mode:
+        raise Exception(
+            'InfluxDBCreateOrg tool invocation not allowed when tool-write-mode is set to False'
+        )
+
+    resolved_url, resolved_token, _ = resolve_influxdb_config(
+        url, token, org=None, require_org=False
+    )
+
+    try:
+        client = get_influxdb_client(url=resolved_url, token=resolved_token, verify_ssl=verify_ssl)
+        orgs_api = client.organizations_api()
+
+        org = orgs_api.create_organization(name=org_name)
+
+        client.close()
+        return {
+            'status': 'success',
+            'organization': {
+                'id': org.id,
+                'name': org.name,
+                'created_at': org.created_at.isoformat() if org.created_at else None,
+            },
+        }
+
+    except Exception as e:
+        logger.error(f'Error creating InfluxDB organization: {str(e)}')
+        return {'status': 'error', 'message': str(e)}
+
+
 def main():
     """Main entry point for the MCP server application."""
     logger.info('Starting Timestream for InfluxDB MCP Server')

{awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info → awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: awslabs.timestream-for-influxdb-mcp-server
-Version: 0.0.11
+Version: 0.0.12
 Summary: An AWS Labs Model Context Protocol (MCP) server for Timestream for InfluxDB
 Project-URL: homepage, https://awslabs.github.io/mcp/
 Project-URL: docs, https://awslabs.github.io/mcp/servers/timestream-for-influxdb-mcp-server/
@@ -38,7 +38,8 @@ An AWS Labs Model Context Protocol (MCP) server for Timestream for InfluxDB. Thi
 - Create, update, list, describe, and delete Timestream for InfluxDB database clusters
 - Manage DB parameter groups
 - Tag management for Timestream for InfluxDB resources
-- Write and query data using InfluxDB's APIs
+- Manage InfluxDB 2 buckets and organizations
+- Write and query data using InfluxDB 2 APIs
 
 
 ## Pre-requisites
@@ -66,6 +67,9 @@ You can modify the settings of your MCP client to run your local server (e.g. fo
       "env": {
         "AWS_PROFILE": "your-aws-profile",
         "AWS_REGION": "us-east-1",
+        "INFLUXDB_URL": "https://your-influxdb-endpoint:8086",
+        "INFLUXDB_TOKEN": "your-influxdb-token",
+        "INFLUXDB_ORG": "your-influxdb-org",
         "FASTMCP_LOG_LEVEL": "ERROR"
       },
       "disabled": false,
@@ -94,9 +98,12 @@ For Windows users, the MCP server configuration format is slightly different:
         "awslabs.timestream-for-influxdb-mcp-server.exe"
       ],
       "env": {
-        "FASTMCP_LOG_LEVEL": "ERROR",
         "AWS_PROFILE": "your-aws-profile",
-        "AWS_REGION": "us-east-1"
+        "AWS_REGION": "us-east-1",
+        "INFLUXDB_URL": "https://your-influxdb-endpoint:8086",
+        "INFLUXDB_TOKEN": "your-influxdb-token",
+        "INFLUXDB_ORG": "your-influxdb-org",
+        "FASTMCP_LOG_LEVEL": "ERROR"
       }
     }
   }
@@ -146,3 +153,11 @@ The Timestream for InfluxDB MCP server provides the following tools:
 
 ##### Query API
 - `InfluxDBQuery`: Query data from InfluxDB using Flux query language
+
+##### Bucket Management
+- `InfluxDBListBuckets`: List all buckets in InfluxDB
+- `InfluxDBCreateBucket`: Create a new bucket in InfluxDB
+
+##### Organization Management
+- `InfluxDBListOrgs`: List all organizations in InfluxDB
+- `InfluxDBCreateOrg`: Create a new organization in InfluxDB

awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/RECORD

@@ -0,0 +1,9 @@
+awslabs/__init__.py,sha256=S-EDeDHGZzllMFhAeuJx_47fv85t2OUMOcsOisD5d_8,796
+awslabs/timestream_for_influxdb_mcp_server/__init__.py,sha256=xahTyEU632c290drFz2I-gBqmF6VJ_OvgyHyoBdHXqo,690
+awslabs/timestream_for_influxdb_mcp_server/server.py,sha256=-sNagSozw734-k_Sl6lbVEFi9ZuyzQmbkxtgZ3IHKTA,54325
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/METADATA,sha256=GHPTJTqhE5D7jhoohPkRXcqBlIA3QbKJYDT2UFrHpyg,7775
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/entry_points.txt,sha256=it4t34A91rT2btYRldC7Gg0iBXbFwK-8Uq8ug7F3GJ4,118
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/licenses/NOTICE,sha256=n2OVRXbV4EeTgFingxNqDwTL4yY1-mVTyT5QN8BeTGU,110
+awslabs_timestream_for_influxdb_mcp_server-0.0.12.dist-info/RECORD,,

awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/RECORD

@@ -1,9 +0,0 @@
-awslabs/__init__.py,sha256=S-EDeDHGZzllMFhAeuJx_47fv85t2OUMOcsOisD5d_8,796
-awslabs/timestream_for_influxdb_mcp_server/__init__.py,sha256=zIpXFDbkw_6_cyXI_2pk8l5tLQsDF9TH5b6EvVL0DCg,690
-awslabs/timestream_for_influxdb_mcp_server/server.py,sha256=V8IU0mhFj9ozZZ64ozXhYhPadskx3oRDdbI-zthYVTo,44306
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/METADATA,sha256=HiArTFwp6PwoDlDEXD6P5fJhW3z2mfEhs-MexK1zOuU,7129
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/entry_points.txt,sha256=it4t34A91rT2btYRldC7Gg0iBXbFwK-8Uq8ug7F3GJ4,118
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/licenses/NOTICE,sha256=n2OVRXbV4EeTgFingxNqDwTL4yY1-mVTyT5QN8BeTGU,110
-awslabs_timestream_for_influxdb_mcp_server-0.0.11.dist-info/RECORD,,