duckdb_demo.py

import duckdb
import os
import pandas as pd

home_folder = os.path.dirname(__file__)

def execute_query(conn, query):
    if conn is None:
        print("Database connection not established.")
        return None
    try:
        result = conn.execute(query)
        return result
    except Exception as e:
        print(f"Error executing query: {e}")
        return None
    

try:
    conn = duckdb.connect(database=f'{home_folder}/my_database.duckdb', read_only=False)
    print("Connection to DuckDB established successfully.")
except Exception as e:
    print(f"Error connecting to database: {e}")
    conn = None

# Example usage
create_table_query = "CREATE TABLE IF NOT EXISTS users (id INTEGER, name VARCHAR);"
insert_table = "INSERT INTO users VALUES (1, 'Alice'), (2, 'Bob');"

select_query = "SELECT * FROM users;"

# Execute create table query
result = execute_query(conn, create_table_query)
print(result)

result = execute_query(conn, insert_table)
print(result)

result_df= execute_query(conn, select_query).fetch_df()
print(result_df.head())

# df to table
# 1. Create a sample DataFrame
new_df = pd.DataFrame({
    'id': [1, 2, 3, 4],
    'name': ['Charlie', 'David', 'Eve', 'Frank']
})
# 2. Register the DataFrame as a DuckDB table
conn.register('new_users', new_df)

# 3. Query the new table
new_query = "SELECT * FROM new_users;"
print("Data from new_users table:")
print(execute_query(conn, new_query).fetchall())

# close the connection
conn.close()