python-databases 0.1.4__tar.gz → 0.1.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (21) hide show
  1. {python_databases-0.1.4/python_databases.egg-info → python_databases-0.1.6}/PKG-INFO +3 -2
  2. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases/elasticsearch_infrastructure/__init__.py +3 -3
  3. python_databases-0.1.6/python_databases/elasticsearch_infrastructure/elasticsearch.py +265 -0
  4. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases/postgresql_infrastructure/__init__.py +1 -1
  5. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases/postgresql_infrastructure/postgresql.py +23 -18
  6. {python_databases-0.1.4 → python_databases-0.1.6/python_databases.egg-info}/PKG-INFO +3 -2
  7. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases.egg-info/requires.txt +2 -1
  8. {python_databases-0.1.4 → python_databases-0.1.6}/requirements.txt +2 -1
  9. python_databases-0.1.6/setup.py +35 -0
  10. {python_databases-0.1.4 → python_databases-0.1.6}/tests/test_postgres.py +8 -6
  11. python_databases-0.1.4/python_databases/elasticsearch_infrastructure/elasticsearch.py +0 -320
  12. python_databases-0.1.4/setup.py +0 -37
  13. {python_databases-0.1.4 → python_databases-0.1.6}/LICENSE +0 -0
  14. {python_databases-0.1.4 → python_databases-0.1.6}/MANIFEST.in +0 -0
  15. {python_databases-0.1.4 → python_databases-0.1.6}/README.md +0 -0
  16. {python_databases-0.1.4 → python_databases-0.1.6}/pyproject.toml +0 -0
  17. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases/__init__.py +0 -0
  18. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases.egg-info/SOURCES.txt +0 -0
  19. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases.egg-info/dependency_links.txt +0 -0
  20. {python_databases-0.1.4 → python_databases-0.1.6}/python_databases.egg-info/top_level.txt +0 -0
  21. {python_databases-0.1.4 → python_databases-0.1.6}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: python-databases
3
- Version: 0.1.4
3
+ Version: 0.1.6
4
4
  Summary: A Python library for managing and interacting with multiple databases seamlessly.
5
5
  Home-page: https://github.com/aviz92/python-databases
6
6
  Author: Avi Zaguri
@@ -11,7 +11,8 @@ Description-Content-Type: text/markdown
11
11
  License-File: LICENSE
12
12
  Requires-Dist: setuptools
13
13
  Requires-Dist: wheel
14
- Requires-Dist: dotenv
14
+ Requires-Dist: python-dotenv
15
+ Requires-Dist: pre-commit
15
16
  Requires-Dist: custom-python-logger
16
17
  Requires-Dist: retrying
17
18
  Requires-Dist: elasticsearch<9
@@ -1,12 +1,12 @@
1
1
  from dotenv import load_dotenv
2
2
 
3
3
  from python_databases.elasticsearch_infrastructure.elasticsearch import (
4
- UrlProtocol,
5
4
  ElasticSearch,
5
+ ElasticSearchCloud,
6
6
  ElasticSearchOnPrem,
7
- ElasticSearchCloud
7
+ UrlProtocol,
8
8
  )
9
9
 
10
10
  load_dotenv()
11
11
 
