armandmasseaugit opened a new pull request, #54259:
URL: https://github.com/apache/airflow/pull/54259
# Enhanced Elasticsearch provider with comprehensive ElasticsearchHook
## Summary
This PR introduces a new comprehensive `ElasticsearchHook` that provides
enhanced functionality for interacting with Elasticsearch clusters, with better
Airflow integration and additional features.
## What changed
### New Features
- **Enhanced ElasticsearchHook**: Complete rewrite with comprehensive
functionality
- **Airflow Connection Integration**: Full support for Airflow connections
with environment variable fallback
- **Bulk Operations**: Support for `bulk()`, `streaming_bulk()`, and
`parallel_bulk()` operations
- **Index Management**: Methods for creating, deleting, and checking index
existence
- **Pandas Integration**: `search_to_pandas()` and `scan_to_pandas()`
methods with lazy imports
- **Context Manager Support**: Automatic connection cleanup with `with`
statement
- **Data Migration**: `reindex()` method for moving data between indices
- **Enhanced Error Handling**: Better connection testing and error reporting
### Deprecations
- **ElasticsearchPythonHook**: Marked as deprecated in favor of the new
`ElasticsearchHook`
### Configuration Improvements
The new hook supports a three-tier configuration priority:
1. Airflow connections (primary)
2. Environment variables (fallback)
3. Default values (last resort)
Supported environment variables:
- `ELASTICSEARCH_HOST`
- `ELASTICSEARCH_PORT`
- `ELASTICSEARCH_USERNAME`
- `ELASTICSEARCH_PASSWORD`
- `ELASTICSEARCH_USE_SSL`
- `ELASTICSEARCH_VERIFY_CERTS`
- `ELASTICSEARCH_TIMEOUT`
- `ELASTICSEARCH_MAX_RETRIES`
## API Reference
### New Methods
```python
# Connection and testing
hook.test_connection() -> bool
hook.get_conn() -> Elasticsearch
# Search operations
hook.search(query, index_name, **kwargs) -> Dict
hook.scan(index, query, **kwargs) -> Generator
# Pandas integration (lazy imports)
hook.search_to_pandas(index, query, **kwargs) -> pd.DataFrame
hook.scan_to_pandas(index, query, **kwargs) -> pd.DataFrame
# Bulk operations
hook.bulk(actions, **kwargs) -> Tuple[int, List]
hook.streaming_bulk(actions, **kwargs) -> Generator
hook.parallel_bulk(actions, **kwargs) -> Generator
# Index management
hook.create_index(index_name, mappings, settings, **kwargs) -> Dict
hook.delete_index(index_name, **kwargs) -> Dict
hook.index_exists(index_name) -> bool
# Data migration
hook.reindex(source_index, target_index, query, **kwargs) -> Tuple[int, List]
# Context manager
with ElasticsearchHook() as hook:
# Automatic cleanup
```
## Usage Examples
### Basic Usage
```python
from airflow.providers.elasticsearch.hooks.elasticsearch import
ElasticsearchHook
# Using default connection
hook = ElasticsearchHook()
# Test connection
if hook.test_connection():
print("Connected successfully!")
```
### Search with Pandas
```python
# Convert search results to DataFrame
df = hook.search_to_pandas(
index="my_index",
query={"match_all": {}}
)
```
### Bulk Operations
```python
actions = [
{"_index": "my_index", "_source": {"field": "value1"}},
{"_index": "my_index", "_source": {"field": "value2"}}
]
success_count, errors = hook.bulk(actions)
```
### Index Management
```python
# Create index with mappings
hook.create_index(
index_name="my_index",
mappings={"properties": {"field": {"type": "text"}}},
settings={"number_of_shards": 1}
)
```
## Breaking Changes
- `ElasticsearchPythonHook` is now deprecated (with runtime warning)
- New users should use `ElasticsearchHook` instead
## Migration Guide
**Before (deprecated):**
```python
from airflow.providers.elasticsearch.hooks.elasticsearch import
ElasticsearchPythonHook
hook = ElasticsearchPythonHook(hosts=["http://localhost:9200";])
```
**After (recommended):**
```python
from airflow.providers.elasticsearch.hooks.elasticsearch import
ElasticsearchHook
hook = ElasticsearchHook(elasticsearch_conn_id="elasticsearch_default")
```
## Tests
- ✅ All existing tests pass
- ✅ Added comprehensive unit tests for new `ElasticsearchHook`
- ✅ Tests cover all new methods and error scenarios
- ✅ Mock-based testing to avoid external dependencies
## Documentation
- ✅ Added comprehensive documentation in
`providers/elasticsearch/docs/hooks/elasticsearch_hook.rst`
- ✅ API reference with detailed method documentation
- ✅ Migration guide from deprecated hook
## Backwards Compatibility
- ✅ All existing functionality remains unchanged
- ✅ `ElasticsearchPythonHook` continues to work (with deprecation warning)
- ✅ No breaking changes for current users
## Dependencies
- ✅ pandas import is lazy (only when needed)
- ✅ No additional required dependencies
- ✅ Optional pandas dependency clearly documented
## Related Issues
None
This enhancement provides a more comprehensive and Airflow-native way to
interact with Elasticsearch, making it easier for users to integrate
Elasticsearch operations into their DAGs while maintaining backwards
compatibility.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
