optimized data writer & check

tutake/api/base_dao.py

@@ -381,10 +381,13 @@ def append(self, records):
                     self.fields = records.fields
                 self.items = self.items + records.items
 
+    def __str__(self):
+        return self.data_frame().__str__()
+
 
 class BatchWriter:
 
-    def __init__(self, engine, table: str, schema, database_dir):
+    def __init__(self, engine, table: str, schema, database_dir=None):
         self.writer = None
         self.engine = engine
         self.table = table
@@ -402,11 +405,12 @@ def _init_writer(self):
         return self.writer
 
     def start(self):
-        self._init_writer()
-        conn = self.engine.connect()
-        result = conn.execute(f"""Select max(id) from {self.table}""").fetchone()
-        if result is not None and result[0] is not None:
-            self.max_id = int(result[0])
+        if self.writer is None:
+            self._init_writer()
+            conn = self.engine.connect()
+            result = conn.execute(f"""Select max(id) from {self.table}""").fetchone()
+            if result is not None and result[0] is not None:
+                self.max_id = int(result[0])
 
     def rollback(self):
         if self.writer:
@@ -438,12 +442,19 @@ def add_records(self, records):
             data = records
         else:
             return
+
+        self.start()
         data['id'] = range(self.max_id, self.max_id + len(data))
         data = data.reindex(columns=['id'] + list(data.columns[:-1]))
         table = pa.Table.from_pandas(df=data, schema=self.schema)
         self.writer.write_table(table)
         self.max_id = self.max_id + len(data)
 
+    def flush(self):
+        conn = self.engine.connect()
+        conn.execute("FORCE CHECKPOINT;")
+        conn.close()
+
     def close(self):
         self.max_id = 0
         if self.writer:

tutake/api/checker.py

@@ -4,7 +4,7 @@
 from tutake.api.base_dao import checker_logger, BatchWriter
 
 
-def _auto_data_repair(self, trade_date, ts_codes):
+def _auto_data_repair(self, trade_date, ts_codes, tushare):
     writer = BatchWriter(self.engine, self.table_name)
     if len(ts_codes) < 100:
         for ts_code in ts_codes:
@@ -56,11 +56,10 @@ def check_by_date(self, method, default_start, force_start=None, date_apply=lamb
             tushare_pd = pd.DataFrame(tushare.items, columns=tushare.fields)
             diff = list(set(tushare_pd[diff_column].unique().tolist()) - set(db[diff_column].unique().tolist()))
             if diff_repair is not None:
-                diff_repair(self, trade_date, diff)
-                continue
+                diff_repair(self, trade_date, diff, tushare)
             else:
                 checker_logger.warning(
-                    f"Not equals data. The date is {trade_date}. tushare size is {tushare.size()}, db size is {db.shape[0]}, diff is {diff}")
+                    f"Not equals data {self.name}. The date is {trade_date}. tushare size is {tushare.size()}, db size is {db.shape[0]}, diff is {diff}")
                 if force_start is None:
                     self.checker.error_point(trade_date=trade_date)
             return

tutake/api/process.py

