Initial upload of database_manager.py

c534334 verified about 1 year ago

15.8 kB

import sqlite3
import pandas as pd
from datetime import datetime
from pathlib import Path
import json
from typing import List, Dict, Any, Optional

class DatabaseManager:
    def __init__(self, db_path: str = "business_database.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize database with required tables"""
        with sqlite3.connect(self.db_path) as conn:
            # Companies table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS companies (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    address TEXT,
                    phone TEXT,
                    website TEXT,
                    email TEXT,
                    rating TEXT,
                    reviews TEXT,
                    category TEXT,
                    hours TEXT,
                    price TEXT,
                    description TEXT,
                    dataset_name TEXT NOT NULL,
                    business_type TEXT,
                    location TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(name, address, dataset_name)
                )
            """)
            
            # Business categories table (for multi-category support)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS company_categories (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    company_id INTEGER,
                    category TEXT,
                    FOREIGN KEY (company_id) REFERENCES companies (id),
                    UNIQUE(company_id, category)
                )
            """)
            
            # Contact collection status
            conn.execute("""
                CREATE TABLE IF NOT EXISTS contact_collection_status (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    company_id INTEGER,
                    website_scraped BOOLEAN DEFAULT FALSE,
                    email_found BOOLEAN DEFAULT FALSE,
                    phone_found BOOLEAN DEFAULT FALSE,
                    last_scraped TIMESTAMP,
                    FOREIGN KEY (company_id) REFERENCES companies (id)
                )
            """)
            
            # Datasets table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS datasets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    description TEXT,
                    business_type TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    total_companies INTEGER DEFAULT 0
                )
            """)
            
            conn.commit()
    
    def create_dataset(self, name: str, description: str = "", business_type: str = "") -> bool:
        """Create a new dataset"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT INTO datasets (name, description, business_type)
                    VALUES (?, ?, ?)
                """, (name, description, business_type))
                conn.commit()
                return True
        except sqlite3.IntegrityError:
            return False  # Dataset already exists
    
    def get_datasets(self) -> List[Dict]:
        """Get all datasets"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT d.*, COUNT(c.id) as company_count
                FROM datasets d
                LEFT JOIN companies c ON d.name = c.dataset_name
                GROUP BY d.id
                ORDER BY d.created_at DESC
            """)
            return [dict(row) for row in cursor.fetchall()]
    
    def add_companies(self, companies: List[Dict], dataset_name: str, business_type: str = "") -> int:
        """Add companies to database, return number of new companies added"""
        added_count = 0
        
        with sqlite3.connect(self.db_path) as conn:
            for company in companies:
                try:
                    # Insert company
                    cursor = conn.execute("""
                        INSERT INTO companies (
                            name, address, phone, website, email, rating, reviews,
                            category, hours, price, description, dataset_name, business_type, location
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """, (
                        company.get('name', ''),
                        company.get('address', ''),
                        company.get('phone', ''),
                        company.get('website', ''),
                        company.get('email', ''),
                        company.get('rating', ''),
                        company.get('reviews', ''),
                        company.get('category', ''),
                        company.get('hours', ''),
                        company.get('price', ''),
                        company.get('description', ''),
                        dataset_name,
                        business_type,
                        company.get('location', '')
                    ))
                    
                    company_id = cursor.lastrowid
                    added_count += 1
                    
                    # Add categories if multiple
                    categories = company.get('categories', [])
                    if not categories and company.get('category'):
                        categories = [company.get('category')]
                    
                    for category in categories:
                        if category:
                            conn.execute("""
                                INSERT OR IGNORE INTO company_categories (company_id, category)
                                VALUES (?, ?)
                            """, (company_id, category))
                    
                    # Initialize contact collection status
                    conn.execute("""
                        INSERT OR IGNORE INTO contact_collection_status (company_id)
                        VALUES (?)
                    """, (company_id,))
                    
                except sqlite3.IntegrityError:
                    # Company already exists, skip
                    continue
            
            # Update dataset company count
            conn.execute("""
                UPDATE datasets 
                SET total_companies = (
                    SELECT COUNT(*) FROM companies WHERE dataset_name = ?
                )
                WHERE name = ?
            """, (dataset_name, dataset_name))
            
            conn.commit()
        
        return added_count
    
    def get_companies(self, dataset_name: str = None, business_type: str = None) -> pd.DataFrame:
        """Get companies from database"""
        with sqlite3.connect(self.db_path) as conn:
            query = """
                SELECT c.*, 
                       GROUP_CONCAT(cc.category, '; ') as all_categories,
                       ccs.website_scraped, ccs.email_found, ccs.phone_found, ccs.last_scraped
                FROM companies c
                LEFT JOIN company_categories cc ON c.id = cc.company_id
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
            """
            params = []
            conditions = []
            
            if dataset_name:
                conditions.append("c.dataset_name = ?")
                params.append(dataset_name)
            
            if business_type:
                conditions.append("c.business_type = ?")
                params.append(business_type)
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            
            query += " GROUP BY c.id ORDER BY c.created_at DESC"
            
            return pd.read_sql_query(query, conn, params=params)
    
    def update_company(self, company_id: int, updates: Dict) -> bool:
        """Update company information"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                if not updates:
                    return True
                
                # Add updated_at to updates
                updates['updated_at'] = datetime.now().isoformat()
                
                # Build the SET clause and values
                set_clause = ", ".join([f"{key} = ?" for key in updates.keys()])
                values = list(updates.values()) + [company_id]
                
                query = f"""
                    UPDATE companies 
                    SET {set_clause}
                    WHERE id = ?
                """
                
                conn.execute(query, values)
                conn.commit()
                return True
        except Exception as e:
            print(f"Error updating company {company_id}: {e}")
            print(f"Updates: {updates}")
            print(f"Values: {values}")
            return False
    
    def delete_company(self, company_id: int) -> bool:
        """Delete a company from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # Delete from contact_collection_status first (foreign key constraint)
                conn.execute("DELETE FROM contact_collection_status WHERE company_id = ?", (company_id,))
                
                # Delete from company_categories
                conn.execute("DELETE FROM company_categories WHERE company_id = ?", (company_id,))
                
                # Delete the company
                conn.execute("DELETE FROM companies WHERE id = ?", (company_id,))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error deleting company {company_id}: {e}")
            return False
    
    def mark_contact_collected(self, company_id: int, email_found: bool = False, phone_found: bool = False):
        """Mark contact collection status"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                UPDATE contact_collection_status 
                SET website_scraped = TRUE, 
                    email_found = ?, 
                    phone_found = ?, 
                    last_scraped = CURRENT_TIMESTAMP
                WHERE company_id = ?
            """, (email_found, phone_found, company_id))
            conn.commit()
    
    def get_companies_for_contact_collection(self, dataset_name: str = None, search_type: str = "both") -> List[Dict]:
        """Get companies that need contact information collection - only those missing emails or phones"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            
            # Build conditions based on search type
            conditions = ["c.website IS NOT NULL AND c.website != ''"]
            
            if search_type == "emails_only":
                conditions.append("(c.email IS NULL OR c.email = '')")
            elif search_type == "phones_only":
                conditions.append("(c.phone IS NULL OR c.phone = '')")
            else:  # both or all
                conditions.append("""
                    (
                        (c.email IS NULL OR c.email = '') OR 
                        (c.phone IS NULL OR c.phone = '') OR
                        (ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
                    )
                """)
            
            query = f"""
                SELECT c.*, ccs.website_scraped, ccs.email_found, ccs.phone_found
                FROM companies c
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE {' AND '.join(conditions)}
            """
            params = []
            
            if dataset_name:
                query += " AND c.dataset_name = ?"
                params.append(dataset_name)
            
            query += " ORDER BY ccs.last_scraped ASC NULLS FIRST"
            
            cursor = conn.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]
    
    def delete_dataset(self, dataset_name: str) -> bool:
        """Delete a dataset and all its companies"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # Get company IDs for this dataset
                cursor = conn.execute("SELECT id FROM companies WHERE dataset_name = ?", (dataset_name,))
                company_ids = [row[0] for row in cursor.fetchall()]
                
                if company_ids:
                    # Delete related records
                    conn.execute("DELETE FROM company_categories WHERE company_id IN ({})".format(
                        ','.join('?' * len(company_ids))), company_ids)
                    conn.execute("DELETE FROM contact_collection_status WHERE company_id IN ({})".format(
                        ','.join('?' * len(company_ids))), company_ids)
                    conn.execute("DELETE FROM companies WHERE dataset_name = ?", (dataset_name,))
                
                # Delete dataset
                conn.execute("DELETE FROM datasets WHERE name = ?", (dataset_name,))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error deleting dataset: {e}")
            return False
    
    def get_dataset_stats(self, dataset_name: str) -> Dict:
        """Get statistics for a dataset"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            
            # Basic stats
            cursor = conn.execute("""
                SELECT 
                    COUNT(*) as total_companies,
                    COUNT(CASE WHEN email IS NOT NULL AND email != '' THEN 1 END) as with_email,
                    COUNT(CASE WHEN phone IS NOT NULL AND phone != '' THEN 1 END) as with_phone,
                    COUNT(CASE WHEN website IS NOT NULL AND website != '' THEN 1 END) as with_website
                FROM companies 
                WHERE dataset_name = ?
            """, (dataset_name,))
            basic_stats = dict(cursor.fetchone())
            
            # Contact collection stats
            cursor = conn.execute("""
                SELECT 
                    COUNT(*) as scraped_count,
                    COUNT(CASE WHEN email_found = 1 THEN 1 END) as email_found_count,
                    COUNT(CASE WHEN phone_found = 1 THEN 1 END) as phone_found_count
                FROM companies c
                JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE c.dataset_name = ? AND ccs.website_scraped = 1
            """, (dataset_name,))
            contact_stats = dict(cursor.fetchone())
            
            # Companies needing contact collection
            cursor = conn.execute("""
                SELECT COUNT(*) as need_contact_collection
                FROM companies c
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE c.dataset_name = ? 
                AND c.website IS NOT NULL AND c.website != ''
                AND (
                    (c.email IS NULL OR c.email = '') OR 
                    (c.phone IS NULL OR c.phone = '') OR
                    (ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
                )
            """, (dataset_name,))
            need_collection = dict(cursor.fetchone())
            
            return {**basic_stats, **contact_stats, **need_collection}

	import sqlite3
	import pandas as pd
	from datetime import datetime
	from pathlib import Path
	import json
	from typing import List, Dict, Any, Optional

	class DatabaseManager:
	def __init__(self, db_path: str = "business_database.db"):
	self.db_path = db_path
	self.init_database()

	def init_database(self):
	"""Initialize database with required tables"""
	with sqlite3.connect(self.db_path) as conn:
	# Companies table
	conn.execute("""
	CREATE TABLE IF NOT EXISTS companies (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	name TEXT NOT NULL,
	address TEXT,
	phone TEXT,
	website TEXT,
	email TEXT,
	rating TEXT,
	reviews TEXT,
	category TEXT,
	hours TEXT,
	price TEXT,
	description TEXT,
	dataset_name TEXT NOT NULL,
	business_type TEXT,
	location TEXT,
	created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
	updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
	UNIQUE(name, address, dataset_name)
	)
	""")

	# Business categories table (for multi-category support)
	conn.execute("""
	CREATE TABLE IF NOT EXISTS company_categories (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	company_id INTEGER,
	category TEXT,
	FOREIGN KEY (company_id) REFERENCES companies (id),
	UNIQUE(company_id, category)
	)
	""")

	# Contact collection status
	conn.execute("""
	CREATE TABLE IF NOT EXISTS contact_collection_status (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	company_id INTEGER,
	website_scraped BOOLEAN DEFAULT FALSE,
	email_found BOOLEAN DEFAULT FALSE,
	phone_found BOOLEAN DEFAULT FALSE,
	last_scraped TIMESTAMP,
	FOREIGN KEY (company_id) REFERENCES companies (id)
	)
	""")

	# Datasets table
	conn.execute("""
	CREATE TABLE IF NOT EXISTS datasets (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	name TEXT UNIQUE NOT NULL,
	description TEXT,
	business_type TEXT,
	created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
	total_companies INTEGER DEFAULT 0
	)
	""")

	conn.commit()

	def create_dataset(self, name: str, description: str = "", business_type: str = "") -> bool:
	"""Create a new dataset"""
	try:
	with sqlite3.connect(self.db_path) as conn:
	conn.execute("""
	INSERT INTO datasets (name, description, business_type)
	VALUES (?, ?, ?)
	""", (name, description, business_type))
	conn.commit()
	return True
	except sqlite3.IntegrityError:
	return False # Dataset already exists

	def get_datasets(self) -> List[Dict]:
	"""Get all datasets"""
	with sqlite3.connect(self.db_path) as conn:
	conn.row_factory = sqlite3.Row
	cursor = conn.execute("""
	SELECT d.*, COUNT(c.id) as company_count
	FROM datasets d
	LEFT JOIN companies c ON d.name = c.dataset_name
	GROUP BY d.id
	ORDER BY d.created_at DESC
	""")
	return [dict(row) for row in cursor.fetchall()]

	def add_companies(self, companies: List[Dict], dataset_name: str, business_type: str = "") -> int:
	"""Add companies to database, return number of new companies added"""
	added_count = 0

	with sqlite3.connect(self.db_path) as conn:
	for company in companies:
	try:
	# Insert company
	cursor = conn.execute("""
	INSERT INTO companies (
	name, address, phone, website, email, rating, reviews,
	category, hours, price, description, dataset_name, business_type, location
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	""", (
	company.get('name', ''),
	company.get('address', ''),
	company.get('phone', ''),
	company.get('website', ''),
	company.get('email', ''),
	company.get('rating', ''),
	company.get('reviews', ''),
	company.get('category', ''),
	company.get('hours', ''),
	company.get('price', ''),
	company.get('description', ''),
	dataset_name,
	business_type,
	company.get('location', '')
	))

	company_id = cursor.lastrowid
	added_count += 1

	# Add categories if multiple
	categories = company.get('categories', [])
	if not categories and company.get('category'):
	categories = [company.get('category')]

	for category in categories:
	if category:
	conn.execute("""
	INSERT OR IGNORE INTO company_categories (company_id, category)
	VALUES (?, ?)
	""", (company_id, category))

	# Initialize contact collection status
	conn.execute("""
	INSERT OR IGNORE INTO contact_collection_status (company_id)
	VALUES (?)
	""", (company_id,))

	except sqlite3.IntegrityError:
	# Company already exists, skip
	continue

	# Update dataset company count
	conn.execute("""
	UPDATE datasets
	SET total_companies = (
	SELECT COUNT(*) FROM companies WHERE dataset_name = ?
	)
	WHERE name = ?
	""", (dataset_name, dataset_name))

	conn.commit()

	return added_count

	def get_companies(self, dataset_name: str = None, business_type: str = None) -> pd.DataFrame:
	"""Get companies from database"""
	with sqlite3.connect(self.db_path) as conn:
	query = """
	SELECT c.*,
	GROUP_CONCAT(cc.category, '; ') as all_categories,
	ccs.website_scraped, ccs.email_found, ccs.phone_found, ccs.last_scraped
	FROM companies c
	LEFT JOIN company_categories cc ON c.id = cc.company_id
	LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
	"""
	params = []
	conditions = []

	if dataset_name:
	conditions.append("c.dataset_name = ?")
	params.append(dataset_name)

	if business_type:
	conditions.append("c.business_type = ?")
	params.append(business_type)

	if conditions:
	query += " WHERE " + " AND ".join(conditions)

	query += " GROUP BY c.id ORDER BY c.created_at DESC"

	return pd.read_sql_query(query, conn, params=params)

	def update_company(self, company_id: int, updates: Dict) -> bool:
	"""Update company information"""
	try:
	with sqlite3.connect(self.db_path) as conn:
	if not updates:
	return True

	# Add updated_at to updates
	updates['updated_at'] = datetime.now().isoformat()

	# Build the SET clause and values
	set_clause = ", ".join([f"{key} = ?" for key in updates.keys()])
	values = list(updates.values()) + [company_id]

	query = f"""
	UPDATE companies
	SET {set_clause}
	WHERE id = ?
	"""

	conn.execute(query, values)
	conn.commit()
	return True
	except Exception as e:
	print(f"Error updating company {company_id}: {e}")
	print(f"Updates: {updates}")
	print(f"Values: {values}")
	return False

	def delete_company(self, company_id: int) -> bool:
	"""Delete a company from database"""
	try:
	with sqlite3.connect(self.db_path) as conn:
	# Delete from contact_collection_status first (foreign key constraint)
	conn.execute("DELETE FROM contact_collection_status WHERE company_id = ?", (company_id,))

	# Delete from company_categories
	conn.execute("DELETE FROM company_categories WHERE company_id = ?", (company_id,))

	# Delete the company
	conn.execute("DELETE FROM companies WHERE id = ?", (company_id,))
	conn.commit()
	return True
	except Exception as e:
	print(f"Error deleting company {company_id}: {e}")
	return False

	def mark_contact_collected(self, company_id: int, email_found: bool = False, phone_found: bool = False):
	"""Mark contact collection status"""
	with sqlite3.connect(self.db_path) as conn:
	conn.execute("""
	UPDATE contact_collection_status
	SET website_scraped = TRUE,
	email_found = ?,
	phone_found = ?,
	last_scraped = CURRENT_TIMESTAMP
	WHERE company_id = ?
	""", (email_found, phone_found, company_id))
	conn.commit()

	def get_companies_for_contact_collection(self, dataset_name: str = None, search_type: str = "both") -> List[Dict]:
	"""Get companies that need contact information collection - only those missing emails or phones"""
	with sqlite3.connect(self.db_path) as conn:
	conn.row_factory = sqlite3.Row

	# Build conditions based on search type
	conditions = ["c.website IS NOT NULL AND c.website != ''"]

	if search_type == "emails_only":
	conditions.append("(c.email IS NULL OR c.email = '')")
	elif search_type == "phones_only":
	conditions.append("(c.phone IS NULL OR c.phone = '')")
	else: # both or all
	conditions.append("""
	(
	(c.email IS NULL OR c.email = '') OR
	(c.phone IS NULL OR c.phone = '') OR
	(ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
	)
	""")

	query = f"""
	SELECT c.*, ccs.website_scraped, ccs.email_found, ccs.phone_found
	FROM companies c
	LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
	WHERE {' AND '.join(conditions)}
	"""
	params = []

	if dataset_name:
	query += " AND c.dataset_name = ?"
	params.append(dataset_name)

	query += " ORDER BY ccs.last_scraped ASC NULLS FIRST"

	cursor = conn.execute(query, params)
	return [dict(row) for row in cursor.fetchall()]

	def delete_dataset(self, dataset_name: str) -> bool:
	"""Delete a dataset and all its companies"""
	try:
	with sqlite3.connect(self.db_path) as conn:
	# Get company IDs for this dataset
	cursor = conn.execute("SELECT id FROM companies WHERE dataset_name = ?", (dataset_name,))
	company_ids = [row[0] for row in cursor.fetchall()]

	if company_ids:
	# Delete related records
	conn.execute("DELETE FROM company_categories WHERE company_id IN ({})".format(
	','.join('?' * len(company_ids))), company_ids)
	conn.execute("DELETE FROM contact_collection_status WHERE company_id IN ({})".format(
	','.join('?' * len(company_ids))), company_ids)
	conn.execute("DELETE FROM companies WHERE dataset_name = ?", (dataset_name,))

	# Delete dataset
	conn.execute("DELETE FROM datasets WHERE name = ?", (dataset_name,))
	conn.commit()
	return True
	except Exception as e:
	print(f"Error deleting dataset: {e}")
	return False

	def get_dataset_stats(self, dataset_name: str) -> Dict:
	"""Get statistics for a dataset"""
	with sqlite3.connect(self.db_path) as conn:
	conn.row_factory = sqlite3.Row

	# Basic stats
	cursor = conn.execute("""
	SELECT
	COUNT(*) as total_companies,
	COUNT(CASE WHEN email IS NOT NULL AND email != '' THEN 1 END) as with_email,
	COUNT(CASE WHEN phone IS NOT NULL AND phone != '' THEN 1 END) as with_phone,
	COUNT(CASE WHEN website IS NOT NULL AND website != '' THEN 1 END) as with_website
	FROM companies
	WHERE dataset_name = ?
	""", (dataset_name,))
	basic_stats = dict(cursor.fetchone())

	# Contact collection stats
	cursor = conn.execute("""
	SELECT
	COUNT(*) as scraped_count,
	COUNT(CASE WHEN email_found = 1 THEN 1 END) as email_found_count,
	COUNT(CASE WHEN phone_found = 1 THEN 1 END) as phone_found_count
	FROM companies c
	JOIN contact_collection_status ccs ON c.id = ccs.company_id
	WHERE c.dataset_name = ? AND ccs.website_scraped = 1
	""", (dataset_name,))
	contact_stats = dict(cursor.fetchone())

	# Companies needing contact collection
	cursor = conn.execute("""
	SELECT COUNT(*) as need_contact_collection
	FROM companies c
	LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
	WHERE c.dataset_name = ?
	AND c.website IS NOT NULL AND c.website != ''
	AND (
	(c.email IS NULL OR c.email = '') OR
	(c.phone IS NULL OR c.phone = '') OR
	(ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
	)
	""", (dataset_name,))
	need_collection = dict(cursor.fetchone())

	return {basic_stats, contact_stats, **need_collection}