12
- __all__ = ['UrlProtocol', 'ElasticSearch', 'ElasticSearchOnPrem', 'ElasticSearchCloud']
12
+ __all__ = ["UrlProtocol", "ElasticSearch", "ElasticSearchOnPrem", "ElasticSearchCloud"]
@@ -0,0 +1,265 @@
1
+ import logging
2
+ import time
3
+ from abc import ABC, abstractmethod
4
+ from copy import deepcopy
5
+ from datetime import UTC, datetime
6
+ from enum import Enum
7
+ from typing import Any
8
+
9
+ import pandas as pd
10
+ from custom_python_logger import get_logger
11
+ from elasticsearch import Elasticsearch, helpers
12
+ from retrying import retry
13
+
14
+
15
+ class UrlProtocol(Enum):
16
+ HTTP = "http"
17
+ HTTPS = "https"
18
+
19
+
20
+ class ElasticSearch(ABC):
21
+ def __init__(
22
+ self,
23
+ elk_hostname: str,
24
+ elasticsearch_port: int | None,
25
+ kibana_port: int | None,
26
+ protocol: UrlProtocol,
27
+ username: str | None,
28
+ password: str | None,
29
+ ) -> None:
30
+ self.logger = get_logger(self.__class__.__name__)
31
+ self._change_elasticsearch_logger()
32
+
33
+ self.protocol = protocol
34
+ self.elk_hostname = elk_hostname
35
+ self.username = username
36
+ self.password = password
37
+
38
+ self.elasticsearch_port = elasticsearch_port
39
+ if not self.elasticsearch_port:
40
+ self.elasticsearch_url = f"{self.protocol.value}://{self.elk_hostname}"
41
+ else:
42
+ self.elasticsearch_url = f"{self.protocol.value}://{self.elk_hostname}:{self.elasticsearch_port}"
43
+ self.kibana_port = kibana_port
44
+
45
+ self.elk_client = None
46
+
47
+ @abstractmethod
48
+ def connect_to_elasticsearch(self) -> None:
49
+ pass
50
+
51
+ @staticmethod
52
+ def _change_elasticsearch_logger(debug_level: int = logging.CRITICAL) -> None:
53
+ tracer = get_logger("elasticsearch")
54
+ tracer.setLevel(debug_level)
55
+
56
+ def _set_list_of_docs_quick(self, list_of_docs: list[dict], request_timeout: int = 600) -> None:
57
+ if failed_response := helpers.streaming_bulk(
58
+ self.elk_client, list_of_docs, raise_on_error=False, request_timeout=request_timeout, chunk_size=1000
59
+ ):
60
+ for item in failed_response:
61
+ if item[1]["index"]["status"] != 201 or item[1]["index"]["_shards"]["failed"] > 0:
62
+ self.logger.debug(f"Failed document: \n{item}")
63
+
64
+ def _set_list_of_docs_safe(self, list_of_docs: list[dict]) -> None:
65
+ success_response, failed_response = helpers.bulk(self.elk_client, list_of_docs)
66
+ if success_response == len(list_of_docs):
67
+ self.logger.info("All documents were reported successfully")
68
+ else:
69
+ self.logger.error(
70
+ f"Not all documents were reported successfully.\n"
71
+ f"The documents that were not reported successfully: {failed_response}"
72
+ )
73
+
74
+ def post_list_of_docs(self, list_of_docs: list[dict], request_timeout: int = 600, quick: bool = False) -> None:
75
+ self.logger.info("Start to report the documents to ELK")
76
+
77
+ if quick:
78
+ self._set_list_of_docs_quick(list_of_docs=list_of_docs, request_timeout=request_timeout)
79
+ else:
80
+ self._set_list_of_docs_safe(list_of_docs=list_of_docs)
81
+ self.logger.info("Finish to report the documents to ELK")
82
+
83
+ @staticmethod
84
+ def _get_doc_with_basic_info() -> dict:
85
+ doc = {}
86
+
87
+ timestamp = datetime.now(tz=UTC)
88
+ date_and_time_str = timestamp.isoformat()
89
+ date = date_and_time_str.split("T")[0]
90
+ time_part = date_and_time_str.split("T")[1].split("+")[0]
91
+
92
+ doc.update(
93
+ {
94
+ "doc_id": timestamp.strftime("%Y%m%dT%H%M%S%f"),
95
+ "timestamp": timestamp,
96
+ "date_and_time_str": date_and_time_str,
97
+ "date_str": date,
98
+ "time_str": time_part,
99
+ }
100
+ )
101
+ return doc
102
+
103
+ @staticmethod
104
+ def _add_list_values_as_str(doc: dict, column: str, value: Any) -> None:
105
+ """Add list values as both list and string representation to the provided doc"""
106
+ if isinstance(value, list) and all(isinstance(item, str) for item in value):
107
+ filtered_value = [item for item in value if item is not None]
108
+ doc[f"{column}"] = filtered_value
109
+ doc[f"{column}_str"] = ", ".join(filtered_value)
110
+
111
+ def _build_document(self, row: dict) -> dict:
112
+ doc = self._get_doc_with_basic_info()
113
+ for custom_field, value in row.items():
114
+ doc[custom_field] = value
115
+
116
+ self._add_list_values_as_str(
117
+ doc=doc,
118
+ column=custom_field,
119
+ value=value,
120
+ )
121
+ return doc
122
+
123
+ @staticmethod
124
+ def _fill_list_of_docs(index_name: str, doc: dict[str, Any]) -> dict:
125
+ return {"_index": index_name, "_source": deepcopy(doc)}
126
+
127
+ def _prepare_documents_for_bulk(self, data: list[dict], doc_index_name: str) -> list[dict]:
128
+ list_of_docs = []
129
+ for index, row in enumerate(data):
130
+ self.logger.debug(f"index: {index + 1}/{len(data)}")
131
+ doc = self._build_document(row=row)
132
+ list_of_docs.append(self._fill_list_of_docs(index_name=doc_index_name, doc=doc))
133
+ return list_of_docs
134
+
135
+ def post_list_of_docs_as_bulk_chunk(
136
+ self, list_of_docs: list[dict], chunk_size: int = 1000, time_sleep: int = 60, quick: bool = False
137
+ ) -> None:
138
+ for i in range(0, len(list_of_docs), chunk_size):
139
+ self.logger.info(f"index: {i} - {i + chunk_size} / {len(list_of_docs)}")
140
+ chunk = list_of_docs[i : i + chunk_size]
141
+ self.post_list_of_docs(list_of_docs=chunk, quick=quick)
142
+ time.sleep(time_sleep)
143
+
144
+ def fill_elk_index_as_bulk( # post_list_of_docs_as_bulk
145
+ self, data: list[dict], doc_index_name: str, chunk_size: int = 1000, time_sleep: int = 1, quick: bool = True
146
+ ) -> None: # use as the main function to send data to the elastic search
147
+ list_of_docs = self._prepare_documents_for_bulk(
148
+ data=data,
149
+ doc_index_name=doc_index_name,
150
+ )
151
+
152
+ self.post_list_of_docs_as_bulk_chunk(
153
+ list_of_docs=list_of_docs,
154
+ chunk_size=chunk_size,
155
+ time_sleep=time_sleep,
156
+ quick=quick,
157
+ )
158
+
159
+ def check_if_index_exists(self, index: str) -> bool:
160
+ if self.elk_client.indices.exists(index=index):
161
+ self.logger.info(f"The index '{index}' exists.")
162
+ return True
163
+ self.logger.info(f"The index '{index}' does not exist.")
164
+ return False
165
+
166
+ def delete_index(self, index: str) -> None:
167
+ self.elk_client.indices.delete(index=index)
168
+ self.logger.info(f"The index '{index}' was deleted.")
169
+
170
+ def check_if_index_exists_and_delete_if_exists(self, index: str) -> bool:
171
+ if self.check_if_index_exists(index=index):
172
+ self.delete_index(index=index)
173
+ return True
174
+ return False
175
+
176
+ def get_documents(self, index: str, query: dict) -> list[dict]:
177
+ try:
178
+ res = self.elk_client.search(index=index, body=query)["hits"]["hits"]
179
+ return [doc["_source"] for doc in res]
180
+ except Exception as e:
181
+ raise Exception(f"Failed to get documents from index '{index}': {str(e)}") from e
182
+
183
+ def convert_dataframes_to_list_of_docs(self, dataframe: pd.DataFrame) -> list:
184
+ data = []
185
+ dataframe_values = dataframe.values.tolist()
186
+ for index, value in enumerate(dataframe_values):
187
+ self.logger.debug(f"index: {index}/{len(dataframe_values)}")
188
+ data.append(value)
189
+ return data
190
+
191
+
192
+ class ElasticSearchOnPrem(ElasticSearch):
193
+ def __init__(
194
+ self,
195
+ elk_hostname: str,
196
+ elasticsearch_port: int | None = 9200,
197
+ kibana_port: int | None = 5602,
198
+ protocol: UrlProtocol = UrlProtocol.HTTPS,
199
+ username: str | None = None,
200
+ password: str | None = None,
201
+ ) -> None:
202
+ super().__init__(
203
+ elk_hostname=elk_hostname,
204
+ elasticsearch_port=elasticsearch_port,
205
+ kibana_port=kibana_port,
206
+ protocol=protocol,
207
+ username=username,
208
+ password=password,
209
+ )
210
+
211
+ self.logger = get_logger(self.__class__.__name__)
212
+
213
+ self.connect_to_elasticsearch()
214
+
215
+ @retry(stop_max_attempt_number=3, wait_fixed=180000)
216
+ def connect_to_elasticsearch(self) -> None:
217
+ if self.username and self.password:
218
+ self.elk_client = Elasticsearch(
219
+ hosts=[self.elasticsearch_url],
220
+ http_auth=(self.username, self.password),
221
+ # ca_certs="/etc/elasticsearch/certs/http_ca.crt",
222
+ verify_certs=False,
223
+ )
224
+ else:
225
+ self.elk_client = Elasticsearch(hosts=[f"{self.elasticsearch_url}"])
226
+
227
+ if self.elk_client.ping():
228
+ self.logger.info("Elasticsearch on-prem Connection Successful")
229
+ else:
230
+ self.logger.error("Elasticsearch on-prem Connection Failed")
231
+ raise Exception(f"Failed to connect to Elasticsearch on-prem on {self.elasticsearch_url}")
232
+
233
+
234
+ class ElasticSearchCloud(ElasticSearch):
235
+ def __init__(
236
+ self,
237
+ elk_hostname: str,
238
+ elasticsearch_port: int | None = 9200,
239
+ kibana_port: int | None = 5602,
240
+ protocol: UrlProtocol = UrlProtocol.HTTPS,
241
+ username: str | None = None,
242
+ password: str | None = None,
243
+ ) -> None:
244
+ super().__init__(
245
+ elk_hostname=elk_hostname,
246
+ elasticsearch_port=elasticsearch_port,
247
+ kibana_port=kibana_port,
248
+ protocol=protocol,
249
+ username=username,
250
+ password=password,
251
+ )
252
+
253
+ self.logger = get_logger(self.__class__.__name__)
254
+
255
+ self.connect_to_elasticsearch()
256
+
257
+ @retry(stop_max_attempt_number=3, wait_fixed=180000)
258
+ def connect_to_elasticsearch(self) -> None:
259
+ self.elk_client = Elasticsearch(cloud_id=self.elk_hostname, http_auth=(self.username, self.password))
260
+
261
+ if self.elk_client.ping():
262
+ self.logger.info("Elasticsearch cloud Connection Successful")
263
+ else:
264
+ self.logger.error("Elasticsearch cloud Connection Failed")
265
+ raise Exception(f"Failed to connect to Elasticsearch cloud on {self.elk_hostname}")
@@ -4,4 +4,4 @@ from python_databases.postgresql_infrastructure.postgresql import PostgreSQL
4
4
 
