DebenOldert 2 роки тому
батько
коміт
df731d5177
5 змінених файлів з 300 додано та 1 видалено
  1. 1 1
      .gitignore
  2. 81 0
      dbi.py
  3. 69 0
      mongodb.py
  4. 146 0
      mysql.py
  5. 3 0
      requirements.txt

+ 1 - 1
.gitignore

@@ -343,5 +343,5 @@ _Pvt_Extensions
 
 # ---> VisualStudioCode
 .settings
-
+.vscode
 

+ 81 - 0
dbi.py

@@ -0,0 +1,81 @@
+from typing import Dict, List, Union, Tuple, Optional, Any, Final, NamedTuple
+import abc
+from dataclasses import dataclass
+import os
+from collections import namedtuple
+
+import pandas as pd
+
+
+class DBI(abc.ABC):
+    OUTPUT = print
+    OUTPUT_KWARGS: dict = {}
+
+    BACKEND = None
+    #WriteReturn: Final = namedtuple('WriteReturn', ['rows', 'id'])
+
+    @dataclass
+    class WriteReturn:
+        rows: int
+        key: Optional[Any]
+
+    def __init__(self, 
+                 host: str, 
+                 user: Optional[str] = None, 
+                 password: Optional[str] = None, 
+                 scheme: Optional[str] = None):
+        self.host: str = os.environ.get(host, 'localhost')
+        self.user: Optional[str] = os.environ.get(user or '', None)
+        self.password: Optional[str] = os.environ.get(password or '', None)
+        self.scheme: Optional[str] = os.environ.get(scheme or '', None)
+
+        self.OUTPUT(f'Connecting to {self.host} -> {self.scheme if self.scheme else "*"}', 
+                    **self.OUTPUT_KWARGS)
+
+    @abc.abstractmethod
+    def write(self, 
+              query) \
+              -> WriteReturn:
+        """Execute a query that writes data to the database (INSERT, UPDATE, DELETE)"""
+        pass
+
+    @abc.abstractmethod
+    def read(self, 
+             query, 
+             single_row=False) \
+             -> Union[pd.DataFrame, Optional[pd.Series]]:
+        """Execute a query that reads data from the database (SELECT, SHOW)"""
+        pass
+
+    @abc.abstractmethod
+    def insert(self, 
+               table: str, 
+               rows: Union[pd.DataFrame, Dict[str, Any]], 
+               ignore: bool = False, 
+               update: bool = False, 
+               schema: Optional[str] = None) \
+               -> WriteReturn:
+        """Prepared function to insert new data to a database"""
+        pass
+    
+    @abc.abstractmethod
+    def update(self, 
+               table: str, 
+               ids: Dict[str, Any], 
+               values: Dict[str, Any], 
+               schema: Optional[str] = None) \
+               -> WriteReturn:
+        """Prepared function to update data in a database"""
+        pass
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Close the database connection"""
+        self.OUTPUT(f'Closing connection to {self.host} -> {self.scheme}', 
+                    **self.OUTPUT_KWARGS)
+
+    @abc.abstractmethod
+    def reconnect(self) -> None:
+        """Reconnect the database connection"""
+        self.OUTPUT(f'Reconnecting to {self.host} -> {self.scheme}', 
+                    **self.OUTPUT_KWARGS)

+ 69 - 0
mongodb.py

@@ -0,0 +1,69 @@
+class MongoDB(Database):
+    class Query:
+        def __init__(self, table: str, select: Dict = None, condition: Dict = None, data: List[Dict] = None):
+            self.table = table
+            self.select = select
+            self.condition = condition
+            self.data = data
+
+    def __init__(self, host: str = 'localhost', user: str = None, password: str = None, scheme: str = ''):
+        super(MongoDB, self).__init__(host, user, password, scheme)
+
+        self.conn = pymongo.MongoClient(f'mongodb://{host}:27017/{scheme or ""}')
+
+    def write(self, query: "MongoDB.Query", single_insert: bool = True) -> bool:
+        super(MongoDB, self).write(query)
+
+        db = self.conn[self.scheme][query.table]
+
+        if single_insert:
+            res = db.insert_one(query.data[0])
+            return res is not None
+        else:
+            res = db.insert_many(query.data)
+            return len(res.inserted_ids) == len(query.data)
+
+    def read_raw(self, query: "MongoDB.Query") -> pymongo.cursor.Cursor:
+        db = self.conn[self.scheme][query.table]
+
+        res = db.find(query.condition, query.select)
+
+        return res
+
+    def read(self, query: "MongoDB.Query", single_row: bool = False) -> Optional[pd.DataFrame]:
+        super(MongoDB, self).read(query, single_row)
+
+        db = self.conn[self.scheme][query.table]
+
+        res = db.find(query.condition, query.select)
+
+        return self.raw_to_dataframe(res, single_row=single_row)
+
+    def update(self, query: "MongoDB.Query") -> bool:
+        db = self.conn[self.scheme][query.table]
+
+        res = db.update_many(query.condition, {'$push': query.data[0]}, upsert=True)
+
+        return res.matched_count > 0
+
+    def close(self) -> None:
+        super(MongoDB, self).close()
+        self.conn.close()
+
+    def reconnect(self) -> None:
+        super(MongoDB, self).reconnect()
+
+    @staticmethod
+    def raw_to_dataframe(raw: pymongo.cursor.Cursor, single_row: bool = False) -> Optional[pd.DataFrame]:
+        if raw is None:
+            return None
+
+        data = list(raw)
+
+        if len(data) == 0:
+            return None
+
+        if single_row:
+            return pd.DataFrame(data).iloc[0]
+        else:
+            return pd.DataFrame(data)

+ 146 - 0
mysql.py

@@ -0,0 +1,146 @@
+from typing import Dict, Hashable, Iterable, List, MutableMapping, Union, Tuple, Optional, Any
+
+from dbi import DBI, pd
+
+class MySQL(DBI):
+    def __init__(self, 
+                 host: str, 
+                 user: Optional[str] = None, 
+                 password: Optional[str] = None, 
+                 scheme: Optional[str] = None):
+        super().__init__(host, user, password, scheme)
+        import mysql.connector
+
+        self.BACKEND = mysql
+
+        self.conn = self.BACKEND.connector.connect(host=self.host, 
+                                                   user=self.user, 
+                                                   password=self.password, 
+                                                   database=self.scheme)
+        self.cursor = self.conn.cursor()
+
+    def write(self, 
+              query: str) \
+              -> DBI.WriteReturn:
+        try:
+            self.cursor.execute(query)
+            self.conn.commit()
+            return DBI.WriteReturn(rows=self.cursor.rowcount,
+                                   key=self.cursor.lastrowid)
+        except (self.BACKEND.connector.InterfaceError, 
+                self.BACKEND.connector.OperationalError):
+            self.reconnect()
+            return self.write(query)
+        except Exception as e:
+            raise Exception(f'MySQL raised an error while writing: {e.args}. SQL: {query}')
+
+    def insert(self, 
+               table: str, 
+               rows: Union[pd.DataFrame, Dict[str, Any], List[Dict[str, Any]]], 
+               ignore: bool = False, 
+               update: bool = False, 
+               schema: Optional[str] = None) \
+               -> DBI.WriteReturn:
+        sql: str = f'INSERT {"IGNORE" if ignore else ""} INTO {"`"+schema+"`." if schema else ""}`{table}` '
+
+        prep_rows: List[Dict[str, Any]] = []
+        if type(rows) is dict:
+            columns: Iterable[str] = rows.keys()
+            prep_rows = [rows]
+        elif type(rows) is pd.DataFrame:
+            columns = rows.columns
+            prep_rows = rows.to_dict(orient='records', index=False) # type: ignore
+        elif type(rows) is list:
+            columns = rows[0].keys()
+            prep_rows = rows
+        else:
+            raise ValueError(f'Cannot extract columns from type "{type(rows).__qualname__}"')
+        
+        sql += f'({",".join([f"`{c}`" for c in columns])}) VALUES '
+
+        # Create each row entry
+        _data: List[str] = []
+        for row in prep_rows:
+            _row: List[str] = []
+            for column in columns:
+                _col = row.get(column)
+                _row.append(self._parse_value(_col))
+
+            _data.append(f'({",".join(_row)})')
+
+        # Merge row entries to single entry
+        sql += f'{",".join(_data)} AS `new` '
+
+        if update:
+            sql += f'ON DUPLICATE KEY UPDATE {",".join([f"`{c}`=`new`.`{c}`" for c in columns])}'
+
+        return self.write(sql)
+
+    def update(self, 
+               table: str, 
+               ids: Dict[str, Any], 
+               values: Dict[str, Any], 
+               schema: Optional[str] = None) \
+               -> DBI.WriteReturn:
+        sql: str = f'UPDATE {"`"+schema+"`." if schema else ""}`{table}` '
+
+        _val: List[str] = []
+        for key in values.keys():
+            _val.append(f'`{key}`={self._parse_value(values.get(key))}')
+
+        sql += f'SET {",".join(_val)} '
+
+        _id: List[str] = []
+
+        for key in ids.keys():
+            _id.append(f'`{key}`={self._parse_value(ids.get(key))}')
+
+        sql += f'WHERE {" AND ".join(_id)}'
+
+        return self.write(sql)
+
+    def read(self, 
+             query: str, 
+             single_row: bool = False) \
+             -> Union[pd.DataFrame, Optional[pd.Series]]:
+        try:
+            self.cursor.execute(query)
+            df = pd.DataFrame(columns=self.cursor.column_names, data=self.cursor.fetchall())
+            if single_row:
+                return df.iloc[0] if len(df) > 0 else None
+            return df
+        except (self.BACKEND.connector.InterfaceError, 
+                self.BACKEND.connector.OperationalError):
+            self.reconnect()
+            return self.read(query, single_row)
+        except Exception as e:
+            raise Exception(f'MySQL raised an error while reading: {e.args}. SQL: {query}')
+
+    def commit(self) -> None:
+        self.conn.commit()
+
+    def close(self) -> None:
+        self.OUTPUT(f'Closing connection to {self.host} -> {self.scheme}', **self.OUTPUT_KWARGS)
+        self.conn.close()
+
+    def reconnect(self) -> None:
+        self.OUTPUT(f'Lost connection to database {self.scheme}. Reconnecting...', **self.OUTPUT_KWARGS)
+        self.conn.reconnect(20, 5)
+        self.OUTPUT(f'Reconnect successful', **self.OUTPUT_KWARGS)
+
+    def change_scheme(self, scheme) -> None:
+        self.close()
+        self.__init__(host=self.host, user=self.user, password=self.password, scheme=scheme)
+
+    @staticmethod
+    def _parse_value(value) -> str:
+        if value is None:
+            return 'NULL'
+        elif type(value) is float:
+            return f'{value:12f}'
+        elif type(value) is int:
+            return f'{value}'
+        elif type(value) is str and value.startswith('~') and value.endswith('~'):
+            return f'{value.strip("~")}'
+        else:
+            return f'\"{str(value)}\"'

+ 3 - 0
requirements.txt

@@ -0,0 +1,3 @@
+mysql-connector-python>=8.0.30
+pymongo>=4.4.1
+pandas>=1.1.2