@@ -283,8 +283,12 @@ def _process_by_func(self, prepare_write, query_parameters, fetch_and_append, wr
             logging.exception(err)
         finally:
             writer.close()
+            last_task = kwargs.get("last_task")
+            if last_task is None or last_task is True:
+                writer.flush()
         return status
 
+
     def _inner_process(self, process_params, fetch_and_append, status: ProcessStatus, writer: BatchWriter = None,
                        retry_cnt=0, entrypoint=None):
         if retry_cnt > self.max_repeat:

tutake/api/process_client.py

@@ -111,20 +111,23 @@ def _get_all_task(self) -> [Task]:
         return apis
 
     def _do_process(self, tasks, entrypoint="scheduler"):
-        def __process(_job_id, _task) -> ProcessStatus:
+        def __process(_job_id, _task, is_last_task=None) -> ProcessStatus:
             if _task is not None:
-                return _task.process(entrypoint=entrypoint)
+                return _task.process(entrypoint=entrypoint, last_task=is_last_task)
             else:
                 return None
 
         start = time.time()
         status_list = []
         if isinstance(tasks, Task):
-            status_list.append(__process(f"tutake_{tasks.name}", tasks))
+            status_list.append(__process(f"tutake_{tasks.name}", tasks, None))
         elif isinstance(tasks, Sequence):
             for task in tasks:
                 try:
-                    status_list.append(__process(f"tutake_{task.name}", task))
+                    if task == tasks[-1]:
+                        status_list.append(__process(f"tutake_{task.name}", task, True))
+                    else:
+                        status_list.append(__process(f"tutake_{task.name}", task, False))
                 except Exception as err:
                     # self.logger.error(f"Exception with {api} process,err is {err}")
                     continue

tutake/api/ts/adj_factor.py

@@ -40,7 +40,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_adj_factor"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -50,7 +49,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareAdjFactor.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareAdjFactor)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareAdjFactor),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'trade_date', 'start_date', 'end_date', 'limit', 'offset']
         self.tushare_fields = ["ts_code", "trade_date", "adj_factor"]
@@ -101,8 +101,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/anns.py

@@ -45,7 +45,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_anns"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -55,7 +54,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareAnns.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareAnns)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareAnns),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'ann_date', 'start_date', 'end_date', 'limit', 'offset']
         self.tushare_fields = ["ts_code", "ann_date", "ann_type", "title", "content", "pub_time", "src_url", "filepath"]
@@ -131,8 +131,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/bak_basic.py

@@ -61,7 +61,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_bak_basic"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -71,7 +70,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareBakBasic.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareBakBasic)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareBakBasic),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['trade_date', 'ts_code', 'limit', 'offset']
         self.tushare_fields = [
@@ -233,8 +233,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/bak_daily.py

@@ -68,7 +68,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_bak_daily"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -78,7 +77,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareBakDaily.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareBakDaily)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareBakDaily),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'trade_date', 'start_date', 'end_date', 'offset', 'limit']
         self.tushare_fields = [
@@ -279,8 +279,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/balancesheet_vip.py

@@ -195,7 +195,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_balancesheet_vip"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -205,7 +204,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareBalancesheetVip.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareBalancesheetVip)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareBalancesheetVip),
+                                  config.get_tutake_data_dir())
 
         query_fields = [
             'ts_code', 'ann_date', 'f_ann_date', 'start_date', 'end_date', 'period', 'report_type', 'comp_type',
@@ -1099,8 +1099,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/cashflow_vip.py

@@ -134,7 +134,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_cashflow_vip"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -144,7 +143,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareCashflowVip.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareCashflowVip)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareCashflowVip),
+                                  config.get_tutake_data_dir())
 
         query_fields = [
             'ts_code', 'ann_date', 'f_ann_date', 'start_date', 'end_date', 'period', 'report_type', 'comp_type',
@@ -718,8 +718,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/ccass_hold.py

@@ -43,7 +43,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_ccass_hold"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -53,7 +52,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareCcassHold.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareCcassHold)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareCcassHold),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'trade_date', 'start_date', 'end_date', 'type', 'hk_hold', 'limit', 'offset']
         self.tushare_fields = ["trade_date", "ts_code", "name", "shareholding", "hold_nums", "hold_ratio"]
@@ -121,8 +121,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/ccass_hold_detail.py

@@ -44,7 +44,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_ccass_hold_detail"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -54,7 +53,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareCcassHoldDetail.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareCcassHoldDetail)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareCcassHoldDetail),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'trade_date', 'start_date', 'end_date', 'hk_code', 'limit', 'offset']
         self.tushare_fields = [
@@ -132,8 +132,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/ci_daily.py

@@ -48,7 +48,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_ci_daily"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -58,7 +57,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareCiDaily.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareCiDaily)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareCiDaily),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['ts_code', 'trade_date', 'start_date', 'end_date', 'limit', 'offset']
         self.tushare_fields = [
@@ -155,8 +155,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """

tutake/api/ts/cn_cpi.py

@@ -50,7 +50,6 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config):
         self.table_name = "tushare_cn_cpi"
         self.database = 'tutake.duckdb'
-        self.database_dir = config.get_tutake_data_dir()
         self.database_url = config.get_data_driver_url(self.database)
         self.engine = create_shared_engine(self.database_url,
                                            connect_args={
@@ -60,7 +59,8 @@ def __init__(self, config):
         session_factory = sessionmaker()
         session_factory.configure(bind=self.engine)
         TushareCnCpi.__table__.create(bind=self.engine, checkfirst=True)
-        self.schema = BaseDao.parquet_schema(TushareCnCpi)
+        self.writer = BatchWriter(self.engine, self.table_name, BaseDao.parquet_schema(TushareCnCpi),
+                                  config.get_tutake_data_dir())
 
         query_fields = ['m', 'start_m', 'end_m', 'limit', 'offset']
         self.tushare_fields = [
@@ -166,8 +166,7 @@ def process(self, **kwargs):
         同步历史数据
         :return:
         """
-        return super()._process(self.fetch_and_append,
-                                BatchWriter(self.engine, self.table_name, self.schema, self.database_dir), **kwargs)
+        return super()._process(self.fetch_and_append, self.writer, **kwargs)
 
     def fetch_and_append(self, **kwargs):
         """