5
5
  load_dotenv()
6
6
 
7
- __all__ = ['PostgreSQL']
7
+ __all__ = ["PostgreSQL"]
@@ -1,21 +1,21 @@
1
- import logging
2
1
  import os
3
- from typing import Optional, Union, Any
2
+ from types import TracebackType
3
+ from typing import Any
4
+
4
5
  import psycopg2
5
6
  from custom_python_logger import get_logger
6
7
  from psycopg2 import extras
7
- from retrying import retry
8
8
 
9
9
 
10
10
  class PostgreSQL:
11
11
  def __init__(
12
12
  self,
13
- name: Optional[str] = None,
14
- host: Optional[str] = None,
15
- port: Optional[str] = None,
16
- username: Optional[str] = None,
17
- password: Optional[str] = None
18
- ):
13
+ name: str | None = None,
14
+ host: str | None = None,
15
+ port: str | None = None,
16
+ username: str | None = None,
17
+ password: str | None = None,
18
+ ) -> None:
19
19
  self.logger = get_logger(self.__class__.__name__)
20
20
 
21
21
  self.name = name or os.getenv("POSTGRESQL_NAME")
@@ -30,21 +30,21 @@ class PostgreSQL:
30
30
  self.create_connection()
31
31
 
32
32
  @property
33
- def client(self):
33
+ def client(self) -> psycopg2.extensions.connection:
34
34
  if not self._client or self._client.closed:
