Refine some db codes (#2751)

* refine some db code

* polish

* add docstring

Author: sneaxiy

python/runtime/db.py

@@ -17,7 +17,6 @@
 import numpy as np
 import runtime.db_writer as db_writer
 import six
-from odps import ODPS, tunnel
 
 
 def parseMySQLDSN(dsn):

python/runtime/db_writer/hive.py

@@ -58,7 +58,6 @@ def _column_list(self):
         return result
 
     def _indexing_table_schema(self, table_schema):
-        cursor = self.conn.cursor()
         column_list = self._column_list()
 
         schema_idx = []
@@ -77,7 +76,7 @@ def _indexing_table_schema(self, table_schema):
 
     def _ordered_row_data(self, row):
         # Use NULL as the default value for hive columns
-        row_data = ["NULL" for i in range(len(self.table_schema))]
+        row_data = ["NULL" for _ in range(len(self.table_schema))]
         for idx, element in enumerate(row):
             row_data[self.schema_idx[idx]] = str(element)
         return CSV_DELIMITER.join(row_data)

python/runtime/db_writer/maxcompute.py

@@ -11,18 +11,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from odps import ODPS, tunnel
 from runtime.db_writer.base import BufferedDBWriter
 
 
 class MaxComputeDBWriter(BufferedDBWriter):
+    """
+    MaxComputeDBWriter is used to write the Python row data into
+    the MaxCompute table.
+
+    Args:
+        conn: the database connection object.
+        table_name (str): the MaxCompute table name.
+        table_schema (list[str]): the column names of the MaxCompute table.
+        buff_size (int): the buffer size to be flushed.
+    """
     def __init__(self, conn, table_name, table_schema, buff_size):
-        return super(MaxComputeDBWriter,
-                     self).__init__(conn, table_name, table_schema, buff_size)
+        super(MaxComputeDBWriter, self).__init__(conn, table_name,
+                                                 table_schema, buff_size)
+
+        # NOTE: import odps here instead of in the front of this file,
+        # so that we do not need the odps package installed in the Docker
+        # image if we do not use MaxComputeDBWriter.
+        from odps import tunnel
+        self.compress = tunnel.CompressOption.CompressAlgorithm.ODPS_ZLIB
 
     def flush(self):
-        compress = tunnel.CompressOption.CompressAlgorithm.ODPS_ZLIB
+        """
+        Flush the row data into the MaxCompute table.
+
+        Returns:
+            None
+        """
         self.conn.write_table(self.table_name,
                               self.rows,
-                              compress_option=compress)
+                              compress_option=self.compress)
         self.rows = []

python/runtime/db_writer/mysql.py

@@ -15,16 +15,32 @@
 
 
 class MySQLDBWriter(BufferedDBWriter):
-    def __init__(self, conn, table_name, table_schema, buff_size):
-        return super().__init__(conn, table_name, table_schema, buff_size)
+    """
+    MySQLDBWriter is used to write the Python row data into
+    the MySQL table.
 
-    def flush(self):
-        statement = '''insert into {} ({}) values({})'''.format(
+    Args:
+        conn: the database connection object.
+        table_name (str): the MySQL table name.
+        table_schema (list[str]): the column names of the MySQL table.
+        buff_size (int): the buffer size to be flushed.
+    """
+    def __init__(self, conn, table_name, table_schema, buff_size):
+        super().__init__(conn, table_name, table_schema, buff_size)
+        self.statement = '''insert into {} ({}) values({})'''.format(
             self.table_name, ", ".join(self.table_schema),
             ", ".join(["%s"] * len(self.table_schema)))
+
+    def flush(self):
+        """
+        Flush the row data into the MySQL table.
+
+        Returns:
+            None
+        """
         cursor = self.conn.cursor()
         try:
-            cursor.executemany(statement, self.rows)
+            cursor.executemany(self.statement, self.rows)
             self.conn.commit()
         finally:
             cursor.close()

python/runtime/diagnostics.py

@@ -10,7 +10,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import copy
 import inspect
 import os
 import re

python/runtime/maxcompute.py

@@ -19,7 +19,19 @@
 # Here we use the sdk to operate the database
 class MaxCompute:
     @staticmethod
-    def connect(database, user, password, host, auth=""):
+    def connect(database, user, password, host):
+        """
+        Create a MaxCompute database connection object.
+
+        Args:
+            database: the MaxCompute project name.
+            user: the MaxCompute AK.
+            password: the MaxCompute SK.
+            host: the MaxCompute endpoint address.
+
+        Returns:
+            A MaxCompute database connection object.
+        """
         return ODPS(user, password, project=database, endpoint=host)
 
     @staticmethod

python/runtime/oss.py

@@ -11,10 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import io
 import os
 import pickle
-import tarfile
 
 import oss2
 import tensorflow as tf