diff --git a/__init__.py b/__init__.py index 7c10e48..0cd8e38 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,52 @@ -#from .nodes.SegGPT import segGPTNode -from .nodes import MSSqlNode -NODE_CLASS_MAPPINGS = { -# **segGPTNode.NODE_CLASS_MAPPINGS, - **MSSqlNode.NODE_CLASS_MAPPINGS -} \ No newline at end of file +import configparser +import os +import pyodbc +from .models.TableInformation import TableInformation +from .nodes.QueryNode import QueryNode +from .nodes.TableNode import TableNode +from .nodes.SelectNode import SelectNode + + +class DatabaseConnection: + def __init__(self, config_file='config.ini'): + self.config_file = config_file + self.conn = None + self.connection_string= None + self._load_config() + + def _load_config(self): + config = configparser.ConfigParser() + current_dir = os.path.dirname(os.path.realpath(__file__)) + config_path = os.path.join(current_dir, self.config_file) + config.read(config_path) + server = config['MSSQL']['server'] + database = config['MSSQL']['database'] + username = config['MSSQL']['username'] + password = config['MSSQL']['password'] + driver = config['MSSQL']['driver'] + integrated_security = config.getboolean('MSSQL', 'integrated_security', fallback=False) + + self.connection_string = f'DRIVER={{{driver}}};SERVER={server};DATABASE={database};' + if integrated_security: + self.connection_string += 'Trusted_Connection=yes;' + else: + self.connection_string += f'UID={username};PWD={password};' + return self.connection_string + def connect(self): + if self.conn is None: + self.conn = pyodbc.connect(self.connection_string) + return self.conn + +# Create a connection and table info instance to pass to the node classes +db_connection = DatabaseConnection('config.ini') +table_info = TableInformation(db_connection.connect()) + +# Node classes +node_classes = [QueryNode, TableNode, SelectNode] + + +# Mapping of node names to classes +NODE_CLASS_MAPPINGS = {node_class.__name__: node_class for node_class in node_classes} + +# Mapping of node names to friendly display names +NODE_DISPLAY_NAME_MAPPINGS = {node_class.__name__: node_class.__name__.replace("Node", " Node") for node_class in node_classes} \ No newline at end of file diff --git a/config.ini b/config.ini new file mode 100644 index 0000000..34d9631 --- /dev/null +++ b/config.ini @@ -0,0 +1,7 @@ +[MSSQL] +server=WINDOWS-I1B2JBN +database=StableDiffusion +username=your_username +password=your_password +integrated_security=True +driver = SQL Server \ No newline at end of file diff --git a/models/TableInformation.py b/models/TableInformation.py new file mode 100644 index 0000000..80a526f --- /dev/null +++ b/models/TableInformation.py @@ -0,0 +1,16 @@ +class TableInformation: + def __init__(self, conn): + self.conn = conn + self.table_info = {} + self.fields = {} + def load_table_info(self): + cursor = self.conn.cursor() + cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_type = 'BASE TABLE'") + tables = [row.table_name for row in cursor.fetchall()] + for table in tables: + cursor.execute(f"SELECT * FROM {table}") + columns = {column[0]: column[1] for column in cursor.description} + # print(columns) + + self.table_info[table] = columns + self.fields = self.table_info['txt2img'] \ No newline at end of file diff --git a/nodes/QueryNode.py b/nodes/QueryNode.py new file mode 100644 index 0000000..bb9abc4 --- /dev/null +++ b/nodes/QueryNode.py @@ -0,0 +1,39 @@ +class QueryNode: + def __init__(self): + self.CATEGORY = "LexNode.MSSQL" + self.FUNCTION = "execute_query" + self.RETURN_TYPES = ("Tuple", ) + self.RETURN_NAMES = ("Query Results", ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "query": ("STRING", {"default": ""}) + }, + } + + def execute_query(self, **kwargs): + global conn + cursor = conn.cursor() + + sql = kwargs.get('query') + cursor.execute(sql) + results = cursor.fetchall() + all_rows = [] + + for result in results: + all_rows.append(result) + + return tuple(all_rows), + + RETURN_TYPES = ("Tuple", ) + RETURN_NAMES = ("Query Results", ) + +NODE_CLASS_MAPPINGS = { + "QueryNode": QueryNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "QueryNode": "Sql Query Node" +} \ No newline at end of file diff --git a/nodes/SelectNode.py b/nodes/SelectNode.py new file mode 100644 index 0000000..db492a2 --- /dev/null +++ b/nodes/SelectNode.py @@ -0,0 +1,71 @@ +import io, torch +import numpy as np +from PIL import Image + +class SelectNode: + def __init__(self, table_name = 'txt2img'): + self.table_name = table_name + self.RETURN_TYPES = self.getFields() + self.RETURN_NAMES = self.getFieldsNames() + table_name ="None" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "id": ("INT", {"default": "1"}) + }, + } + + @property + def RETURN_TYPES(self): + return self.getFields() + + @property + def RETURN_NAMES(self): + return self.getFieldsNames() + + FUNCTION = "execute_query" + CATEGORY = "LexNodes.MSSQL" + + def getFields(self): + global table_info + fields = table_info[self.table_name] + return tuple(fields.keys()) + + def getFieldsNames(self): + global table_info + fields = table_info[self.table_name] + return tuple(fields.values()) + + def execute_query(self, **kwargs): + global conn,table_info + cursor = conn.cursor() + + columns = ', '.join(table_info[self.table_name].keys()) + sql = f"SELECT {columns} FROM {self.table_name} WHERE Id = {kwargs.get('id')}" + cursor.execute(sql) + results = cursor.fetchall() + all_rows = [] + + for result in results: + result_dict = {} + for column, value in zip(table_info[self.table_name].keys(), result): + if column =='Image': # Check if the value is a bytearray + image = Image.open(io.BytesIO(value)) + # Perform additional image processing + image = image.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + result_dict[column] = image + else: + result_dict[column] = value + all_rows.append(tuple(result_dict.values())) + + return all_rows[0] +NODE_CLASS_MAPPINGS = { + "SelectNode": SelectNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SelectNode": "Sql Select Node" +} \ No newline at end of file diff --git a/nodes/TableNode.py b/nodes/TableNode.py new file mode 100644 index 0000000..c7cecf8 --- /dev/null +++ b/nodes/TableNode.py @@ -0,0 +1,109 @@ +import datetime +import pyodbc +import io, torch +import numpy as np +from PIL import Image + +class TableNode: + table_name="txt2img" + fields = {} + def __init__(self, db_connection, table_info): + self.conn = db_connection.connect() + self.table_info = table_info + self.table_name = "txt2img" + print(self.table_info) + + + @classmethod + def INPUT_TYPES(cls): + print() + tables = list(cls.table_info.keys()) + fields = cls.table_info[cls.table_name] + + # Create an input for each field + inputs = {"Table": (tables,)} + for field, field_type in fields.items(): + if field_type == int: + input_type = "INT" + default_value = 0 + elif field_type == str: + input_type = "STRING" + default_value = "" + elif field_type == bool: + input_type = "STRING" + default_value = "False" + elif field_type == float: + input_type = "FLOAT" + default_value = 0.0 + elif field_type == bytearray: + input_type = "IMAGE" + default_value = None + else: + input_type = "STRING" # default type + default_value = "" + + inputs[field] = (input_type, {"default": default_value}) + + print({"required": inputs}) + return {"required": inputs} + + RETURN_TYPES = ("STRING", "INT",) + FUNCTION = "execute_query" + CATEGORY = "LexNode.MSSQL" + + def execute_query(self, **kwargs): + cursor = self.conn.cursor() + + # Set the table name to the value of the "Table" field and remove it from kwargs + self.table_name = kwargs.pop('Table', self.table_name) + + if kwargs.get('id', None) == 0: # Check if id is 0 + # Remove 'id' from kwargs + kwargs.pop('id', None) + + # Prepare the SQL statement for inserting a new record + kwargs['DateAdded'] = datetime.datetime.now() # Add current datetime + columns = ', '.join(kwargs.keys()) + placeholders = ', '.join('?' for _ in kwargs) + sql = f"INSERT INTO {self.table_name} ({columns}) VALUES ({placeholders})" + + values = [] + for value in kwargs.values(): + if isinstance(value, torch.Tensor): # Check if the value is a PyTorch tensor + # Convert the tensor to a numpy array + i = 255. * value[0].cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + byte_array = io.BytesIO() + img.save(byte_array, format='JPEG') + value = pyodbc.Binary(byte_array.getvalue()) + elif isinstance(value, datetime.datetime): # Check if the value is a datetime object + value = value.strftime('%Y-%m-%d %H:%M:%S') # Format the datetime object to string + values.append(value) + + cursor.execute(sql, tuple(values)) + self.conn.commit() # Don't forget to commit the changes + cursor.execute("SELECT @@IDENTITY AS 'Identity'") + id_of_new_row = cursor.fetchone()[0] + return ["Insert operation completed.", id_of_new_row] + else: + for field, value in kwargs.items(): + if field in self.table_info[self.table_name]: # Ensure the field exists in the table + if isinstance(value, torch.Tensor): # Check if the value is a numpy array + # Ensure the numpy array can be represented as an image + i = 255. * value[0].cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + # Convert the numpy array to a byte array + byte_arr = io.BytesIO() + img.save(byte_arr, format='JPEG') + value = pyodbc.Binary(byte_arr.getvalue()) + cursor.execute(f"UPDATE {self.table_name} SET {field} = ? WHERE {field} = ?", (value, value)) + self.conn.commit() # Don't forget to commit the changes + return ["Update operation completed.", kwargs.get('id')] + +NODE_CLASS_MAPPINGS = { + "TableNode": TableNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TableNode": "Sql Table Node" + } \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 82a77ad..c292fdb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ numpy opencv-python git+https://github.com/facebookresearch/detectron2.git -pyodbc \ No newline at end of file +pyodbc \ No newline at end of file