35
35
  self.logger.warning("Re-establishing PostgreSQL connection")
36
36
  self.create_connection()
37
37
  return self._client
38
38
 
39
39
  @property
40
- def cursor(self):
40
+ def cursor(self) -> psycopg2.extensions.cursor:
41
41
  if not self._cursor or self._cursor.closed:
42
42
  self.logger.warning("Re-creating PostgreSQL cursor")
43
43
  self.create_connection()
44
44
  return self._cursor
45
45
 
46
46
  # @retry(stop_max_attempt_number=3, wait_fixed=180000)
47
- def create_connection(self):
47
+ def create_connection(self) -> None:
48
48
  self.logger.info(f"Create connection to {self.host}:{self.port}")
49
49
 
50
50
  self._client = psycopg2.connect(
@@ -56,15 +56,15 @@ class PostgreSQL:
56
56
  )
57
57
  self._cursor = self._client.cursor(cursor_factory=extras.DictCursor)
58
58
 
59
- def execute_query(self, query: str, params: Optional[Union[tuple, dict[str, Any]]] = None) -> None:
59
+ def execute_query(self, query: str, params: tuple | dict[str, Any] | None = None) -> None:
60
60
  self.logger.debug(f"Fetch rows from DB by query {query} with params {params}")
61
61
  try:
62
62
  self.cursor.execute(query, params)
63
63
  except psycopg2.ProgrammingError as e:
64
64
  raise psycopg2.ProgrammingError(f"Error executing query: {query}") from e
65
- self.logger.debug(f"Finished to fetch rows by query")
65
+ self.logger.debug("Finished to fetch rows by query")
66
66
 
67
- def fetch_all(self, query: str, params: Optional[Union[tuple, dict[str, Any]]] = None) -> list[dict[str, Any]]:
67
+ def fetch_all(self, query: str, params: tuple | dict[str, Any] | None = None) -> list[dict[str, Any]]:
68
68
  self.execute_query(query, params)
69
69
  try:
70
70
  results = self.cursor.fetchall()
@@ -73,12 +73,17 @@ class PostgreSQL:
73
73
  self.logger.debug(f"Fetched {len(results)} rows")
74
74
  return results
