michael-s-molina commented on code in PR #31303:
URL: https://github.com/apache/superset/pull/31303#discussion_r1878176634
##########
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:
Review Comment:
Could you add log messages here in the same way you do for the other
methods? Something like:
```
logger.info(f"Dropping foreign key X for table Y...")
```
##########
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:
If count is zero you'll get `ZeroDivisionError`.
##########
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%)")
+ logger.info(
+ f"End: {GREEN}{callable.__name__}{RESET} batch operation
{GREEN}succesfully{RESET} executed"
+ )
+
+
+def add_columns(table_name: str, *columns: Column) -> None:
"""
- if not table_has_column(table_name, column.name):
- print(f"Adding column '{column.name}' to table '{table_name}'.\n")
- op.add_column(table_name, column)
- else:
- print(f"Column '{column.name}' already exists in table
'{table_name}'.\n")
+ Adds new columns to an existing database table.
+
+ For each column on columns list variable, it checks if the column already
exists in the table. If a column already exists, it logs an informational
+ message and skips the addition of that column. Otherwise, it proceeds to
add the new column to the table.
+
+ The operation is performed within a batch operation context which allows a
more performant approach
Review Comment:
Could you apply the same for similar function docs?
```suggestion
The operation is performed using Alembic's batch_alter_table.
```
##########
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%)")
+ logger.info(
+ f"End: {GREEN}{callable.__name__}{RESET} batch operation
{GREEN}succesfully{RESET} executed"
+ )
+
+
+def add_columns(table_name: str, *columns: Column) -> None:
"""
- if not table_has_column(table_name, column.name):
- print(f"Adding column '{column.name}' to table '{table_name}'.\n")
- op.add_column(table_name, column)
- else:
- print(f"Column '{column.name}' already exists in table
'{table_name}'.\n")
+ Adds new columns to an existing database table.
+
+ For each column on columns list variable, it checks if the column already
exists in the table. If a column already exists, it logs an informational
Review Comment:
Could you apply the same for similar function docs?
```suggestion
If a column already exists, it logs an informational
```
##########
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")
Review Comment:
This entry should replace the one at the beginning of the file.
##########
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%)")
+ logger.info(
+ f"End: {GREEN}{callable.__name__}{RESET} batch operation
{GREEN}succesfully{RESET} executed"
+ )
+
+
+def add_columns(table_name: str, *columns: Column) -> None:
"""
- if not table_has_column(table_name, column.name):
- print(f"Adding column '{column.name}' to table '{table_name}'.\n")
- op.add_column(table_name, column)
- else:
- print(f"Column '{column.name}' already exists in table
'{table_name}'.\n")
+ Adds new columns to an existing database table.
+
+ For each column on columns list variable, it checks if the column already
exists in the table. If a column already exists, it logs an informational
+ message and skips the addition of that column. Otherwise, it proceeds to
add the new column to the table.
+
+ The operation is performed within a batch operation context which allows a
more performant approach
+
+ :param table_name: The name of the table to which the columns will be
added.
+ :param columns: A list of SQLAlchemy Column objects that define the name,
type, and other attributes of the columns to be added.
+ """
+
+ cols_to_add = []
+ for col in columns:
+ if table_has_column(table_name=table_name, column_name=col.name):
+ logger.info(
+ f"Column {LRED}{col.name}{RESET} already present on table
{LRED}{table_name}{RESET} Skipping..."
+ )
+ else:
+ cols_to_add.append(col)
+
+ with op.batch_alter_table(table_name) as batch_op:
+ for col in cols_to_add:
+ logger.info(
+ f"Adding column {GREEN}{col.name}{RESET} to table
{GREEN}{table_name}{RESET}"
+ )
+ batch_op.add_column(col)
+
+
+def drop_columns(table_name: str, *columns: str) -> None:
+ """
+ Drops specified columns from an existing database table.
+
+ For each column, it first checks if the column exists in the table. If a
column does not exist, it logs an informational
+ message and skips the dropping of that column. Otherwise, it proceeds to
remove the column from the table.
+
+ The operation is performed within a batch operation context which allows a
more performant approach
+
+ :param table_name: The name of the table from which the columns will be
removed.
+ :param columns: A list of column names to be dropped.
+ """
+
+ cols_to_drop = []
+ for col in columns:
+ if not table_has_column(table_name=table_name, column_name=col):
+ logger.info(
+ f"Column {LRED}{col}{RESET} is not present on table
{LRED}{table_name}{RESET} Skipping..."
+ )
+ else:
+ cols_to_drop.append(col)
+
+ with op.batch_alter_table(table_name) as batch_op:
+ for col in cols_to_drop:
+ logger.info(
+ f"Dropping column {GREEN}{col}{RESET} from table
{GREEN}{table_name}{RESET}"
+ )
+ batch_op.drop_column(col)
+
+
+def create_index(table_name: str, index_name: str, *columns: str) -> None:
+ """
+ Creates an index on specified columns of an existing database table.
+
+ This function checks if an index with the given name already exists on the
specified table.
+ If so, it logs an informational message and skips the index creation
process.
+ Otherwise, it proceeds to create a new index with the specified name on
the given columns of the table.
+
+ The operation is performed within a batch operation context which allows a
more performant approach
+
+ :param table_name: The name of the table on which the index will be
created.
+ :param index_name: The name of the index to be created.
+ :param columns: A list column names where the index will be created
+ """
+
+ if table_has_index(table=table_name, index=index_name):
+ logger.info(
+ f"Table {LRED}{table_name}{RESET} already has index
{LRED}{index_name}{RESET} Skipping..."
+ )
+ return
+
+ logger.info(
+ f"Creating index {GREEN}{index_name}{RESET} on table
{GREEN}{table_name}{RESET}"
+ )
+
+ with op.batch_alter_table(table_name) as batch_op:
+ batch_op.create_index(index_name=index_name, columns=columns)
Review Comment:
Do you need to use a batch operation here given that it's a single
operation? Same for `drop_index`.
--
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: notifications-unsubscr...@superset.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For additional commands, e-mail: notifications-h...@superset.apache.org
