From 3afa7431bc43e33f99fabd3331992f557beeeec8 Mon Sep 17 00:00:00 2001 From: gsd Date: Mon, 4 Aug 2025 19:31:08 +0300 Subject: [PATCH] postgres support --- tf2idb_pgsql.py | 237 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 tf2idb_pgsql.py diff --git a/tf2idb_pgsql.py b/tf2idb_pgsql.py new file mode 100644 index 0000000..a5b77fc --- /dev/null +++ b/tf2idb_pgsql.py @@ -0,0 +1,237 @@ +import argparse +import psycopg2 +from time import time +import traceback + +class DB_Connector: + def __init__(self, host, port, db, user, password, schema, dryrun = False): + self.dryrun = dryrun + self.connector = psycopg2.connect(host = host, port = port, database=db, user = user, password=password, options="-c search_path="+schema) + #self.connector = pymysql.connect(host = host, port = port, database=db, user = user, password=password, cursorclass=pymysql.cursors.DictCursor) + + def execute(self, sql, payload = ()): + if self.dryrun: + print(sql, payload) + return + + with self.connector.cursor() as cursor: + print(sql, payload) + cursor.execute(sql, payload) + + def commit(self): + if self.dryrun: + print("--commit data") + return + + self.connector.commit() + + +class tf2idb: + def __init__(self, TF_FOLDER, DB_CONNECTOR, ITEMS_GAME = "scripts/items/items_game.txt"): + import vdf + import collections + import copy + + def dict_merge(dct, merge_dct): + """ Recursive dict merge. Inspired by :meth:dict.update(), instead of + updating only top-level keys, dict_merge recurses down into dicts nested + to an arbitrary depth, updating keys. The merge_dct is merged into + dct. + :param dct: dict onto which the merge is executed + :param merge_dct: dct merged into dct + :return: None + """ + for k, v in merge_dct.items(): + if (k == 'used_by_classes' or k == 'model_player_per_class'): #handles Demoman vs demoman... Valve pls + v = dict((k2.lower(), v2) for k2, v2 in v.items()) + if (k in dct and isinstance(dct[k], dict) and isinstance(v, collections.abc.Mapping)): + dict_merge(dct[k], v) + else: + dct[k] = copy.deepcopy(v) + + def resolve_prefabs(item, prefabs): + # generate list of prefabs + prefab_list = item.get('prefab', '').split() + for prefab in prefab_list: + subprefabs = prefabs[prefab].get('prefab', '').split() + prefab_list.extend(p for p in subprefabs if p not in prefab_list) + + # iterate over prefab list and merge, nested prefabs first + result = {} + for prefab in ( prefabs[p] for p in reversed(prefab_list) ): + dict_merge(result, prefab) + + dict_merge(result, item) + return result, prefab_list + + data = None + + ITEMS_GAME = TF_FOLDER + "/" + ITEMS_GAME + + dbc = DB_CONNECTOR + + with open(ITEMS_GAME) as f: + data = vdf.parse(f) + data = data['items_game'] + + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_class') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_item_attributes') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_item') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_particles') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_equip_conflicts') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_equip_regions') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_capabilities') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_attributes') + dbc.execute('DROP TABLE IF EXISTS new_tf2idb_qualities') + + dbc.execute('CREATE TABLE new_tf2idb_class (id INT NOT NULL , class VARCHAR(64) NOT NULL , slot VARCHAR(64) NULL, PRIMARY KEY (id, class));') + dbc.execute('CREATE TABLE new_tf2idb_item_attributes (id INT NOT NULL, attribute INT NOT NULL, value VARCHAR(128) NOT NULL,static INT NULL, PRIMARY KEY (id, attribute));' + ) + dbc.execute('CREATE TABLE new_tf2idb_item (id INT PRIMARY KEY NOT NULL, name VARCHAR(128) NOT NULL, item_name VARCHAR(128) NULL, class VARCHAR(128) NOT NULL, slot VARCHAR(128) NULL, quality VARCHAR(128) NOT NULL, tool_type VARCHAR(128) NULL, min_ilevel INT NULL, max_ilevel INT NULL, baseitem INT NULL, holiday_restriction VARCHAR(128) NULL, has_string_attribute INT NULL, propername INT NULL )' + ) + dbc.execute('CREATE TABLE new_tf2idb_particles (id INT PRIMARY KEY NOT NULL , name VARCHAR(128) NOT NULL )') + dbc.execute('CREATE TABLE new_tf2idb_equip_conflicts (name VARCHAR(128) NOT NULL , region VARCHAR(128) NOT NULL , PRIMARY KEY (name, region))') + dbc.execute('CREATE TABLE new_tf2idb_equip_regions (id INT NOT NULL , region VARCHAR(128) NOT NULL , PRIMARY KEY (id, region))') + dbc.execute('CREATE TABLE new_tf2idb_capabilities (id INT NOT NULL , capability VARCHAR(128) NOT NULL )') + dbc.execute('CREATE TABLE new_tf2idb_attributes (' + 'id INT PRIMARY KEY NOT NULL,' + 'name VARCHAR(128) NOT NULL,' + 'attribute_class VARCHAR(128) NULL,' + 'attribute_type VARCHAR(128) NULL,' + 'description_string VARCHAR(128) NULL,' + 'description_format VARCHAR(128) NULL,' + 'effect_type VARCHAR(128) NULL,' + 'hidden INT NULL,' + 'stored_as_integer INT NULL,' + 'armory_desc VARCHAR(128) NULL,' + 'is_set_bonus INT NULL,' + 'is_user_generated INT NULL,' + 'can_affect_recipe_component_name INT NULL,' + 'apply_tag_to_item_definition VARCHAR(128) NULL' + ')' + ) + dbc.execute('CREATE TABLE new_tf2idb_qualities (name VARCHAR(128) PRIMARY KEY NOT NULL , value INT NOT NULL )') + + nonce = int(time()) + dbc.execute('CREATE INDEX tf2idb_item_attributes_%i ON new_tf2idb_item_attributes (attribute ASC)' % nonce) + dbc.execute('CREATE INDEX tf2idb_class_%i ON new_tf2idb_class (class ASC)' % nonce) + dbc.execute('CREATE INDEX tf2idb_item_%i ON new_tf2idb_item (slot ASC)' % nonce) + + + # qualities + for qname,qdata in data['qualities'].items(): + dbc.execute('INSERT INTO new_tf2idb_qualities (name, value) VALUES (%s,%s)', (qname, qdata['value'])) + + # particles + for particle_type, particle_list in data['attribute_controlled_attached_particles'].items(): + for k,v in particle_list.items(): + dbc.execute('INSERT INTO new_tf2idb_particles (id,name) VALUES (%s,%s)', (k, v['system']) ) #TODO add the other fields too + + # attributes + attribute_type = {} + for k,v in data['attributes'].items(): + at = v.get('attribute_type') + if at: + atype = at + else: + if v.get('stored_as_integer'): + atype = 'integer' + else: + atype = 'float' + attribute_type[v['name'].lower()] = (k, atype) + dbc.execute('INSERT INTO new_tf2idb_attributes (id,name,attribute_class,attribute_type,description_string,description_format,effect_type,hidden,stored_as_integer,armory_desc,is_set_bonus, is_user_generated,can_affect_recipe_component_name,apply_tag_to_item_definition) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)', + (k,v.get('name'),v.get('attribute_class'),v.get('attribute_type'),v.get('description_string'),v.get('description_format'), + v.get('effect_type'),v.get('hidden'),v.get('stored_as_integer'),v.get('armory_desc'),v.get('is_set_bonus'), + v.get('is_user_generated'),v.get('can_affect_recipe_component_name'),v.get('apply_tag_to_item_definition') + ) + ) + + # conflicts + for k,v in data['equip_conflicts'].items(): + for region in v.keys(): + dbc.execute('INSERT INTO new_tf2idb_equip_conflicts (name,region) VALUES (%s,%s)', (k, region)) + + # items + for id,v in data['items'].items(): + if id == 'default': + continue + i, prefabs_used = resolve_prefabs(v, data['prefabs']) + baseitem = 1 if 'baseitem' in i else 0 + + try: + tool = None + if 'tool' in i: + tool = i['tool'].get('type') + + has_string_attribute = 0 + if 'static_attrs' in i: + for name,value in i['static_attrs'].items(): + aid,atype = attribute_type[name.lower()] + if atype == 'string': + has_string_attribute = 1 + dbc.execute('INSERT INTO new_tf2idb_item_attributes (id,attribute,value,static) VALUES (%s,%s,%s,%s)', (id,aid,value,1)) + + if 'attributes' in i: + for name,info in i['attributes'].items(): + aid,atype = attribute_type[name.lower()] + if atype == 'string': + has_string_attribute = 1 + dbc.execute('INSERT INTO new_tf2idb_item_attributes (id,attribute,value,static) VALUES (%s,%s,%s,%s)', (id,aid,info['value'],0)) + + dbc.execute('INSERT INTO new_tf2idb_item ' + '(id,name,item_name,class,slot,quality,tool_type,min_ilevel,max_ilevel,baseitem,holiday_restriction,has_string_attribute,propername) ' + 'VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)', + (id,i['name'],i.get('item_name'),i['item_class'],i.get('item_slot'),i.get('item_quality', ''), tool, i.get('min_ilevel'), i.get('max_ilevel'),baseitem, + i.get('holiday_restriction'), has_string_attribute, i.get('propername')) + ) + + if 'used_by_classes' in i: + for prof, val in i['used_by_classes'].items(): + dbc.execute('INSERT INTO new_tf2idb_class (id,class,slot) VALUES (%s,%s,%s)', (id, prof.lower(), val if val != '1' else None)) + + region_field = i.get('equip_region') or i.get('equip_regions') + if region_field: + if type(region_field) is str: + region_field = {region_field: 1} + for region in region_field.keys(): + dbc.execute('INSERT INTO new_tf2idb_equip_regions (id,region) VALUES (%s,%s)', (id, region)) + + # capabilties + for capability,val in i.get('capabilities', {}).items(): + dbc.execute('INSERT INTO new_tf2idb_capabilities (id,capability) VALUES (%s,%s)', (id, (capability if val != '0' else '!'+capability))) + + except: + traceback.print_exc() + print(id) + raise + + def replace_table(name): + dbc.execute('DROP TABLE IF EXISTS %s' % name) + dbc.execute('ALTER TABLE new_%s RENAME TO %s' % (name,name)) + + replace_table('tf2idb_class') + replace_table('tf2idb_item_attributes') + replace_table('tf2idb_item') + replace_table('tf2idb_particles') + replace_table('tf2idb_equip_conflicts') + replace_table('tf2idb_equip_regions') + replace_table('tf2idb_capabilities') + replace_table('tf2idb_attributes') + replace_table('tf2idb_qualities') + + dbc.commit() + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("tf_directory", type=str, default="") + parser.add_argument("--host", type=str, default="localhost") + parser.add_argument("--port", type=int, default=5432) + parser.add_argument("db", type=str, default="") + parser.add_argument("schema", type=str) + parser.add_argument("user", type=str, default="") + parser.add_argument("password", type=str, default="") + parser.add_argument("--dry-run", action="store_true", default=False) + args = parser.parse_args() + c = DB_Connector(args.host, args.port, args.db, args.user, args.password, args.schema, args.dry_run) + tf2idb(args.tf_directory, c) +