75
75
 
76
- def close(self):
76
+ def close(self) -> None:
77
77
  self.cursor.close()
78
78
  self.client.close()
79
79
 
80
- def __enter__(self):
80
+ def __enter__(self) -> "PostgreSQL":
81
81
  return self
82
82
 
83
- def __exit__(self, exc_type, exc_val, exc_tb):
83
+ def __exit__(
84
+ self,
85
+ exc_type: type[BaseException] | None,
86
+ exc_val: BaseException | None,
87
+ exc_tb: TracebackType | None,
88
+ ) -> None:
84
89
  self.close()
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: python-databases
3
- Version: 0.1.4
3
+ Version: 0.1.6
4
4
  Summary: A Python library for managing and interacting with multiple databases seamlessly.
5
5
  Home-page: https://github.com/aviz92/python-databases
6
6
  Author: Avi Zaguri
@@ -11,7 +11,8 @@ Description-Content-Type: text/markdown
11
11
  License-File: LICENSE
12
12
  Requires-Dist: setuptools
13
13
  Requires-Dist: wheel
14
- Requires-Dist: dotenv
14
+ Requires-Dist: python-dotenv
15
+ Requires-Dist: pre-commit
15
16
  Requires-Dist: custom-python-logger
16
17
  Requires-Dist: retrying
17
18
  Requires-Dist: elasticsearch<9
@@ -1,6 +1,7 @@
1
1
  setuptools
2
2
  wheel
3
- dotenv
3
+ python-dotenv
4
+ pre-commit
4
5
  custom-python-logger
5
6
  retrying
6
7
  elasticsearch<9
@@ -1,6 +1,7 @@
1
1
  setuptools
2
2
  wheel
3
- dotenv
3
+ python-dotenv
4
+ pre-commit
4
5
  custom-python-logger
5
6
  retrying
6
7
  elasticsearch<9
@@ -0,0 +1,35 @@
1
+ from setuptools import find_packages, setup
2
+
3
+ package_version = "0.1.6"
4
+
5
+ package_name = "python-databases"
6
+ package_description = "A Python library for managing and interacting with multiple databases seamlessly."
7
+
8
+ package_name_ = package_name.replace("-", "_")
9
+ package_long_description_content_type = "text/markdown"
10
+ package_url = f"https://github.com/aviz92/{package_name}"
11
+ package_python_requires = ">=3.11"
12
+ package_author = "Avi Zaguri"
13
+
14
+ with open("requirements.txt") as file:
15
+ package_install_requires = [line.strip() for line in file.readlines() if line.strip() and not line.startswith("#")]
16
+
17
+ with open("README.md") as file:
18
+ package_long_description = file.read()
19
+
20
+ setup(
21
+ name=package_name,
22
+ version=package_version,
23
+ packages=find_packages(include=[package_name_, f"{package_name_}.*"]),
24
+ install_requires=package_install_requires,
25
+ author=package_author,
26
+ author_email="",
27
+ description=package_description,
28
+ long_description=package_long_description,
29
+ long_description_content_type=package_long_description_content_type,
30
+ url=package_url,
31
+ project_urls={
32
+ "Repository": package_url,
33
+ },
34
+ python_requires=package_python_requires,
35
+ )
@@ -1,7 +1,9 @@
1
1
  # test_postgresql.py
2
2
 
3
+ from unittest.mock import MagicMock, patch
4
+
3
5
  import pytest
4
- from unittest.mock import patch, MagicMock
6
+
5
7
  from python_databases.postgresql_infrastructure import PostgreSQL
6
8
 
7
9
 
@@ -15,21 +17,21 @@ def mock_connection():
15
17
  yield mock_conn, mock_cursor
16
18
 
17
19
 
18
- def test_connection_created(mock_connection):
20
+ def test_connection_created(mock_connection: MagicMock): # pylint: disable=W0621
19
21
  conn, _ = mock_connection
20
22
  db = PostgreSQL(name="test_db", host="localhost", port="5432", username="user", password="pass")
21
23
  assert db.client == conn
22
24
  assert conn.cursor.call_count == 2
23
25
 
24
26
 
25
- def test_execute_query_success(mock_connection):
27
+ def test_execute_query_success(mock_connection: MagicMock): # pylint: disable=W0621
26
28
  _, cursor = mock_connection
27
29
  db = PostgreSQL(name="test_db", host="localhost", port="5432", username="user", password="pass")
28
30
  db.execute_query("SELECT 1", None)
29
31
  cursor.execute.assert_called_with("SELECT 1", None)
30
32
 
31
33
 
32
- def test_fetch_all_success(mock_connection):
34
+ def test_fetch_all_success(mock_connection: MagicMock): # pylint: disable=W0621
33
35
  _, cursor = mock_connection
