fustor-source-elasticsearch 0.1.4__py3-none-any.whl

fustor_source_elasticsearch/__init__.py

@@ -0,0 +1,311 @@
+"""
+Fuagent source driver for Elasticsearch.
+"""
+import logging
+import threading
+from typing import Any, Dict, Iterator, Tuple
+from datetime import datetime
+
+from elasticsearch import Elasticsearch, AsyncElasticsearch, AuthenticationException, AuthorizationException
+
+from fustor_core.drivers import SourceDriver
+from fustor_core.models.config import SourceConfig, PasswdCredential, ApiKeyCredential
+from fustor_core.exceptions import DriverError
+from fustor_event_model.models import EventBase, InsertEvent
+
+logger = logging.getLogger("fustor_agent.driver.elasticsearch")
+
+class ElasticsearchDriver(SourceDriver):
+    _instances: Dict[str, 'ElasticsearchDriver'] = {}
+    _lock = threading.Lock()
+    
+    def __new__(cls, id: str, config: SourceConfig):
+        # Generate unique signature: URI + credential to ensure permission isolation
+        signature = f"{config.uri}#{hash(str(config.credential))}"
+        
+        with ElasticsearchDriver._lock:
+            if signature not in ElasticsearchDriver._instances:
+                instance = super().__new__(cls)
+                ElasticsearchDriver._instances[signature] = instance
+            return ElasticsearchDriver._instances[signature]
+    
+    def __init__(self, id: str, config: SourceConfig):
+        # Prevent re-initialization of shared instances
+        if hasattr(self, '_initialized'):
+            return
+        
+        super().__init__(id, config)
+        self.uri = self.config.uri
+        self.credential = self.config.credential
+        self.driver_params = self.config.driver_params
+        
+        self._initialized = True
+
+    def _get_es_client(self) -> Elasticsearch:
+        return self._get_sync_es_client(self.uri, self.credential)
+
+    @staticmethod
+    def _get_sync_es_client(uri: str, credential: Any) -> Elasticsearch:
+        auth_params = {}
+        if isinstance(credential, PasswdCredential) and credential.user:
+            auth_params = (credential.user, credential.passwd or '')
+        elif isinstance(credential, ApiKeyCredential) and credential.key:
+            auth_params = {'api_key': credential.key}
+        
+        return Elasticsearch(
+            hosts=[uri], 
+            basic_auth=auth_params if isinstance(auth_params, tuple) else None, 
+            api_key=auth_params.get('api_key') if isinstance(auth_params, dict) else None
+        )
+
+    @staticmethod
+    async def _get_async_es_client(uri: str, credential: Any) -> AsyncElasticsearch:
+        auth_params = {}
+        if isinstance(credential, PasswdCredential) and credential.user:
+            auth_params = (credential.user, credential.passwd or '')
+        elif isinstance(credential, ApiKeyCredential) and credential.key:
+            auth_params = {'api_key': credential.key}
+
+        return AsyncElasticsearch(
+            hosts=[uri],
+            basic_auth=auth_params if isinstance(auth_params, tuple) else None,
+            api_key=auth_params.get('api_key') if isinstance(auth_params, dict) else None
+        )
+
+    def get_snapshot_iterator(self, **kwargs) -> Iterator[EventBase]:
+        index_name = self.driver_params.get("index_name")
+        timestamp_field = self.driver_params.get("timestamp_field")
+        if not index_name or not timestamp_field:
+            raise DriverError("'index_name' and 'timestamp_field' are required driver parameters.")
+
+        client = self._get_es_client()
+        logger.info(f"Starting snapshot for Elasticsearch index '{index_name}'.")
+        
+        pit = client.open_point_in_time(index=index_name, keep_alive="1m")
+        try:
+            search_after = None
+            snapshot_time = int(datetime.now().timestamp() * 1000)
+            while True:
+                resp = client.search(
+                    index=index_name, 
+                    size=kwargs.get("batch_size", 100), 
+                    sort=[{timestamp_field: "asc"}, {"_doc": "asc"}], 
+                    pit={"id": pit['id'], "keep_alive": "1m"}, 
+                    search_after=search_after
+                )
+                hits = resp['hits']['hits']
+                if not hits:
+                    break
+                
+                rows = [_normalize_doc(h) for h in hits]
+                fields = list(rows[0].keys()) if rows else []
+                yield InsertEvent(event_schema=index_name, table=index_name, rows=rows, fields=fields, index=snapshot_time)
+                
+                search_after = hits[-1]['sort']
+        finally:
+            if pit:
+                client.close_point_in_time(id=pit['id'])
+            logger.info(f"Snapshot for Elasticsearch index '{index_name}' finished.")
+
+    def is_position_available(self, position: int) -> bool:
+        """
+        Checks if the Elasticsearch position (timestamp) is available for resuming.
+        Since Elasticsearch doesn't have the same concept of positions as databases,
+        we generally consider positions available but may need to check if they're too old.
+        """
+        # For now, assume all positions are available as ES allows timestamp-based queries
+        # TODO In the future, we might check if the position is older than retention period
+        return True
+
+    def get_message_iterator(self, start_position: int=-1, **kwargs) -> Iterator[EventBase]:
+        
+        def _iterator_func() -> Iterator[EventBase]:
+            index_name = self.driver_params.get("index_name")
+            timestamp_field = self.driver_params.get("timestamp_field")
+            polling_interval = self.driver_params.get("polling_interval_sec", 5)
+            if not index_name or not timestamp_field:
+                raise DriverError("'index_name' and 'timestamp_field' are required driver parameters.")
+
+            client = self._get_es_client()
+            start_timestamp = start_position if start_position!=-1 else int(datetime.now().timestamp() * 1000)
+            last_ts_iso = datetime.fromtimestamp(start_timestamp / 1000).isoformat()
+            stop_event = kwargs.get("stop_event")
+
+            while not stop_event.is_set():
+                resp = client.search(
+                    index=index_name, 
+                    query={"range": {timestamp_field: {"gt": last_ts_iso}}},
+                    sort=[{timestamp_field: "asc"}], 
+                    size=100
+                )
+                hits = resp['hits']['hits']
+                if not hits:
+                    stop_event.wait(timeout=polling_interval)
+                    continue
+
+                for hit in hits:
+                    if stop_event.is_set():
+                        break
+                    doc = _normalize_doc(hit)
+                    ts_str = doc.get(timestamp_field)
+                    if ts_str:
+                        dt = datetime.fromisoformat(ts_str.replace('Z', '+00:00'))
+                        current_ts_ms = int(dt.timestamp() * 1000)
+                        yield UpdateEvent(event_schema=index_name, table=index_name, rows=[doc], fields=list(doc.keys()), index=current_ts_ms)
+                        last_ts_iso = dt.isoformat()
+        
+        return _iterator_func()
+
+    @classmethod
+    async def get_available_fields(cls, **kwargs) -> Dict[str, Any]:
+        uri = kwargs.get("uri")
+        credential_data = kwargs.get("credential")
+        driver_params = kwargs.get("driver_params", {})
+        index_name = driver_params.get("index_name")
+
+        if uri is None or credential_data is None or index_name is None:
+            raise DriverError("'uri', 'credential', and 'driver_params.index_name' are required.")
+
+        client = await cls._get_async_es_client(uri, credential_data)
+        try:
+            mapping = await client.indices.get_mapping(index=index_name)
+            properties = mapping[index_name]["mappings"].get("properties", {})
+            
+            flat_properties = {}
+            def flatten(props, prefix=''):
+                for key, value in props.items():
+                    if "properties" in value:
+                        flatten(value["properties"], f"{prefix}{key}.")
+                    else:
+                        flat_properties[f"{prefix}{key}"] = {"type": value.get("type", "object")}
+            
+            flatten(properties)
+            flat_properties["_id"] = {"type": "keyword"}
+            flat_properties["_index"] = {"type": "keyword"}
+
+            return {"properties": flat_properties}
+        except Exception as e:
+            logger.error(f"Failed to get available fields for ES index '{index_name}': {e}")
+            raise DriverError(f"Failed to get mapping for index '{index_name}': {e}")
+        finally:
+            await client.close()
+
+    @classmethod
+    async def test_connection(cls, **kwargs) -> Tuple[bool, str]:
+        uri = kwargs.get("uri")
+        credential_data = kwargs.get("credential")
+        if uri is None or credential_data is None:
+            return False, "'uri' and 'credential' are required."
+
+        try:
+            client = await cls._get_async_es_client(uri, credential_data)
+            if await client.ping():
+                return True, "Successfully connected to Elasticsearch."
+            else:
+                return False, "Connection to Elasticsearch failed."
+        except (AuthenticationException, AuthorizationException) as e:
+            return False, f"Authentication/Authorization failed: {e}"
+        except Exception as e:
+            return False, f"An unexpected error occurred: {e}"
+        finally:
+            if 'client' in locals():
+                await client.close()
+
+    @classmethod
+    async def check_privileges(cls, **kwargs) -> Tuple[bool, str]:
+        uri = kwargs.get("uri")
+        credential_data = kwargs.get("credential")
+        driver_params = kwargs.get("driver_params", {})
+        index_name = driver_params.get("index_name")
+
+        if uri is None or credential_data is None or index_name is None:
+            return False, "'uri', 'credential', and 'driver_params.index_name' are required."
+
+        client = await cls._get_async_es_client(uri, credential_data)
+        try:
+            response = await client.security.has_privileges(
+                body={"index": [{"names": [index_name], "privileges": ["read"]}]}
+            )
+            if response.get("has_all_requested"):
+                return True, "User has sufficient privileges for the index."
+            else:
+                return False, f"User lacks 'read' privilege for index '{index_name}'."
+        except Exception as e:
+            logger.warning(f"Could not verify privileges via security API (may not be enabled): {e}. Assuming success.")
+            return True, "Could not verify privileges via security API; assuming success."
+        finally:
+            await client.close()
+
+    @classmethod
+    async def get_wizard_steps(cls) -> Dict[str, Any]:
+        return {
+            "steps": [
+                {
+                    "step_id": "connection",
+                    "title": "Connection Details",
+                    "schema": {
+                        "type": "object",
+                        "properties": {
+                            "uri": {"type": "string", "title": "Elasticsearch URI"},
+                            "credential": {
+                                "type": "object",
+                                "title": "Credentials",
+                                "oneOf": [
+                                    {"$ref": "#/components/schemas/PasswdCredential"},
+                                    {"$ref": "#/components/schemas/ApiKeyCredential"}
+                                ]
+                            }
+                        },
+                        "required": ["uri", "credential"]
+                    },
+                    "validations": ["test_connection"]
+                },
+                {
+                    "step_id": "index_config",
+                    "title": "Index Configuration",
+                    "schema": {
+                        "type": "object",
+                        "properties": {
+                            "driver_params": {
+                                "type": "object",
+                                "title": "Driver Parameters",
+                                "properties": {
+                                    "index_name": {"type": "string", "title": "Index Name"},
+                                    "timestamp_field": {"type": "string", "title": "Timestamp Field"}
+                                },
+                                "required": ["index_name", "timestamp_field"]
+                            }
+                        },
+                        "required": ["driver_params"]
+                    },
+                    "validations": ["check_privileges", "discover_fields_no_cache"]
+                }
+            ],
+            "components": {
+                "schemas": {
+                    "PasswdCredential": {
+                        "type": "object",
+                        "title": "Username/Password",
+                        "properties": {
+                            "user": {"type": "string", "title": "Username"},
+                            "passwd": {"type": "string", "title": "Password", "format": "password"}
+                        },
+                        "required": ["user"]
+                    },
+                    "ApiKeyCredential": {
+                        "type": "object",
+                        "title": "API Key",
+                        "properties": {
+                            "key": {"type": "string", "title": "API Key", "format": "password"}
+                        },
+                        "required": ["key"]
+                    }
+                }
+            }
+        }
+
+def _normalize_doc(hit: Dict[str, Any]) -> Dict[str, Any]:
+    doc = hit.get('_source', {})
+    doc['_id'] = hit.get('_id')
+    doc['_index'] = hit.get('_index')
+    return doc

