#!/usr/bin/env python3 """ ClickHouse Database Manager Centralized class for ClickHouse database connections and operations """ import os import sys from datetime import datetime from typing import Dict, Any, List, Optional, Tuple, Union from clickhouse_driver import Client import threading import time import pytz # Add parent directories to path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.logger import db_logger class ClickHouseManager: """ Centralized ClickHouse database manager with connection pooling and common operations """ def __init__(self, host: str = "localhost", port: int = 9000, database: str = "wallets_info", user: str = "default", password: str = ""): """Initialize ClickHouse manager with connection settings""" self.host = host self.port = port self.database = database self.user = user self.password = password # Thread-local storage for connections self._local = threading.local() # Connection settings self.connection_settings = { 'connect_timeout': 3, 'send_timeout': 10, 'receive_timeout': 10, 'max_block_size': 1000, 'max_insert_block_size': 1000, 'use_numpy': False } # Test initial connection self._test_connection() def _get_client(self) -> Client: """Get thread-local ClickHouse client""" if not hasattr(self._local, 'client') or self._local.client is None: try: self._local.client = Client( host=self.host, port=self.port, database=self.database, user=self.user, password=self.password, settings=self.connection_settings ) db_logger.debug(f"Created new ClickHouse client for thread {threading.current_thread().ident}") except Exception as e: db_logger.error(f"Failed to create ClickHouse client: {e}") raise return self._local.client def _test_connection(self): """Test ClickHouse connection""" try: client = self._get_client() result = client.execute("SELECT 1") db_logger.info(f"✅ ClickHouse connection successful: {result}") except Exception as e: db_logger.error(f"❌ ClickHouse connection failed: {e}") raise ConnectionError(f"Cannot connect to ClickHouse: {e}") def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Tuple]: """Execute a ClickHouse query with optional parameters""" try: client = self._get_client() if params: result = client.execute(query, params) else: result = client.execute(query) return result except Exception as e: db_logger.error(f"Query execution failed: {e}") db_logger.error(f"Query: {query}") if params: db_logger.error(f"Params: {params}") raise def execute_insert(self, table: str, columns: List[str], data: List[Tuple]) -> bool: """Execute bulk insert into ClickHouse table""" if not data: db_logger.warning("No data provided for insert") return False try: client = self._get_client() columns_str = ", ".join(columns) query = f"INSERT INTO {table} ({columns_str}) VALUES" client.execute(query, data) db_logger.info(f"✅ Bulk inserted {len(data)} rows into {table}") return True except Exception as e: db_logger.error(f"Bulk insert failed for table {table}: {e}") return False def execute_update(self, table: str, column: str , data: str, wallet : str) -> bool: """Execute bulk insert into ClickHouse table""" if not data: db_logger.warning("No data provided for update") return False try: client = self._get_client() query = f"ALTER TABLE {table} UPDATE {column} = toDateTime(%(value)s) WHERE wallet_address = %(wallet)s" params = { "value": data, "wallet": wallet } print(query,params) client.execute(query, params) return True except Exception as e: print(f"Bulk insert failed for table {table}: {e}") return False # ========================================== # TOKEN DATA OPERATIONS # ========================================== def get_token_as_processed(self,token): try: query = "SELECT processed FROM tokens_info.tokens WHERE token_address = %(token_address)s LIMIT 1" result = self.execute_query(query, {'token_address': token}) return result[0][0] if result else None except Exception as e: db_logger.error(f"Error getting pool for token {token}: {e}", exc_info=True) return None def get_token_pool(self,token): try: query = "SELECT pool FROM tokens_info.tokens WHERE token_address = %(token_address)s LIMIT 1" result = self.execute_query(query, {'token_address': token}) return result[0][0] if result else None except Exception as e: db_logger.error(f"Error getting pool for token {token}: {e}", exc_info=True) return None def set_token_as_quality(self,token,quality): try: query = "ALTER TABLE tokens_info.tokens UPDATE quality = %(quality)s WHERE token_address = %(token_address)s" self.execute_query(query, {'token_address': token, 'quality': quality}) db_logger.info(f"✅ Set token quality in ClickHouse: {token} to {quality}") except Exception as e: db_logger.error(f"Error setting token as processed {token}: {e}", exc_info=True) def set_token_as_processed(self,token): try: query = "ALTER TABLE tokens_info.tokens UPDATE processed = 1 WHERE token_address = %(token_address)s" self.execute_query(query, {'token_address': token}) db_logger.info(f"✅ Set token as processed in ClickHouse: {token}") except Exception as e: db_logger.error(f"Error setting token as processed {token}: {e}", exc_info=True) # ========================================== # WALLETS_DATA TABLE OPERATIONS # ========================================== def wallet_exists(self, wallet_address: str) -> bool: """Check if wallet exists in wallets_data table""" try: query = "SELECT 1 FROM wallets_data WHERE wallet_address = %(wallet_address)s LIMIT 1" result = self.execute_query(query, {'wallet_address': wallet_address}) return len(result) > 0 except Exception as e: db_logger.error(f"Error checking wallet existence for {wallet_address}: {e}") return False def set_wallet_as_processed(self, wallet: str, token: str): """Mark wallet as processed in ClickHouse database""" try: warsaw_tz = pytz.timezone("Europe/Warsaw") now = datetime.now(warsaw_tz).strftime('%Y-%m-%d %H:%M:%S') self.execute_update('wallets_data', 'processed_time', now , wallet) db_logger.info(f"💶Updated wallet processed [{wallet}] -> : TRUE ") except Exception as e: db_logger.error(f"💶Error updating wallet info in ClickHouse DB: {e}", exc_info=True) def get_wallet_processed_time(self, wallet_address: str) -> Optional[int]: """Get wallet processed time as Unix timestamp""" try: query = "SELECT toUnixTimestamp(processed_time) FROM wallets_data WHERE wallet_address = %(wallet_address)s LIMIT 1" result = self.execute_query(query, {'wallet_address': wallet_address}) if result: return result[0][0] return None except Exception as e: db_logger.error(f"Error getting processed time for {wallet_address}: {e}") return None def get_wallet_type(self, wallet_address: str) -> Optional[str]: """Get wallet type from wallets_data table""" try: query = "SELECT wallet_type FROM wallets_data WHERE wallet_address = %(wallet_address)s LIMIT 1" result = self.execute_query(query, {'wallet_address': wallet_address}) if result: return result[0][0] return None except Exception as e: db_logger.error(f"Error getting wallet type for {wallet_address}: {e}") return None def get_last_signature(self, wallet_address: str) -> Optional[str]: """Get wallet type from wallets_data table""" try: query = "SELECT last_signature FROM wallets_data WHERE wallet_address = %(wallet_address)s LIMIT 1" result = self.execute_query(query, {'wallet_address': wallet_address}) if result: return result[0][0] return None except Exception as e: db_logger.error(f"Error getting wallet type for {wallet_address}: {e}") return None def set_last_signature(self, wallet_address: str, signature: str) -> bool: """Set last signature for wallet in wallets_data table""" try: # Check if wallet exists exists = self.wallet_exists(wallet_address) if not exists: # Create new wallet entry with the signature db_logger.info(f"Creating new wallet entry for {wallet_address} with signature {signature}") wallet_data = { 'wallet_address': wallet_address, 'last_signature': signature, 'processed_time': datetime.now(), 'token_amount': 0, 'tokens_in_MC': 0, 'f_tokens_ph': 0, 'f_hold_time': 0, 'f_minbuy': 0, 'f_before_mig': 0, 'f_created_tokens': 0, 'wallet_type': 'UNKNOWN' } result = self.insert_wallet(wallet_data) if result: db_logger.info(f"✅ Set last signature for new wallet {wallet_address}: {signature}") return result # Update existing wallet's signature query = "ALTER TABLE wallets_data UPDATE last_signature=%(signature)s WHERE wallet_address = %(wallet_address)s" self.execute_query(query, {'wallet_address': wallet_address, "signature": signature}) db_logger.info(f"✅ Updated last signature for {wallet_address}: {signature}") return True except Exception as e: db_logger.error(f"❌ Error setting last signature for {wallet_address}: {e}") return False def set_token_extracted(self, wallet_address: str, token: str) -> bool: """Set token extracted for wallet in wallets_data table""" try: # Update existing wallet's signature query = "ALTER TABLE wallets_data UPDATE token_extracted=%(token_extracted)s WHERE wallet_address = %(wallet_address)s" self.execute_query(query, {'wallet_address': wallet_address, "token_extracted": token}) db_logger.info(f"✅ Updated token_extracted for {wallet_address}: {token}") return True except Exception as e: db_logger.error(f"❌ Error setting last signature for {wallet_address}: {e}") return False def insert_wallet(self, wallet_data: Dict[str, Any]) -> bool: """Insert new wallet into wallets_data table""" try: # Prepare data as tuple insert_data = [( wallet_data.get('wallet_address', ''), wallet_data.get('last_signature', ''), wallet_data.get('processed_time', datetime.now()), int(wallet_data.get('token_amount', 0)), int(wallet_data.get('tokens_in_MC', 0)), int(wallet_data.get('f_tokens_ph', 0)), int(wallet_data.get('f_hold_time', 0)), int(wallet_data.get('f_minbuy', 0)), int(wallet_data.get('f_before_mig', 0)), int(wallet_data.get('f_created_tokens', 0)), str(wallet_data.get('wallet_type', 'UNKNOWN')) )] columns = [ 'wallet_address', 'last_signature', 'processed_time', 'token_amount', 'tokens_in_MC', 'f_tokens_ph', 'f_hold_time', 'f_minbuy', 'f_before_mig', 'f_created_tokens', 'wallet_type' ] return self.execute_insert('wallets_data', columns, insert_data) except Exception as e: db_logger.error(f"Error inserting wallet {wallet_data.get('wallet_address', 'unknown')}: {e}") return False def update_wallet_type(self, wallet_address: str, wallet_type: str) -> bool: """Update wallet type in wallets_data table""" try: if not self.wallet_exists(wallet_address): db_logger.warning(f"Wallet {wallet_address} not found for type update") return False query = """ ALTER TABLE wallets_data UPDATE wallet_type = %(wallet_type)s WHERE wallet_address = %(wallet_address)s """ self.execute_query(query, { 'wallet_type': wallet_type, 'wallet_address': wallet_address }) db_logger.info(f"Updated wallet type for {wallet_address} to {wallet_type}") return True except Exception as e: db_logger.error(f"Error updating wallet type for {wallet_address}: {e}") return False def update_wallet_stats(self, wallet_address: str, stats_data: Dict[str, Any]) -> bool: """Update or create wallet statistics in wallets_data table""" try: # Check if wallet exists exists = self.wallet_exists(wallet_address) if not exists: # Create new wallet with stats data db_logger.info(f"Creating new wallet entry for {wallet_address}") wallet_data = { 'wallet_address': wallet_address, 'last_signature': '', 'processed_time': datetime.now(), 'token_amount': int(stats_data.get('token_amount', 0)), 'tokens_in_MC': int(stats_data.get('tokens_in_MC', 0)), 'f_tokens_ph': int(stats_data.get('f_tokens_ph', 0)), 'f_hold_time': int(stats_data.get('f_hold_time', 0)), 'f_minbuy': int(stats_data.get('f_minbuy', 0)), 'f_before_mig': int(stats_data.get('f_before_mig', 0)), 'f_created_tokens': int(stats_data.get('f_created_tokens', 0)), 'wallet_type': str(stats_data.get('wallet_type', 'UNKNOWN')) } return self.insert_wallet(wallet_data) # Update existing wallet update_query = """ ALTER TABLE wallets_data UPDATE token_amount = %(token_amount)s, tokens_in_MC = %(tokens_in_MC)s, f_tokens_ph = %(f_tokens_ph)s, f_hold_time = %(f_hold_time)s, f_minbuy = %(f_minbuy)s, f_before_mig = %(f_before_mig)s, f_created_tokens = %(f_created_tokens)s, wallet_type = %(wallet_type)s WHERE wallet_address = %(wallet_address)s """ update_params = { 'wallet_address': wallet_address, 'token_amount': int(stats_data.get('token_amount', 0)), 'tokens_in_MC': int(stats_data.get('tokens_in_MC', 0)), 'f_tokens_ph': int(stats_data.get('f_tokens_ph', 0)), 'f_hold_time': int(stats_data.get('f_hold_time', 0)), 'f_minbuy': int(stats_data.get('f_minbuy', 0)), 'f_before_mig': int(stats_data.get('f_before_mig', 0)), 'f_created_tokens': int(stats_data.get('f_created_tokens', 0)), 'wallet_type': str(stats_data.get('wallet_type', 'UNKNOWN')) } self.execute_query(update_query, update_params) db_logger.info(f"Updated wallet stats for {wallet_address}: {stats_data.get('wallet_type', 'UNKNOWN')}") return True except Exception as e: db_logger.error(f"Error updating wallet stats for {wallet_address}: {e}") return False def processed_more_than_week_ago(self, wallet_address: str) -> bool: """Check if wallet was processed more than a week ago""" try: processed_time = self.get_wallet_processed_time(wallet_address) if processed_time is None: return False current_time = time.time() week_seconds = 7 * 24 * 60 * 60 # 604800 seconds return (current_time - processed_time) > week_seconds except Exception as e: db_logger.error(f"Error checking week-ago status for {wallet_address}: {e}") return False # ========================================== # TRANSACTIONS TABLE OPERATIONS # ========================================== def get_wallet_transactions(self, wallet_address: str) -> List[Dict[str, Any]]: """Get all transactions for a wallet from transactions table""" try: query = """ SELECT wallet_address, token_extracted, signature, type, profit, blocktime, delta_sol, delta_token, fee, MC FROM transactions WHERE wallet_address = %(wallet_address)s ORDER BY blocktime DESC """ result = self.execute_query(query, {'wallet_address': wallet_address}) keys = ['wallet', 'mint_address', 'signature', 'type', 'profit', 'blocktime', 'delta_sol', 'delta_token', 'fee', 'MC'] return [dict(zip(keys, row)) for row in result] except Exception as e: db_logger.error(f"Error getting transactions for {wallet_address}: {e}") return [] #TODO check if transaction exists in db def insert_transactions(self, transactions: List[Dict[str, Any]]) -> bool: """Insert multiple transactions into transactions table""" if not transactions: db_logger.warning("No transactions provided for insert") return False try: data_to_insert = [] for tx in transactions: data_to_insert.append(( tx.get('wallet_address', ''), tx.get('token_extracted', ''), tx.get('signature', ''), tx.get('type', ''), tx.get('profit'), int(tx.get('blocktime', 0)), float(tx.get('delta_sol', 0)), float(tx.get('delta_token', 0)), float(tx.get('fee', 0)), float(tx.get('MC', 0)) )) columns = [ 'wallet_address', 'token_extracted', 'signature', 'type', 'profit', 'blocktime', 'delta_sol', 'delta_token', 'fee', 'MC' ] return self.execute_insert('transactions', columns, data_to_insert) except Exception as e: db_logger.error(f"Error inserting transactions: {e}") return False # ========================================== # UTILITY METHODS # ========================================== def get_existing_wallets(self, wallet_addresses: List[str]) -> set: """Get set of existing wallet addresses from provided list""" if not wallet_addresses: return set() try: # Use IN clause for bulk checking placeholders = "', '".join(wallet_addresses) query = f"SELECT wallet_address FROM wallets_data WHERE wallet_address IN ('{placeholders}')" result = self.execute_query(query) return {row[0] for row in result} except Exception as e: db_logger.error(f"Error checking existing wallets: {e}") return set() def get_table_count(self, table_name: str) -> int: """Get record count for a table""" try: query = f"SELECT COUNT(*) FROM {table_name}" result = self.execute_query(query) return result[0][0] if result else 0 except Exception as e: db_logger.error(f"Error getting count for table {table_name}: {e}") return 0 def close_connection(self): """Close thread-local connection""" if hasattr(self._local, 'client') and self._local.client: try: self._local.client.disconnect() self._local.client = None db_logger.debug("Closed ClickHouse connection") except Exception as e: db_logger.error(f"Error closing ClickHouse connection: {e}") # Global instance for easy importing clickhouse_manager = ClickHouseManager()