34
36
  cursor.fetchall.return_value = [{"id": 1}]
35
37
  db = PostgreSQL(name="test_db", host="localhost", port="5432", username="user", password="pass")
@@ -38,9 +40,9 @@ def test_fetch_all_success(mock_connection):
38
40
  assert results == [{"id": 1}]
39
41
 
40
42
 
41
- def test_context_manager_closes_connection(mock_connection):
43
+ def test_context_manager_closes_connection(mock_connection: MagicMock): # pylint: disable=W0621
42
44
  conn, cursor = mock_connection
43
- with PostgreSQL(name="test_db", host="localhost", port="5432", username="user", password="pass") as db:
45
+ with PostgreSQL(name="test_db", host="localhost", port="5432", username="user", password="pass") as _:
44
46
  pass
45
47
  cursor.close.assert_called_once()
46
48
  conn.close.assert_called_once()
@@ -1,320 +0,0 @@
1
- import logging
2
- import time
3
- from abc import ABC, abstractmethod
4
- from copy import deepcopy
5
- from datetime import datetime, timezone
6
- from enum import Enum
7
- from typing import Optional, Any
8
- import pandas as pd
9
- from custom_python_logger import get_logger
10
- from elasticsearch import Elasticsearch, helpers
11
- from retrying import retry
12
-
13
-
14
- class UrlProtocol(Enum):
15
- HTTP = 'http'
16
- HTTPS = 'https'
17
-
18
-
19
- class ElasticSearch(ABC):
20
- def __init__(
21
- self,
22
- elk_hostname: str,
23
- elasticsearch_port: Optional[int],
24
- kibana_port: Optional[int],
25
- protocol: UrlProtocol,
26
- username: Optional[str],
27
- password: Optional[str]
28
- ):
29
- self.logger = get_logger(self.__class__.__name__)
30
- self._change_elasticsearch_logger()
31
-
32
- self.protocol = protocol
33
- self.elk_hostname = elk_hostname
34
- self.username = username
35
- self.password = password
36
-
37
- self.elasticsearch_port = elasticsearch_port
38
- if not self.elasticsearch_port:
39
- self.elasticsearch_url = f'{self.protocol.value}://{self.elk_hostname}'
40
- else:
41
- self.elasticsearch_url = f'{self.protocol.value}://{self.elk_hostname}:{self.elasticsearch_port}'
42
- self.kibana_port = kibana_port
43
-
44
- self.elk_client = None
45
-
46
- def _change_elasticsearch_logger(self):
47
- tracer = get_logger('elasticsearch')
48
- tracer.setLevel(logging.CRITICAL) # or desired level
49
- # tracer.addHandler(logging.FileHandler('indexer.log'))
50
-
51
- @abstractmethod
52
- def connect_to_elasticsearch(self):
53
- pass
54
-
55
- def set_list_of_docs(self, list_of_docs: list[dict], request_timeout: int = 600, quick: bool = True) -> None:
56
- self.logger.info('Start to report the documents to ELK')
57
-
58
- if quick:
59
- if failed_response := helpers.streaming_bulk(
60
- self.elk_client,
61
- list_of_docs,
62
- raise_on_error=False,
63
- request_timeout=request_timeout,
64
- chunk_size=1000
65
- ):
66
- for item in failed_response:
67
- if item[1]['index']['status'] != 201 or item[1]['index']['_shards']['failed'] > 0:
68
- self.logger.debug(f'Failed document: \n{item}')
69
- else:
70
- success_response, failed_response = helpers.bulk(self.elk_client, list_of_docs)
71
- if success_response == len(list_of_docs):
72
- self.logger.info('All documents were reported successfully')
73
- else:
74
- self.logger.error(
75
- f'Not all documents were reported successfully.\n'
76
- f'The documents that were not reported successfully: {failed_response}'
77
- )
78
- self.logger.info('Finish to report the documents to ELK')
79
-
80
- def fill_list_of_docs( # noqa
81
- self,
82
- list_of_docs: list[dict],
83
- index_name: str,
84
- doc: dict[str, Any]
85
- ) -> None:
86
- list_of_docs.append(
87
- {
88
- "_index": index_name,
89
- "_source": deepcopy(doc)
90
- }
91
- )
92
-
93
- def _add_basic_info(
94
- self,
95
- doc: dict,
96
- doc_id: str,
97
- timestamp: datetime,
98
- date_and_time: str,
99
- username: str = None
100
- ) -> dict:
101
- doc.update(
102
- {
103
- 'doc_id': doc_id,
104
- 'timestamp': timestamp,
105
- 'date_str': f'{date_and_time}'.replace('-', '_')
106
- }
107
- )
108
- if username:
109
- doc['username'] = username
110
- return doc
111
-
112
- def _add_list_values_as_str(self, doc: dict, value: Any, column: str) -> None:
113
- """Add list values as both list and string representation to the provided doc"""
114
- if type(value) is list and all(type(item) is str for item in value):
115
- filtered_value = list(filter(lambda item: item is not None, value))
116
- doc[f'{column}'] = filtered_value
117
- doc[f'{column}_str'] = ', '.join(filtered_value)
118
-
119
- def _build_document(
120
- self,
121
- row: dict,
122
- doc_id: str,
123
- timestamp: datetime,
124
- date_and_time: str,
125
- username: str = None
126
- ) -> dict:
127
- doc = self._add_basic_info(
128
- doc={},
129
- doc_id=doc_id,
130
- timestamp=timestamp,
131
- date_and_time=date_and_time,
132
- username=username
133
- )
134
-
135
- # Add custom fields
136
- for custom_field, value in row.items():
137
- doc[custom_field] = value
138
- self._add_list_values_as_str(
139
- doc=doc,
140
- value=value,
141
- column=custom_field
142
- )
143
-
144
- return doc
145
-
146
- def _prepare_documents_for_bulk(
147
- self,
148
- data: list[dict],
149
- doc_id: str,
150
- doc_index_name: str,
151
- timestamp: datetime,
152
- date_and_time: str,
153
- username: str = None
154
- ) -> list[dict]:
155
- list_of_docs = []
156
- for index, row in enumerate(data):
157
- self.logger.debug(f'index: {index + 1}/{len(data)}')
158
-
159
- doc = self._build_document(
160
- row=row,
161
- doc_id=doc_id,
162
- timestamp=timestamp,
163
- date_and_time=date_and_time,
164
- username=username
165
- )
166
-
167
- self.fill_list_of_docs(
168
- list_of_docs=list_of_docs,
169
- index_name=doc_index_name,
170
- doc=doc
171
- )
172
-
173
- return list_of_docs
174
-
175
- def fill_elk_index_as_bulk_chunk(self, list_of_docs: list[dict], chunk_size=1000, time_sleep: int = 60) -> None:
176
- for i in range(0, len(list_of_docs), chunk_size):
177
- self.logger.info(f'index: {i} - {i + chunk_size} / {len(list_of_docs)}')
178
- chunk = list_of_docs[i:i + chunk_size]
179
- self.set_list_of_docs(list_of_docs=chunk)
180
- time.sleep(time_sleep)
181
-
182
- def fill_elk_index_as_bulk(
183
- self,
184
- data: list[dict],
185
- doc_index_name: str,
186
- doc_id: Optional[str] = None,
187
- timestamp: datetime = None,
188
- date_and_time: Optional[str] = None,
189
- username: Optional[str] = None,
190
- chunk_size=1000,
191
- time_sleep: int = 1
192
- ) -> None: # use as the main function to send data to the elastic search
193
- timestamp = timestamp or datetime.now(tz=timezone.utc)
194
- list_of_docs = self._prepare_documents_for_bulk(
195
- data=data,
196
- doc_index_name=doc_index_name,
197
- timestamp=timestamp,
198
- doc_id=doc_id or timestamp.strftime('%Y_%m_%dT%H_%M_%S_%f'),
199
- date_and_time=date_and_time or timestamp.isoformat(),
200
- username=username
201
- )
202
-
203
- self.fill_elk_index_as_bulk_chunk(
204
- list_of_docs=list_of_docs,
205
- chunk_size=chunk_size,
206
- time_sleep=time_sleep
207
- )
208
-
209
- def convert_dataframes_to_list_of_docs(self, dataframe: pd.DataFrame) -> list:
210
- data = []
211
- dataframe_values = dataframe.values.tolist()
212
- for index, value in enumerate(dataframe_values):
213
- self.logger.debug(f'index: {index}/{len(dataframe_values)}')
214
- data.append(value)
215
- return data
216
-
217
- def check_if_index_exists(self, index: str) -> bool:
218
- if self.elk_client.indices.exists(index=index):
219
- self.logger.info(f"The index '{index}' exists.")
220
- return True
221
- else:
222
- self.logger.info(f"The index '{index}' does not exist.")
223
- return False
224
-
225
- def delete_index(self, index: str) -> None:
226
- self.elk_client.indices.delete(index=index)
227
- self.logger.info(f"The index '{index}' was deleted.")
228
-
229
- def check_if_index_exists_and_delete_if_exists(self, index: str) -> bool:
230
- if self.check_if_index_exists(index=index):
231
- self.delete_index(index=index)
232
- return True
233
- else:
234
- self.logger.info(f"The index '{index}' does not exist.")
235
- return False
236
-
237
- def get_documents(self, index: str, query: dict) -> list[dict]:
238
- try:
239
- res = self.elk_client.search(index=index, body=query)["hits"]["hits"]
240
- return [doc["_source"] for doc in res]
241
- except Exception as e:
242
- raise Exception(f"Failed to get documents from index '{index}': {str(e)}")
243
-
244
-
245
- class ElasticSearchOnPrem(ElasticSearch):
246
- def __init__(
247
- self,
248
- elk_hostname: str,
249
- elasticsearch_port: Optional[int] = 9200,
250
- kibana_port: Optional[int] = 5602,
251
- protocol: UrlProtocol = UrlProtocol.HTTPS,
252
- username: Optional[str] = None,
253
- password: Optional[str] = None
254
- ):
255
- super().__init__(
256
- elk_hostname=elk_hostname,
257
- elasticsearch_port=elasticsearch_port,
258
- kibana_port=kibana_port,
259
- protocol=protocol,
260
- username=username,
261
- password=password
262
- )
263
-
264
- self.logger = get_logger(self.__class__.__name__)
265
-
266
- self.connect_to_elasticsearch()
267
-
268
- @retry(stop_max_attempt_number=3, wait_fixed=180000)
269
- def connect_to_elasticsearch(self):
270
- if self.username and self.password:
271
- self.elk_client = Elasticsearch(
272
- hosts=[self.elasticsearch_url],
273
- http_auth=(self.username, self.password),
274
- # ca_certs="/etc/elasticsearch/certs/http_ca.crt",
275
- verify_certs=False, timeout=60
276
- )
277
- else:
278
- self.elk_client = Elasticsearch(hosts=[f'{self.elasticsearch_url}'], timeout=20)
279
-
280
- if self.elk_client.ping():
281
- self.logger.info("Elasticsearch on-prem Connection Successful")
282
- else:
283
- self.logger.error("Elasticsearch on-prem Connection Failed")
284
- raise Exception(
285
- f'Failed to connect to Elasticsearch on-prem on {self.elasticsearch_url}')
286
-
287
-
288
- class ElasticSearchCloud(ElasticSearch):
289
- def __init__(
290
- self,
291
- elk_hostname: str,
292
- elasticsearch_port: Optional[int] = 9200,
293
- kibana_port: Optional[int] = 5602,
294
- protocol: UrlProtocol = UrlProtocol.HTTPS,
295
- username: Optional[str] = None,
296
- password: Optional[str] = None
297
- ):
298
- super().__init__(
299
- elk_hostname=elk_hostname,
300
- elasticsearch_port=elasticsearch_port,
301
- kibana_port=kibana_port,
302
- protocol=protocol,
303
- username=username,
304
- password=password
305
- )
306
-
307
- self.logger = get_logger(self.__class__.__name__)
308
-
309
- self.connect_to_elasticsearch()
310
-
311
- @retry(stop_max_attempt_number=3, wait_fixed=180000)
312
- def connect_to_elasticsearch(self):
313
- self.elk_client = Elasticsearch(cloud_id=self.elk_hostname, http_auth=(self.username, self.password))
314
-
315
- if self.elk_client.ping():
316
- self.logger.info("Elasticsearch cloud Connection Successful")
317
- else:
318
- self.logger.error("Elasticsearch cloud Connection Failed")
319
- raise Exception(
320
- f'Failed to connect to Elasticsearch cloud on {self.elk_hostname}')
@@ -1,37 +0,0 @@
1
- from setuptools import setup, find_packages
2
-
3
- package_version = '0.1.4'
4
-
5
- package_name = 'python-databases'
6
- package_description = 'A Python library for managing and interacting with multiple databases seamlessly.'
7
-
8
- package_name_ = package_name.replace('-', '_')
9
- package_long_description_content_type = 'text/markdown'
10
- package_url = f'https://github.com/aviz92/{package_name}'
11
- package_python_requires = '>=3.11'
12
- package_author = 'Avi Zaguri'
13
-
14
- with open('requirements.txt', 'r') as file:
15
- package_install_requires = [
16
- line.strip() for line in file.readlines() if line.strip() and not line.startswith('#')
17
- ]
18
-
19
- with open('README.md', 'r') as file:
20
- package_long_description = file.read()
21
-
22
- setup(
23
- name=package_name,
24
- version=package_version,
25
- packages=find_packages(include=[package_name_, f'{package_name_}.*']),
26
- install_requires=package_install_requires,
27
- author=package_author,
28
- author_email='',
29
- description=package_description,
30
- long_description=package_long_description,
31
- long_description_content_type=package_long_description_content_type,
32
- url=package_url,
33
- project_urls={
34
- 'Repository': package_url,
35
- },
36
- python_requires=package_python_requires,
37
- )