diff --git a/src/masoniteorm/schema/Schema.py b/src/masoniteorm/schema/Schema.py index eaedabf8..6f2f9ece 100644 --- a/src/masoniteorm/schema/Schema.py +++ b/src/masoniteorm/schema/Schema.py @@ -62,6 +62,7 @@ def __init__( self._dry = dry self.connection = connection self.connection_class = connection_class + self._connection_driver = None self._connection = None self.grammar = grammar self.platform = platform @@ -132,7 +133,7 @@ def create(self, table): self._blueprint = Blueprint( self.grammar, - connection=self.new_connection(), + connection=self.get_connection(), table=Table(table), action="create", platform=self.platform, @@ -148,7 +149,7 @@ def create_table_if_not_exists(self, table): self._blueprint = Blueprint( self.grammar, - connection=self.new_connection(), + connection=self.get_connection(), table=Table(table), action="create_table_if_not_exists", platform=self.platform, @@ -174,7 +175,7 @@ def table(self, table): self._blueprint = Blueprint( self.grammar, - connection=self.new_connection(), + connection=self.get_connection(), table=TableDiff(table), action="alter", platform=self.platform, @@ -212,15 +213,25 @@ def get_connection_information(self): } def new_connection(self): + """Explicitly creates a new connection.""" if self._dry: return - self._connection = ( + return ( self.connection_class(**self.get_connection_information()) .set_schema(self.schema) .make_connection() ) + def get_connection(self): + """Create""" + if self._dry: + return + + if self._connection: + return self._connection + + self._connection = self.new_connection() return self._connection def has_column(self, table, column, query_only=False): @@ -238,11 +249,11 @@ def has_column(self, table, column, query_only=False): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def get_columns(self, table, dict=True): table = self.platform().get_current_schema( - self.new_connection(), table, schema=self.get_schema() + self.get_connection(), table, schema=self.get_schema() ) result = {} if dict: @@ -264,7 +275,7 @@ def drop_table(self, table, query_only=False): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def drop(self, *args, **kwargs): return self.drop_table(*args, **kwargs) @@ -276,7 +287,7 @@ def drop_table_if_exists(self, table, exists=False, query_only=False): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def rename(self, table, new_name): sql = self.platform().compile_rename_table(table, new_name) @@ -285,7 +296,7 @@ def rename(self, table, new_name): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def truncate(self, table, foreign_keys=False): sql = self.platform().compile_truncate( @@ -296,7 +307,7 @@ def truncate(self, table, foreign_keys=False): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def get_schema(self): """Gets the schema set on the migration class""" @@ -315,7 +326,7 @@ def get_all_tables(self): self._sql = sql return sql - result = self.new_connection().query(sql, ()) + result = self.get_connection().query(sql, ()) return ( list(map(lambda t: list(t.values())[0], result)) if result else [] @@ -338,7 +349,7 @@ def has_table(self, table, query_only=False): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def enable_foreign_key_constraints(self): sql = self.platform().enable_foreign_key_constraints() @@ -347,7 +358,7 @@ def enable_foreign_key_constraints(self): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) def disable_foreign_key_constraints(self): sql = self.platform().disable_foreign_key_constraints() @@ -356,4 +367,17 @@ def disable_foreign_key_constraints(self): self._sql = sql return sql - return bool(self.new_connection().query(sql, ())) + return bool(self.get_connection().query(sql, ())) + + def query_builder(self): + """Get a query builder for the schema connection""" + from ..query import QueryBuilder + + return QueryBuilder( + connection=self.connection, + connection_class=self.connection_class, + connection_driver=self._connection_driver, + connection_details=self.connection_details, + schema=self.schema, + dry=self.dry, + ) diff --git a/tests/schema/test_schema.py b/tests/schema/test_schema.py new file mode 100644 index 00000000..6d3c73a4 --- /dev/null +++ b/tests/schema/test_schema.py @@ -0,0 +1,43 @@ +import unittest + +from src.masoniteorm.query import QueryBuilder +from src.masoniteorm.schema import Schema +from src.masoniteorm.schema.platforms import SQLitePlatform +from tests.integrations.config.database import DATABASES +from tests.utils import MockSQLiteConnection + + +class TestSchema(unittest.TestCase): + maxDiff = None + + def get_schema(self): + return Schema( + connection="dev", + connection_class=MockSQLiteConnection, + connection_details=DATABASES, + platform=SQLitePlatform, + ) + + def test_connection_is_cached(self): + schema = self.get_schema() + first_connection = schema.get_connection() + self.assertIsNotNone(first_connection) + second_connection = schema.get_connection() + self.assertIsNotNone(second_connection) + self.assertEqual(first_connection, second_connection) + + def test_new_connection_is_not_cached(self): + schema = self.get_schema() + first_connection = schema.get_connection() + self.assertIsNotNone(first_connection) + new_connection = schema.new_connection() + self.assertIsNotNone(new_connection) + second_connection = schema.get_connection() + self.assertIsNotNone(second_connection) + self.assertEqual(first_connection, second_connection) + self.assertNotEqual(new_connection, first_connection) + + def test_can_get_query_builder(self): + schema = self.get_schema() + builder = schema.query_builder() + self.assertIsInstance(builder, QueryBuilder)