import gzip
import os
import pickle
import tempfile
import zipfile
import shutil
from copy import deepcopy
from functools import partial
from hashlib import md5
from itertools import islice
import psycopg
from aisdb.aisdb import decoder
from aisdb import sqlpath
from aisdb.database.dbconn import PostgresDBConn
[docs]
class FileChecksums:
"""
Initializes a FileChecksums object with a specified database connection.
:param dbconn: A required parameter of type PostgresDBConn that represents the database connection.
:param dbconn: A required parameter of type PostgresDBConn that represents the database connection.
:return: None
"""
def __init__(self, *, dbconn):
"""
:param dbconn: A required parameter of type PostgresDBConn that represents the database connection.
:param dbconn: A required parameter of type PostgresDBConn that represents the database connection.
:return: None
"""
assert isinstance(dbconn, (PostgresDBConn))
assert isinstance(dbconn, (PostgresDBConn))
self.dbconn = dbconn
self.checksums_table()
self.tmp_dir = tempfile.mkdtemp()
if not os.path.isdir(self.tmp_dir):
os.mkdir(self.tmp_dir)
[docs]
def checksums_table(self):
"""
Creates a checksums table in the database if it doesn't exist.
This method creates a table named 'hashmap' in the database, if it doesn't already exist.
The table contains two columns: 'hash' of type TEXT and 'bytes' of type BYTEA for PostgresDBConn.
The table contains two columns: 'hash' of type TEXT and 'bytes' of type BYTEA for PostgresDBConn.
:param self: an instance of the current object.
:return: None
"""
cur = self.dbconn.cursor()
if isinstance(self.dbconn, PostgresDBConn):
cur.execute("""
CREATE TABLE IF NOT EXISTS
hashmap(
hash TEXT PRIMARY KEY,
bytes BYTEA
);""")
cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_map on hashmap(hash)")
self.dbconn.commit()
[docs]
def insert_checksum(self, checksum):
"""
Inserts a checksum and corresponding pickled None value into the hashmap table.
:param checksum: The checksum to be inserted into the hashmap table.
:returns: None
"""
if isinstance(self.dbconn, PostgresDBConn):
self.dbconn.execute(
"INSERT INTO hashmap VALUES ($1,$2) ON CONFLICT DO NOTHING",
[checksum, pickle.dumps(None)],
)
self.dbconn.commit()
[docs]
def checksum_exists(self, checksum):
"""
Check if a given checksum exists in the database.
:param checksum: The checksum value to check.
:return: True if the checksum exists in the database, False otherwise.
"""
cur = self.dbconn.cursor()
if isinstance(self.dbconn, PostgresDBConn):
cur.execute("SELECT * FROM hashmap WHERE hash = %s", [checksum])
res = cur.fetchone()
if res is None or res is False:
return False
return True
[docs]
def get_md5(self, path, f):
"""
Calculates the MD5 hash digest of a file.
:param path: The path of the file.
:param f: The file object to calculate the digest for.
:return: The MD5 hash digest of the file content.
"""
if path[-4:].lower() == ".csv":
_ = f.read(1600) # skip the header (~1.6kb)
digest = md5(f.read(1000)).hexdigest()
return digest
def _fast_unzip(zipf, dirname):
"""
This function unzips a compressed file archive quickly; it handles both .zip and .gz file formats.
:param zipf: The path to the compressed file archive.
:param dirname: The directory to extract the contents of the archive.
"""
if zipf.lower()[-4:] == ".zip":
exists = set(sorted(os.listdir(dirname)))
with zipfile.ZipFile(zipf, "r") as zip_ref:
contents = set(zip_ref.namelist())
members = list(contents - exists)
try:
zip_ref.extractall(path=dirname, members=members)
except zipfile.BadZipFile as e:
print("Bad file found!")
elif zipf.lower()[-3:] == ".gz":
unzip_file = os.path.join(dirname, zipf.rsplit(os.path.sep, 1)[-1][:-3])
with gzip.open(zipf, "rb") as f1, open(unzip_file, "wb") as f2:
f2.write(f1.read())
else:
raise ValueError("unknown zip file type")
[docs]
def fast_unzip(zip_filenames, dirname):
"""
Unzips multiple zip files to a specified directory.
:param zip_filenames: List of zip file names to be extracted.
:param dirname: Directory path where the files should be extracted to.
"""
print(f"unzipping files to {dirname} ... ")
fcn = partial(_fast_unzip, dirname=dirname)
for file in zip_filenames:
fcn(file)
[docs]
def process_raw_files(
dbconn,
dbindex,
raw_files,
source,
timescaledb,
raw_insertion,
vacuum,
verbose,
workers,
type_preference,
not_zipped,
unzipped,
not_zipped_checksums,
unzipped_checksums,
skip_checksum,
):
if not raw_files:
if verbose:
print("All files returned an existing checksum. Cleaning temporary data...")
for tmpfile in unzipped:
os.remove(tmpfile)
return []
if verbose:
print("creating tables...")
if isinstance(dbconn, PostgresDBConn):
# The canonical SQL files below own the ais_global_* schema.
raw_files = [str(f) for f in raw_files] # Ensure file paths are strings
if timescaledb:
with open(
os.path.join(sqlpath, "timescale_createtable_dynamic.sql"), "r"
) as f:
create_dynamic_table_stmt = f.read()
with open(
os.path.join(sqlpath, "timescale_createtable_static.sql"), "r"
) as f:
create_static_table_stmt = f.read()
dbconn.execute(create_dynamic_table_stmt)
dbconn.execute(create_static_table_stmt)
if not raw_insertion:
dbconn.drop_indexes(verbose=verbose)
dbconn.commit()
else:
with open(
os.path.join(sqlpath, "psql_createtable_dynamic_noindex.sql"), "r"
) as f:
create_dynamic_table_stmt = f.read()
with open(os.path.join(sqlpath, "psql_createtable_static.sql"), "r") as f:
create_static_table_stmt = f.read()
dbconn.execute(create_dynamic_table_stmt.format("global"))
dbconn.execute(create_static_table_stmt.format("global"))
dbconn.commit()
completed_files = decoder(
psql_conn_string=dbconn.connection_string,
files=raw_files,
source=source,
verbose=verbose,
workers=workers,
type_preference=type_preference,
allow_swap=False,
)
# Surface silent ingestion failures: the Rust decoder returns only the
# files that completed; if every input errored, raise instead of
# reporting success over an empty insert.
if raw_files and not completed_files:
raise RuntimeError(
f"decoder failed for all {len(raw_files)} input files; "
"see stderr for per-file errors"
)
else:
raise ValueError("Unsupported DB connection")
if verbose and not skip_checksum:
print("saving checksums...")
for filename, signature in zip(
not_zipped + unzipped, not_zipped_checksums + unzipped_checksums
):
if filename in completed_files:
dbindex.insert_checksum(signature)
else:
if verbose:
print(f"error processing {filename}, skipping checksum...")
dbindex.dbconn.commit()
if verbose:
print("cleaning temporary data...")
try:
for tmpfile in unzipped:
os.remove(tmpfile)
print(f"Cleaning temp dir: {dbindex.tmp_dir}")
shutil.rmtree(dbindex.tmp_dir, ignore_errors=True)
except Exception as e:
print(f"Error cleaning temporary files: {e}")
if isinstance(dbconn, PostgresDBConn):
if not raw_insertion and timescaledb:
dbconn.rebuild_indexes(verbose=verbose)
dbconn.execute("ANALYZE")
dbconn.commit()
dbconn.aggregate_static_msgs(verbose=verbose)
if not raw_insertion and vacuum:
print("finished parsing data\nvacuuming...")
if vacuum is not True:
raise ValueError("vacuum arg must be a boolean for PostgreSQL")
# VACUUM cannot run inside a transaction block
dbconn.commit()
previous_autocommit = dbconn.conn.autocommit
dbconn.conn.autocommit = True
try:
dbconn.conn.execute("VACUUM")
finally:
dbconn.conn.autocommit = previous_autocommit
return completed_files
[docs]
def decode_msgs(
filepaths,
dbconn,
source,
vacuum=False,
skip_checksum=True,
workers=4,
type_preference="all",
raw_insertion=True,
verbose=True,
timescaledb=False,
):
"""
Decode messages from filepaths and insert them into a database.
:param filepaths: list of file paths to decode
:param dbconn: database connection to use for insertion
:param source: source identifier for the decoded messages
:param vacuum: whether to vacuum the database after insertion (default is False)
:param skip_checksum: whether to skip checksum validation (default is True)
:param workers: number of parallel workers to use (default is 4)
:param type_preference: preferred file type to be used (default is "all")
:param raw_insertion: whether to insert messages without indexing them (default is True)
:param verbose: whether to print verbose output (default is True)
:param timescaledb: whether to insert data to a database with timescale extension (default is False)
:return: None
"""
if not isinstance(dbconn, PostgresDBConn): # pragma: no cover
raise ValueError(
f"db argument must be a DBConn database connection. got {dbconn}"
)
if len(filepaths) == 0: # pragma: no cover
raise ValueError("must supply atleast one filepath.")
dbindex = FileChecksums(dbconn=dbconn)
# handle zipfiles
zipped = {
f for f in filepaths if f.lower()[-4:] == ".zip" or f.lower()[-3:] == ".gz"
}
not_zipped = sorted(list(set(filepaths) - set(zipped)))
not_zipped_checksums = []
unzipped_checksums = []
zipped_checksums = []
unzipped = []
_skipped = []
if verbose:
print("generating file checksums...")
for item in deepcopy(not_zipped):
with open(os.path.abspath(item), "rb") as f:
signature = dbindex.get_md5(item, f)
if skip_checksum:
continue
if dbindex.checksum_exists(signature):
_skipped.append(item)
not_zipped.remove(item)
if verbose:
print(f"found matching checksum, skipping {item}")
else:
not_zipped_checksums.append(signature)
for item in deepcopy(zipped):
with open(os.path.abspath(item), "rb") as f:
signature = dbindex.get_md5(item, f)
if skip_checksum:
# Process all files, still collect checksum for later insertion
zipped_checksums.append(signature)
else:
if dbindex.checksum_exists(signature):
_skipped.append(item)
zipped.remove(item)
if verbose:
print(f"found matching checksum, skipping {item}")
else:
zipped_checksums.append(signature)
# Track per-batch failures so a fully-failed run raises instead of
# silently reporting success.
batches_attempted = 0
batches_failed = 0
# Process not zipped files
for not_zip_file, checksum in zip(
not_zipped,
not_zipped_checksums if not skip_checksum else [None] * len(not_zipped),
):
print(f"Cleaning temp dir: {dbindex.tmp_dir}")
shutil.rmtree(dbindex.tmp_dir, ignore_errors=True)
os.makedirs(dbindex.tmp_dir, exist_ok=True)
try:
current_not_zipped = [not_zip_file]
current_not_zipped_checksums = [checksum] if not skip_checksum else []
raw_files = current_not_zipped # not_zipped + unzipped no longer needed
completed_files = process_raw_files(
dbconn,
dbindex,
raw_files,
source,
timescaledb,
raw_insertion,
vacuum,
verbose,
workers,
type_preference,
current_not_zipped,
[], # nothing unzipped here
current_not_zipped_checksums,
[], # only not_zipped checksums
skip_checksum,
)
except Exception as e:
print(f"Failed to process {not_zip_file}: {e}")
batches_failed += 1
continue
finally:
batches_attempted += 1
# Process zipped files
def batched(iterable, n):
"""Batch data into lists of length n."""
it = iter(iterable)
while True:
batch = list(islice(it, n))
if not batch:
break
yield batch
for zip_batch in batched(list(zip(zipped, zipped_checksums)), workers):
print(f"\nProcessing zip batch: {[zf for zf, _ in zip_batch]}")
shutil.rmtree(dbindex.tmp_dir, ignore_errors=True)
os.makedirs(dbindex.tmp_dir, exist_ok=True)
current_unzipped = []
current_unzipped_checksums = []
for zip_file, checksum in zip_batch:
try:
_fast_unzip(zip_file, dbindex.tmp_dir)
except Exception as e:
print(f"Failed to unzip {zip_file}: {e}")
continue
# collect all .csv files just unzipped
current_unzipped = sorted(
[
os.path.join(dbindex.tmp_dir, f)
for f in os.listdir(dbindex.tmp_dir)
if f.endswith(".csv")
]
)
# calculate checksums if needed
if not skip_checksum:
for item in current_unzipped:
with open(os.path.abspath(item), "rb") as f:
signature = dbindex.get_md5(item, f)
current_unzipped_checksums.append(signature)
raw_files = not_zipped + current_unzipped
try:
completed_files = process_raw_files(
dbconn,
dbindex,
raw_files,
source,
timescaledb,
raw_insertion,
vacuum,
verbose,
workers,
type_preference,
not_zipped,
current_unzipped,
not_zipped_checksums,
current_unzipped_checksums,
skip_checksum,
)
except Exception as e:
print(f"Failed to process batch {[zf for zf, _ in zip_batch]}: {e}")
batches_failed += 1
continue
finally:
batches_attempted += 1
# Clean temp dir after batch
shutil.rmtree(dbindex.tmp_dir, ignore_errors=True)
if batches_attempted and batches_failed == batches_attempted:
raise RuntimeError(
f"all {batches_attempted} decode batches failed; no data was inserted"
)
if batches_failed:
print(f"WARNING: {batches_failed}/{batches_attempted} decode batches failed")