diff --git a/core/etl/extractor/mysql.py b/core/etl/extractor/mysql.py new file mode 100644 index 0000000000000000000000000000000000000000..e94add8aa24b0b2d56ca1ba241c86c7bf155d819 --- /dev/null +++ b/core/etl/extractor/mysql.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +# ---------------------------------------------------------------------- +# Mysql Data Extractor +# ---------------------------------------------------------------------- +# Copyright (C) 2007-2019 The NOC Project +# See LICENSE for details +# ---------------------------------------------------------------------- + +# Python modules +from __future__ import absolute_import +import os + +# NOC modules +from noc.core.etl.extractor.sql import SQLExtractor + + +class MySQLExtractor(SQLExtractor): + """ + Mysql SQL extractor. + Requres pymysql + + Configuration variables + *MYSQL_HOST* - Mysql host + *MYSQL_PORT* - Mysql port + *MYSQL_USER* - Mysql database user + *MYSQL_PASSWORD* - Mysql database password + *MYSQL_DB* - Mysql database + *MYSQL_CHARSET* - Mysql charset + """ + + def __init__(self, *args, **kwargs): + super(MySQLExtractor, self).__init__(*args, **kwargs) + self.connect = None + + def get_cursor(self): + import pymysql + + if not self.connect: + # Alter environment + old_env = os.environ.copy() # Save environment + os.environ.update(self.system.config) + # Connect to database + self.logger.info("Connecting to database") + self.connect = pymysql.connect( + host=str(self.config.get("MYSQL_HOST")), + port=int(self.config.get("MYSQL_PORT")), + user=str(self.config.get("MYSQL_USER")), + password=str(self.config.get("MYSQL_PASSWORD")), + db=str(self.config.get("MYSQL_DB")), + charset=str( + self.config.get("MYSQL_CHARSET") + if self.config.get("MYSQL_CHARSET") + else "utf8mb4" + ), + ) + + os.environ = old_env # Restore environment + cursor = self.connect.cursor() + return cursor + + def iter_data(self): + cursor = self.get_cursor() + # Fetch data + self.logger.info("Fetching data") + for query, params in self.get_sql(): + cursor.execute(query, params) + for row in cursor: + yield row