22 lines
759 B
Python
22 lines
759 B
Python
import duckdb
|
|
from cognitive_architecture.root_dir import get_absolute_path
|
|
|
|
class DuckDBAdapter():
|
|
def __init__(self):
|
|
db_path = get_absolute_path("./data/cognee")
|
|
db_location = db_path + "/cognee.duckdb"
|
|
|
|
self.db_client = duckdb.connect(db_location)
|
|
|
|
def get_datasets(self):
|
|
tables = self.db_client.sql("SELECT DISTINCT schema_name FROM duckdb_tables();").to_df().to_dict()
|
|
|
|
return list(
|
|
filter(
|
|
lambda table_name: table_name.endswith('staging') is False,
|
|
tables["schema_name"].values()
|
|
)
|
|
)
|
|
|
|
def get_files_metadata(self, dataset_name: str):
|
|
return self.db_client.sql(f"SELECT * FROM {dataset_name}.file_metadata;").to_df().to_dict()
|