fustor_source_elasticsearch-0.1.4.dist-info/METADATA

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: fustor-source-elasticsearch
+Version: 0.1.4
+Summary: An Elasticsearch source for Fustor Agent
+Author-email: Huajin Wang <wanghuajin999@163.com>
+Requires-Python: >=3.11
+Requires-Dist: elasticsearch>=9.2.0
+Requires-Dist: fustor-core

fustor_source_elasticsearch-0.1.4.dist-info/RECORD

@@ -0,0 +1,7 @@
+fustor_source_elasticsearch/__init__.py,sha256=ww7iz8OK1XrzM9XxBOqTx8hd9fa35lDm5MZrLUl7Boc,13836
+fustor_source_elasticsearch/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fustor_source_elasticsearch-0.1.4.dist-info/METADATA,sha256=xC3-coD6pANIVGITzw52Uv6uOoYwHXOVq5WBkqqwRV8,258
+fustor_source_elasticsearch-0.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+fustor_source_elasticsearch-0.1.4.dist-info/entry_points.txt,sha256=6ra2kkIbMjl9GGlMiCZOEmQauXZ69eotLbhw6R19LWM,95
+fustor_source_elasticsearch-0.1.4.dist-info/top_level.txt,sha256=5yN8-ZDiTMdXAdqhUizGREqzj9bvxRz0MUZ_96bDzfg,28
+fustor_source_elasticsearch-0.1.4.dist-info/RECORD,,

fustor_source_elasticsearch-0.1.4.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

fustor_source_elasticsearch-0.1.4.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[fustor_agent.drivers.sources]
+elasticsearch = fustor_source_elasticsearch:ElasticsearchDriver

fustor_source_elasticsearch-0.1.4.dist-info/top_level.txt

@@ -0,0 +1 @@
+fustor_source_elasticsearch