luizotavio32 commented on code in PR #31303:
URL: https://github.com/apache/superset/pull/31303#discussion_r1878692788
##########
superset/migrations/shared/utils.py:
##########
@@ -185,15 +193,211 @@ def has_table(table_name: str) -> bool:
return table_exists
-def add_column_if_not_exists(table_name: str, column: sa.Column) -> None:
+def drop_fks_for_table(table_name: str) -> None:
+ """
+ Drop all foreign key constraints for a table if it exist and the database
+ is not sqlite.
+
+ :param table_name: The table name to drop foreign key constraints for
+ """
+ connection = op.get_bind()
+ inspector = Inspector.from_engine(connection)
+
+ if isinstance(connection.dialect, SQLiteDialect):
+ return # sqlite doesn't like constraints
+
+ if has_table(table_name):
+ foreign_keys = inspector.get_foreign_keys(table_name)
+
+ for fk in foreign_keys:
+ op.drop_constraint(fk["name"], table_name, type_="foreignkey")
+
+
+logger = logging.getLogger("alembic")
+
+
+def create_table(table_name: str, *columns: SchemaItem) -> None:
+ """
+ Creates a database table with the specified name and columns.
+
+ This function checks if a table with the given name already exists in the
database.
+ If the table exists, it logs an informational message and skips the
creation process.
+ Otherwise, it proceeds to create a new table using the provided name and
schema columns.
+
+ :param table_name: The name of the table to be created.
+ :param columns: A variable number of arguments representing the schema
just like when calling alembic's method create_table()
+ """
+
+ if has_table(table_name=table_name):
+ logger.info(f"Table {LRED}{table_name}{RESET} already exists
Skipping...")
+ return
+
+ logger.info(f"Creating table {GREEN}{table_name}{RESET}...")
+ op.create_table(table_name, *columns)
+ logger.info(f"Table {GREEN}{table_name}{RESET} created")
+
+
+def drop_table(table_name: str) -> None:
"""
- Adds a column to a table if it does not already exist.
+ Drops a database table with the specified name.
- :param table_name: Name of the table.
- :param column: SQLAlchemy Column object.
+ This function checks if a table with the given name exists in the database.
+ If the table does not exist, it logs an informational message and skips
the dropping process.
+ If the table exists, it first attempts to drop all foreign key constraints
associated with the table
+ (handled by `drop_fks_for_table`) and then proceeds to drop the table.
+
+ :param table_name: The name of the table to be dropped.
+ """
+
+ if not has_table(table_name=table_name):
+ logger.info(f"Table {GREEN}{table_name}{RESET} doesn't exist
Skipping...")
+ return
+
+ logger.info(f"Dropping table {GREEN}{table_name}{RESET}...")
+ drop_fks_for_table(table_name)
+ op.drop_table(table_name=table_name)
+ logger.info(f"Table {GREEN}{table_name}{RESET} dropped")
+
+
+def batch_operation(
+ callable: Callable[[int, int], None], count: int, batch_size: int
+) -> None:
+ """
+ Executes an operation by dividing a task into smaller batches and tracking
progress.
+
+ This function is designed to process a large number of items in smaller
batches. It takes a callable
+ that performs the operation on each batch. The function logs the progress
of the operation as it processes
+ through the batches.
+
+ :param callable: A callable function that takes two integer arguments:
+ the start index and the end index of the current batch.
+ :param count: The total number of items to process.
+ :param batch_size: The number of items to process in each batch.
+ """
+ for offset in range(0, count, batch_size):
+ percentage = (offset / count) * 100 if count else 0
+ logger.info(f"Progress: {offset:,}/{count:,} ({percentage:.2f}%)")
+ callable(offset, min(offset + batch_size, count))
+
+ logger.info(f"Progress: {count:,}/{count:,} (100%)")
Review Comment:
Should I add a condition to skip this informational if count = 0 or print:
`Progress: 1/1 (100%)` ?
##########
superset/migrations/shared/utils.py:
##########
@@ -185,15 +193,211 @@ def has_table(table_name: str) -> bool:
return table_exists
-def add_column_if_not_exists(table_name: str, column: sa.Column) -> None:
+def drop_fks_for_table(table_name: str) -> None:
+ """
+ Drop all foreign key constraints for a table if it exist and the database
+ is not sqlite.
+
+ :param table_name: The table name to drop foreign key constraints for
+ """
+ connection = op.get_bind()
+ inspector = Inspector.from_engine(connection)
+
+ if isinstance(connection.dialect, SQLiteDialect):
+ return # sqlite doesn't like constraints
+
+ if has_table(table_name):
+ foreign_keys = inspector.get_foreign_keys(table_name)
+
+ for fk in foreign_keys:
+ op.drop_constraint(fk["name"], table_name, type_="foreignkey")
+
+
+logger = logging.getLogger("alembic")
+
+
+def create_table(table_name: str, *columns: SchemaItem) -> None:
+ """
+ Creates a database table with the specified name and columns.
+
+ This function checks if a table with the given name already exists in the
database.
+ If the table exists, it logs an informational message and skips the
creation process.
+ Otherwise, it proceeds to create a new table using the provided name and
schema columns.
+
+ :param table_name: The name of the table to be created.
+ :param columns: A variable number of arguments representing the schema
just like when calling alembic's method create_table()
+ """
+
+ if has_table(table_name=table_name):
+ logger.info(f"Table {LRED}{table_name}{RESET} already exists
Skipping...")
+ return
+
+ logger.info(f"Creating table {GREEN}{table_name}{RESET}...")
+ op.create_table(table_name, *columns)
+ logger.info(f"Table {GREEN}{table_name}{RESET} created")
+
+
+def drop_table(table_name: str) -> None:
"""
- Adds a column to a table if it does not already exist.
+ Drops a database table with the specified name.
- :param table_name: Name of the table.
- :param column: SQLAlchemy Column object.
+ This function checks if a table with the given name exists in the database.
+ If the table does not exist, it logs an informational message and skips
the dropping process.
+ If the table exists, it first attempts to drop all foreign key constraints
associated with the table
+ (handled by `drop_fks_for_table`) and then proceeds to drop the table.
+
+ :param table_name: The name of the table to be dropped.
+ """
+
+ if not has_table(table_name=table_name):
+ logger.info(f"Table {GREEN}{table_name}{RESET} doesn't exist
Skipping...")
+ return
+
+ logger.info(f"Dropping table {GREEN}{table_name}{RESET}...")
+ drop_fks_for_table(table_name)
+ op.drop_table(table_name=table_name)
+ logger.info(f"Table {GREEN}{table_name}{RESET} dropped")
+
+
+def batch_operation(
+ callable: Callable[[int, int], None], count: int, batch_size: int
+) -> None:
+ """
+ Executes an operation by dividing a task into smaller batches and tracking
progress.
+
+ This function is designed to process a large number of items in smaller
batches. It takes a callable
+ that performs the operation on each batch. The function logs the progress
of the operation as it processes
+ through the batches.
+
+ :param callable: A callable function that takes two integer arguments:
+ the start index and the end index of the current batch.
+ :param count: The total number of items to process.
+ :param batch_size: The number of items to process in each batch.
+ """
+ for offset in range(0, count, batch_size):
+ percentage = (offset / count) * 100 if count else 0
+ logger.info(f"Progress: {offset:,}/{count:,} ({percentage:.2f}%)")
+ callable(offset, min(offset + batch_size, count))
+
+ logger.info(f"Progress: {count:,}/{count:,} (100%)")
Review Comment:
Should I add a condition to skip this informational if count == 0 or print:
`Progress: 1/1 (100%)` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
